Fixed disconnection events handling

Signed-off-by: Seb Julliand <sebjulliand@gmail.com>

Author: sebjulliand

src/Instance.ts

@@ -1,7 +1,7 @@
 import { EventEmitter } from "stream";
 import * as vscode from "vscode";
 import { ILELibrarySettings } from "./api/CompileTools";
-import IBMi, { ConnectionResult } from "./api/IBMi";
+import IBMi, { ConnectionResult, DisconnectedCallback } from "./api/IBMi";
 import { BaseStorage } from "./api/configuration/storage/BaseStorage";
 import { CodeForIStorage } from "./api/configuration/storage/CodeForIStorage";
 import { ConnectionStorage } from "./api/configuration/storage/ConnectionStorage";
@@ -79,27 +79,32 @@ export default class Instance {
 
     let result: ConnectionResult;
 
-    const timeoutHandler = async (conn: IBMi) => {
-      if (conn) {
-        const choice = await vscode.window.showWarningMessage(`Connection lost`, {
-          modal: true,
-          detail: `Connection to ${conn.currentConnectionName} has dropped. Would you like to reconnect?`
-        }, `Yes`, `No, get logs`);
-
-        let reconnect = choice === `Yes`;
-        let collectLogs = choice === `No, get logs`;
-
-        if (collectLogs) {
-          const logs = this.output.content;
-          vscode.workspace.openTextDocument({ content: logs, language: `plaintext` }).then(doc => {
-            vscode.window.showTextDocument(doc);
-          });
-        }
+    const onDisconnected: DisconnectedCallback = async (connection, error) => {
+      if (connection.connectionSuccessful) {
+        this.setConnection();
+        this.fire(`disconnected`);
 
-        this.disconnect();
+        if (error) {
+          const choice = await vscode.window.showWarningMessage(`Connection lost: ${error.description || error.message || error.level}`, {
+            modal: true,
+            detail: `Connection to ${connection.currentConnectionName} has dropped. Would you like to reconnect?`
+          }, `Yes`, `No, get logs`);
 
-        if (reconnect) {
-          await this.connect({ ...options, reconnecting: true });
+          const reconnect = choice === `Yes`;
+          const collectLogs = choice === `No, get logs`;
+
+          if (collectLogs) {
+            const logs = this.output.content;
+            vscode.workspace.openTextDocument({ content: logs, language: `plaintext` }).then(doc => {
+              vscode.window.showTextDocument(doc);
+            });
+          }
+
+          this.disconnect();
+
+          if (reconnect) {
+            await this.connect({ ...options, reconnecting: true });
+          }
         }
       }
     };
@@ -119,8 +124,8 @@ export default class Instance {
               options.data,
               {
                 callbacks: {
-                  timeoutCallback: timeoutHandler,
-                  onConnectedOperations: options.onConnectedOperations || [],
+                  onDisconnected,
+                  onConnectedOperations: options.onConnectedOperations,
                   uiErrorHandler: handleConnectionResults,
                   progress: (message) => { p.report(message) },
                   message: messageCallback,
@@ -183,11 +188,6 @@ export default class Instance {
     }
 
     if (connection) {
-      connection.setDisconnectedCallback(async () => {
-        this.setConnection();
-        this.fire(`disconnected`);
-      });
-
       this.connection = connection;
       this.storage.setConnectionName(connection.currentConnectionName);
       await IBMi.GlobalStorage.setLastConnection(connection.currentConnectionName);

src/api/IBMi.ts

@@ -29,6 +29,14 @@ export interface ConnectionResult {
   errorCodes?: ConnectionErrorCode[]
 }
 
+export type SSHError = {
+  level: 'client-socket' | 'client-ssh'
+  message?: string
+  description?: string
+}
+
+export type DisconnectedCallback = (conn: IBMi, error?: SSHError) => Promise<void>;
+
 const remoteApps = [ // All names MUST also be defined as key in 'remoteFeatures' below!!
   {
     path: `/usr/bin/`,
@@ -45,11 +53,9 @@ const remoteApps = [ // All names MUST also be defined as key in 'remoteFeatures
   }
 ];
 
-type DisconnectCallback = (conn: IBMi) => Promise<void>;
-
 interface ConnectionCallbacks {
   onConnectedOperations?: Function[],
-  timeoutCallback?: (conn: IBMi) => Promise<void>,
+  onDisconnected?: DisconnectedCallback,
   uiErrorHandler: (connection: IBMi, error: ConnectionErrorCode, data?: any) => Promise<boolean>,
   progress: (detail: { message: string }) => void,
   message: (type: ConnectionMessageType, message: string) => void,
@@ -116,6 +122,8 @@ export default class IBMi {
   private currentAsp: string | undefined;
   private libraryAsps = new Map<string, number>();
 
+  connectionSuccessful = false;
+
   /**
    * @deprecated Will be replaced with {@link IBMi.getAllIAsps} in v3.0.0
    */
@@ -147,15 +155,6 @@ export default class IBMi {
     process.stdout.write(text);
   };
 
-  private disconnectedCallback: (DisconnectCallback) | undefined;
-
-  /**
-   * Will only be called once per connection.
-   */
-  setDisconnectedCallback(callback: DisconnectCallback) {
-    this.disconnectedCallback = callback;
-  }
-
   /**
    * getConfigFile can return pre-defined configuration files,
    * but can lazy load new configuration files as well.
@@ -336,19 +335,17 @@ export default class IBMi {
         };
       }
 
-      if (callbacks.timeoutCallback) {
-        const timeoutCallbackWrapper = () => {
-          // Don't call the callback function if it was based on a user cancellation request.
-          if (!wasCancelled) {
-            callbacks.timeoutCallback!(this);
-          }
+      // Trigger callbacks unless the connection is cancelled
+      const callbackWrapper = (error?: SSHError) => {
+        if (!wasCancelled) {
+          callbacks.onDisconnected?.(this, error);
         }
+      };
 
-        // Register handlers after we might have to abort due to bad configuration.
-        this.client.connection!.once(`timeout`, timeoutCallbackWrapper);
-        this.client.connection!.once(`end`, timeoutCallbackWrapper);
-        this.client.connection!.once(`error`, timeoutCallbackWrapper);
-      }
+      //end: disconnected by user
+      this.client.connection.once(`end`, callbackWrapper);
+      //error: connection dropped for some reason (details given in the SSHError type) 
+      this.client.connection.once(`error`, callbackWrapper);
 
       callbacks.progress({
         message: `Checking home directory.`
@@ -956,8 +953,7 @@ export default class IBMi {
       }
 
       if (!options.reconnecting) {
-        const delayedOperations: Function[] = callbacks.onConnectedOperations ? [...callbacks.onConnectedOperations] : [];
-        for (const operation of delayedOperations) {
+        for (const operation of callbacks.onConnectedOperations || []) {
           await operation();
         }
       }
@@ -977,6 +973,8 @@ export default class IBMi {
         maximumArgsLength: this.maximumArgsLength
       });
 
+      this.connectionSuccessful = true;
+
       return {
         success: true
       };
@@ -1172,11 +1170,8 @@ export default class IBMi {
     }
 
     if (this.client) {
+      this.client.dispose();
       this.client = undefined;
-
-      if (failedToConnect === false && this.disconnectedCallback) {
-        this.disconnectedCallback(this);
-      }
     }
   }