fix: further error handling enhancements

smorrisj · smorrisj · commit 6e4644024224 · 2024-10-15T20:58:47.000-04:00
diff --git a/client/src/connection/ssh/auth.ts b/client/src/connection/ssh/auth.ts
@@ -3,7 +3,15 @@
 import { l10n, window } from "vscode";
 
 import { readFileSync } from "fs";
-import { NextAuthHandler, ParsedKey, Prompt, utils } from "ssh2";
+import {
+  AgentAuthMethod,
+  KeyboardInteractiveAuthMethod,
+  ParsedKey,
+  PasswordAuthMethod,
+  Prompt,
+  PublicKeyAuthMethod,
+  utils,
+} from "ssh2";
 
 /**
  * Abstraction for presenting authentication prompts to the user.
@@ -118,13 +126,13 @@ export class AuthHandler {
    * @param resolve a function that resolves the promise that is waiting for the password
    * @param username the user name to use for the connection
    */
-  passwordAuth = async (cb: NextAuthHandler, username: string) => {
+  passwordAuth = async (username: string): Promise<PasswordAuthMethod> => {
     const pw = await this._authPresenter.presentPasswordPrompt(username);
-    cb({
+    return {
       type: "password",
       password: pw,
       username: username,
-    });
+    };
   };
 
   /**
@@ -133,8 +141,10 @@ export class AuthHandler {
    * @param resolve a function that resolves the promise that is waiting for authentication
    * @param username the user name to use for the connection
    */
-  keyboardInteractiveAuth = async (cb: NextAuthHandler, username: string) => {
-    cb({
+  keyboardInteractiveAuth = async (
+    username: string,
+  ): Promise<KeyboardInteractiveAuthMethod> => {
+    return {
       type: "keyboard-interactive",
       username: username,
       prompt: (_name, _instructions, _instructionsLang, prompts, finish) => {
@@ -146,20 +156,20 @@ export class AuthHandler {
             finish(answers);
           });
       },
-    });
+    };
   };
 
   /**
    * Authenticate to the server using the ssh-agent. See the extension Docs for more information on how to set up the ssh-agent.
    * @param cb ssh2 NextHandler callback instance. This is used to pass the authentication information to the ssh server.
    * @param username the user name to use for the connection
    */
-  sshAgentAuth = (cb: NextAuthHandler, username: string) => {
-    cb({
+  sshAgentAuth = (username: string): AgentAuthMethod => {
+    return {
       type: "agent",
       agent: process.env.SSH_AUTH_SOCK,
       username: username,
-    });
+    };
   };
 
   /**
@@ -172,10 +182,9 @@ export class AuthHandler {
    * @param username the user name to use for the connection
    */
   privateKeyAuth = async (
-    cb: NextAuthHandler,
     privateKeyFilePath: string,
     username: string,
-  ) => {
+  ): Promise<PublicKeyAuthMethod> => {
     // first, try to parse the key file without a passphrase
     const parsedKeyResult = this._keyParser.parseKey(privateKeyFilePath);
     const hasParseError = parsedKeyResult instanceof Error;
@@ -186,7 +195,6 @@ export class AuthHandler {
     // key is encrypted, prompt for passphrase
     if (passphraseRequired) {
       const passphrase = await this._authPresenter.presentPassphrasePrompt();
-
       //parse the keyfile using the passphrase
       const passphrasedKeyContentsResult = this._keyParser.parseKey(
         privateKeyFilePath,
@@ -195,23 +203,24 @@ export class AuthHandler {
 
       if (passphrasedKeyContentsResult instanceof Error) {
         throw passphrasedKeyContentsResult;
+      } else {
+        return {
+          type: "publickey",
+          key: passphrasedKeyContentsResult,
+          passphrase: passphrase,
+          username: username,
+        };
       }
-
-      cb({
-        type: "publickey",
-        key: passphrasedKeyContentsResult,
-        passphrase: passphrase,
-        username: username,
-      });
     } else {
       if (hasParseError) {
         throw parsedKeyResult;
+      } else {
+        return {
+          type: "publickey",
+          key: parsedKeyResult,
+          username: username,
+        };
       }
-      cb({
-        type: "publickey",
-        key: parsedKeyResult,
-        username: username,
-      });
     }
   };
 }
diff --git a/client/src/connection/ssh/index.ts b/client/src/connection/ssh/index.ts
@@ -3,12 +3,16 @@
 import { l10n } from "vscode";
 
 import {
+  AgentAuthMethod,
   AuthHandlerMiddleware,
   AuthenticationType,
   Client,
   ClientChannel,
   ConnectConfig,
+  KeyboardInteractiveAuthMethod,
   NextAuthHandler,
+  PasswordAuthMethod,
+  PublicKeyAuthMethod,
 } from "ssh2";
 
 import { BaseConfig, RunResult } from "..";
@@ -87,14 +91,17 @@ export class SSHSession extends Session {
         return;
       }
 
+      const authHandlerFn = this.handleSSHAuthentication();
       const cfg: ConnectConfig = {
         host: this._config.host,
         port: this._config.port,
         username: this._config.username,
         readyTimeout: CONNECT_READY_TIMEOUT,
         keepaliveInterval: KEEPALIVE_INTERVAL,
         keepaliveCountMax: KEEPALIVE_UNANSWERED_THRESHOLD,
-        authHandler: this.handleSSHAuthentication,
+        authHandler: (methodsLeft, partialSuccess, callback) => (
+          authHandlerFn(methodsLeft, partialSuccess, callback), undefined
+        ),
       };
 
       if (!this._conn) {
@@ -305,71 +312,74 @@ export class SSHSession extends Session {
    */
   private clearAuthState = (): void => {
     this._sessionReady = false;
+    this._authsLeft = [];
   };
 
-  private handleSSHAuthentication: AuthHandlerMiddleware = (
-    authsLeft: AuthenticationType[],
-    partialSuccess: boolean, //used in scenarios which require multiple auth methods to denote partial success
-    nextAuth: NextAuthHandler,
-  ) => {
-    if (!authsLeft) {
-      nextAuth("none"); //sending none will prompt the server to send supported auth methods
-      return;
-    }
-
-    if (authsLeft.length === 0) {
-      this._reject?.(
-        new Error(l10n.t("Could not authenticate to the SSH server.")),
-      );
-      this.clearAuthState();
-      return false; //returning false will stop the authentication process
-    }
+  private handleSSHAuthentication = (): AuthHandlerMiddleware => {
+    //The ssh2 library supports sending false to stop the authentication process
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const END_AUTH = false as unknown as AuthenticationType;
+    return async (
+      authsLeft: AuthenticationType[],
+      partialSuccess: boolean, //used in scenarios which require multiple auth methods to denote partial success
+      nextAuth: NextAuthHandler,
+    ) => {
+      if (!authsLeft) {
+        return nextAuth({ type: "none", username: this._config.username }); //sending none will prompt the server to send supported auth methods
+      } else {
+        if (authsLeft.length === 0) {
+          return nextAuth(END_AUTH);
+        }
 
-    if (this._authsLeft.length === 0 || partialSuccess) {
-      this._authsLeft = authsLeft;
-    }
+        if (this._authsLeft.length === 0 || partialSuccess) {
+          this._authsLeft = authsLeft;
+        }
 
-    const authMethod = this._authsLeft.shift();
-
-    switch (authMethod) {
-      case "publickey": {
-        //user set a keyfile path in profile config
-        if (this._config.privateKeyFilePath) {
-          this._authHandler
-            .privateKeyAuth(
-              nextAuth,
-              this._config.privateKeyFilePath,
-              this._config.username,
-            )
-            .catch((e) => {
-              this._reject?.(e);
-              return false;
-            });
-        } else if (process.env.SSH_AUTH_SOCK) {
-          this._authHandler.sshAgentAuth(nextAuth, this._config.username);
+        const authMethod = this._authsLeft.shift();
+
+        try {
+          let authPayload:
+            | PublicKeyAuthMethod
+            | AgentAuthMethod
+            | PasswordAuthMethod
+            | KeyboardInteractiveAuthMethod;
+
+          switch (authMethod) {
+            case "publickey": {
+              //user set a keyfile path in profile config
+              if (this._config.privateKeyFilePath) {
+                authPayload = await this._authHandler.privateKeyAuth(
+                  this._config.privateKeyFilePath,
+                  this._config.username,
+                );
+              } else if (process.env.SSH_AUTH_SOCK) {
+                authPayload = this._authHandler.sshAgentAuth(
+                  this._config.username,
+                );
+              }
+              break;
+            }
+            case "password": {
+              authPayload = await this._authHandler.passwordAuth(
+                this._config.username,
+              );
+              break;
+            }
+            case "keyboard-interactive": {
+              authPayload = await this._authHandler.keyboardInteractiveAuth(
+                this._config.username,
+              );
+              break;
+            }
+            default:
+              nextAuth(authMethod);
+          }
+          return nextAuth(authPayload);
+        } catch (e) {
+          this._reject?.(e);
+          return nextAuth(END_AUTH);
         }
-        break;
       }
-      case "password": {
-        this._authHandler
-          .passwordAuth(nextAuth, this._config.username)
-          .catch((e) => {
-            this._reject?.(e);
-            return false;
-          });
-        break;
-      }
-      case "keyboard-interactive": {
-        this._authHandler
-          .keyboardInteractiveAuth(nextAuth, this._config.username)
-          .catch((e) => {
-            this._reject?.(e);
-            return false;
-          });
-        break;
-      }
-      default:
-        nextAuth(authMethod);
-    }
+    };
   };
 }
diff --git a/client/test/connection/ssh/auth.test.ts b/client/test/connection/ssh/auth.test.ts
