From 823a0e3f0d0cdb416bbd8f27ecd2fe7fc1d73288 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 15 Nov 2024 17:56:11 -0500
Subject: [PATCH] implement custom error to show more accurate error message

---
 packages/core/src/awsService/ec2/model.ts  | 13 ++++---------
 packages/core/src/shared/extensions/ssh.ts |  5 ++++-
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/packages/core/src/awsService/ec2/model.ts b/packages/core/src/awsService/ec2/model.ts
index 5280d951d56..1635e6aad8c 100644
--- a/packages/core/src/awsService/ec2/model.ts
+++ b/packages/core/src/awsService/ec2/model.ts
@@ -23,6 +23,7 @@ import { DefaultIamClient } from '../../shared/clients/iamClient'
 import { ErrorInformation } from '../../shared/errors'
 import {
     sshAgentSocketVariable,
+    SSHError,
     startSshAgent,
     startVscodeRemote,
     testSshConnection,
@@ -154,13 +155,6 @@ export class Ec2Connecter implements vscode.Disposable {
         }
     }
 
-    public throwGeneralConnectionError(selection: Ec2Selection, error: Error) {
-        this.throwConnectionError('Unable to connect to target instance. ', selection, {
-            code: 'EC2SSMConnect',
-            cause: error,
-        })
-    }
-
     public async checkForStartSessionError(selection: Ec2Selection): Promise<void> {
         await this.checkForInstanceStatusError(selection)
 
@@ -189,7 +183,7 @@ export class Ec2Connecter implements vscode.Disposable {
             const response = await this.ssmClient.startSession(selection.instanceId)
             await this.openSessionInTerminal(response, selection)
         } catch (err: unknown) {
-            this.throwGeneralConnectionError(selection, err as Error)
+            this.throwConnectionError('', selection, err as Error)
         }
     }
 
@@ -209,7 +203,8 @@ export class Ec2Connecter implements vscode.Disposable {
             )
             await startVscodeRemote(remoteEnv.SessionProcess, remoteEnv.hostname, '/', remoteEnv.vscPath, remoteUser)
         } catch (err) {
-            this.throwGeneralConnectionError(selection, err as Error)
+            const message = err instanceof SSHError ? 'Testing SSH connection to instance failed' : ''
+            this.throwConnectionError(message, selection, err as Error)
         } finally {
             await this.ssmClient.terminateSession(testSession)
         }
diff --git a/packages/core/src/shared/extensions/ssh.ts b/packages/core/src/shared/extensions/ssh.ts
index e88580f6a32..6f5d3ad13cf 100644
--- a/packages/core/src/shared/extensions/ssh.ts
+++ b/packages/core/src/shared/extensions/ssh.ts
@@ -13,11 +13,14 @@ import { ArrayConstructor, NonNullObject } from '../utilities/typeConstructors'
 import { Settings } from '../settings'
 import { VSCODE_EXTENSION_ID } from '../extensions'
 import { SSM } from 'aws-sdk'
+import { ToolkitError } from '../errors'
 
 const localize = nls.loadMessageBundle()
 
 export const sshAgentSocketVariable = 'SSH_AUTH_SOCK'
 
+export class SSHError extends ToolkitError {}
+
 export function getSshConfigPath(): string {
     const sshConfigDir = path.join(fs.getUserHomeDir(), '.ssh')
     return path.join(sshConfigDir, 'config')
@@ -136,7 +139,7 @@ export async function testSshConnection(
         })
     } catch (error) {
         getLogger().error('SSH connection test failed: %O', error)
-        throw error
+        throw new SSHError('SSH connection test failed', { cause: error as Error })
     }
 }