feat(ec2): dry run connection script to surface errors earlier.

packages/core/src/awsService/ec2/model.ts

@@ -13,15 +13,21 @@ import { SsmClient } from '../../shared/clients/ssmClient'
 import { Ec2Client } from '../../shared/clients/ec2Client'
 import {
     VscodeRemoteConnection,
+    createBoundProcess,
     ensureDependencies,
     getDeniedSsmActions,
     openRemoteTerminal,
     promptToAddInlinePolicy,
 } from '../../shared/remoteSession'
 import { DefaultIamClient } from '../../shared/clients/iamClient'
 import { ErrorInformation } from '../../shared/errors'
-import { sshAgentSocketVariable, startSshAgent, startVscodeRemote } from '../../shared/extensions/ssh'
-import { createBoundProcess } from '../../codecatalyst/model'
+import {
+    sshAgentSocketVariable,
+    SSHError,
+    startSshAgent,
+    startVscodeRemote,
+    testSshConnection,
+} from '../../shared/extensions/ssh'
 import { getLogger } from '../../shared/logger/logger'
 import { CancellationError, Timeout } from '../../shared/utilities/timeoutUtils'
 import { showMessageWithCancel } from '../../shared/utilities/messages'
@@ -149,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)
 
@@ -184,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)
         }
     }
 
@@ -193,11 +192,21 @@ export class Ec2Connecter implements vscode.Disposable {
 
         const remoteUser = await this.getRemoteUser(selection.instanceId)
         const remoteEnv = await this.prepareEc2RemoteEnvWithProgress(selection, remoteUser)
-
+        const testSession = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession')
         try {
+            await testSshConnection(
+                remoteEnv.SessionProcess,
+                remoteEnv.hostname,
+                remoteEnv.sshPath,
+                remoteUser,
+                testSession
+            )
             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)
         }
     }
 
@@ -208,12 +217,19 @@ export class Ec2Connecter implements vscode.Disposable {
         return remoteEnv
     }
 
+    private async startSSMSession(instanceId: string): Promise<SSM.StartSessionResponse> {
+        const ssmSession = await this.ssmClient.startSession(instanceId, 'AWS-StartSSHSession')
+        await this.addActiveSession(instanceId, ssmSession.SessionId!)
+        return ssmSession
+    }
+
     public async prepareEc2RemoteEnv(selection: Ec2Selection, remoteUser: string): Promise<Ec2RemoteEnv> {
         const logger = this.configureRemoteConnectionLogger(selection.instanceId)
         const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
         const keyPair = await this.configureSshKeys(selection, remoteUser)
-        const hostNamePrefix = 'aws-ec2-'
-        const sshConfig = new SshConfig(ssh, hostNamePrefix, 'ec2_connect', keyPair.getPrivateKeyPath())
+        const hostnamePrefix = 'aws-ec2-'
+        const hostname = `${hostnamePrefix}${selection.instanceId}`
+        const sshConfig = new SshConfig(ssh, hostnamePrefix, 'ec2_connect', keyPair.getPrivateKeyPath())
 
         const config = await sshConfig.ensureValid()
         if (config.isErr()) {
@@ -222,8 +238,8 @@ export class Ec2Connecter implements vscode.Disposable {
 
             throw err
         }
-        const ssmSession = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession')
-        await this.addActiveSession(selection.instanceId, ssmSession.SessionId!)
+
+        const ssmSession = await this.startSSMSession(selection.instanceId)
 
         const vars = getEc2SsmEnv(selection, ssm, ssmSession)
         const envProvider = async () => {
@@ -236,7 +252,7 @@ export class Ec2Connecter implements vscode.Disposable {
         })
 
         return {
-            hostname: `${hostNamePrefix}${selection.instanceId}`,
+            hostname,
             envProvider,
             sshPath: ssh,
             vscPath: vsc,

packages/core/src/codecatalyst/model.ts

@@ -19,7 +19,6 @@ import { getLogger } from '../shared/logger'
 import { AsyncCollection, toCollection } from '../shared/utilities/asyncCollection'
 import { getCodeCatalystSpaceName, getCodeCatalystProjectName, getCodeCatalystDevEnvId } from '../shared/vscode/env'
 import { sshAgentSocketVariable, startSshAgent, startVscodeRemote } from '../shared/extensions/ssh'
-import { ChildProcess } from '../shared/utilities/processUtils'
 import { isDevenvVscode } from './utils'
 import { Timeout } from '../shared/utilities/timeoutUtils'
 import { Commands } from '../shared/vscode/commands2'
@@ -28,7 +27,7 @@ import { fileExists } from '../shared/filesystemUtilities'
 import { CodeCatalystAuthenticationProvider } from './auth'
 import { ToolkitError } from '../shared/errors'
 import { Result } from '../shared/utilities/result'
-import { VscodeRemoteConnection, ensureDependencies } from '../shared/remoteSession'
+import { EnvProvider, VscodeRemoteConnection, createBoundProcess, ensureDependencies } from '../shared/remoteSession'
 import { SshConfig, sshLogFileLocation } from '../shared/sshConfig'
 import { fs } from '../shared'
 
@@ -111,28 +110,6 @@ export function createCodeCatalystEnvProvider(
     }
 }
 
-type EnvProvider = () => Promise<NodeJS.ProcessEnv>
-
-/**
- * Creates a new {@link ChildProcess} class bound to a specific dev environment. All instances of this
- * derived class will have SSM session information injected as environment variables as-needed.
- */
-export function createBoundProcess(envProvider: EnvProvider): typeof ChildProcess {
-    type Run = ChildProcess['run']
-    return class SessionBoundProcess extends ChildProcess {
-        public override async run(...args: Parameters<Run>): ReturnType<Run> {
-            const options = args[0]
-            const envVars = await envProvider()
-            const spawnOptions = {
-                ...options?.spawnOptions,
-                env: { ...envVars, ...options?.spawnOptions?.env },
-            }
-
-            return super.run({ ...options, spawnOptions })
-        }
-    }
-}
-
 export async function cacheBearerToken(bearerToken: string, devenvId: string): Promise<void> {
     await fs.writeFile(bearerTokenCacheLocation(devenvId), `${bearerToken}`, 'utf8')
 }

packages/core/src/shared/extensions/ssh.ts

@@ -12,11 +12,15 @@ import { ChildProcess } from '../utilities/processUtils'
 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')
@@ -119,6 +123,26 @@ export class RemoteSshSettings extends Settings.define('remote.SSH', remoteSshTy
     }
 }
 
+export async function testSshConnection(
+    ProcessClass: typeof ChildProcess,
+    hostname: string,
+    sshPath: string,
+    user: string,
+    session: SSM.StartSessionResponse
+): Promise<void> {
+    try {
+        const env = { SESSION_ID: session.SessionId, STREAM_URL: session.StreamUrl, TOKEN: session.TokenValue }
+        await new ProcessClass(sshPath, ['-T', `${user}@${hostname}`, 'echo connected && exit']).run({
+            spawnOptions: {
+                env,
+            },
+        })
+    } catch (error) {
+        getLogger().error('SSH connection test failed: %O', error)
+        throw new SSHError('SSH connection test failed', { cause: error as Error })
+    }
+}
+
 export async function startVscodeRemote(
     ProcessClass: typeof ChildProcess,
     hostname: string,

packages/core/src/shared/remoteSession.ts

@@ -77,7 +77,7 @@ interface DependencyPaths {
     readonly ssh: string
 }
 
-type EnvProvider = () => Promise<NodeJS.ProcessEnv>
+export type EnvProvider = () => Promise<NodeJS.ProcessEnv>
 
 export interface VscodeRemoteConnection {
     readonly sshPath: string
@@ -251,3 +251,23 @@ export async function getDeniedSsmActions(client: IamClient, roleArn: string): P
 
     return deniedActions
 }
+
+/**
+ * Creates a new {@link ChildProcess} class bound to a specific dev environment. All instances of this
+ * derived class will have SSM session information injected as environment variables as-needed.
+ */
+export function createBoundProcess(envProvider: EnvProvider): typeof ChildProcess {
+    type Run = ChildProcess['run']
+    return class SessionBoundProcess extends ChildProcess {
+        public override async run(...args: Parameters<Run>): ReturnType<Run> {
+            const options = args[0]
+            const envVars = await envProvider()
+            const spawnOptions = {
+                ...options?.spawnOptions,
+                env: { ...envVars, ...options?.spawnOptions?.env },
+            }
+
+            return super.run({ ...options, spawnOptions })
+        }
+    }
+}