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

## Problem
Follow up to
#6018 (comment)

## Solution
- Run `ssh` within the same env as it will be run on real connection. 
- Log any resulting errors, and inform user where the process failed. 
- Also part of this PR is moving some functions to `remoteSession.ts`
that are general enough to be there.

Error msg:
<img width="1518" alt="Screenshot 2024-11-15 at 5 53 51 PM"
src="https://github.com/user-attachments/assets/e8e56887-b792-43e3-9d94-73d0e5246766">


---

<!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md -->

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: Justin M. Keyes <[email protected]>

Author: Hweinstock

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

@@ -8,15 +8,26 @@ import * as path from 'path'
 import * as nls from 'vscode-nls'
 import fs from '../fs/fs'
 import { getLogger } from '../logger'
-import { ChildProcess } from '../utilities/processUtils'
+import { ChildProcess, ChildProcessResult } 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 { ErrorInformation, ToolkitError } from '../errors'
 
 const localize = nls.loadMessageBundle()
 
 export const sshAgentSocketVariable = 'SSH_AUTH_SOCK'
 
+export class SshError extends ToolkitError {
+    constructor(message: string, options: ErrorInformation) {
+        super(message, {
+            ...options,
+            code: SshError.name,
+        })
+    }
+}
+
 export function getSshConfigPath(): string {
     const sshConfigDir = path.join(fs.getUserHomeDir(), '.ssh')
     return path.join(sshConfigDir, 'config')
@@ -119,6 +130,30 @@ 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<ChildProcessResult | never> {
+    try {
+        const env = { SESSION_ID: session.SessionId, STREAM_URL: session.StreamUrl, TOKEN: session.TokenValue }
+        const result = await new ProcessClass(sshPath, [
+            '-T',
+            `${user}@${hostname}`,
+            'echo "test connection succeeded" && exit',
+        ]).run({
+            spawnOptions: {
+                env,
+            },
+        })
+        return result
+    } catch (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 remote 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 })
+        }
+    }
+}

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

@@ -4,7 +4,14 @@
  */
 import * as assert from 'assert'
 import { ChildProcess } from '../../../shared/utilities/processUtils'
-import { startSshAgent } from '../../../shared/extensions/ssh'
+import { startSshAgent, testSshConnection } from '../../../shared/extensions/ssh'
+import { createBoundProcess } from '../../../shared/remoteSession'
+import { createExecutableFile, createTestWorkspaceFolder } from '../../testUtil'
+import { WorkspaceFolder } from 'vscode'
+import path from 'path'
+import { SSM } from 'aws-sdk'
+import { fs } from '../../../shared/fs/fs'
+import { isWin } from '../../../shared/vscode/env'
 
 describe('SSH Agent', function () {
     it('can start the agent on windows', async function () {
@@ -29,3 +36,77 @@ describe('SSH Agent', function () {
         assert.strictEqual(await getStatus(), 'Running')
     })
 })
+
+function echoEnvVarsCmd(varNames: string[]) {
+    const toShell = (s: string) => (isWin() ? `%${s}%` : `$${s}`)
+    return `echo "${varNames.map(toShell).join(' ')}"`
+}
+
+/**
+ * Trim noisy windows ChildProcess result to final line for easier testing.
+ */
+function assertOutputContains(rawOutput: string, expectedString: string): void | never {
+    const output = rawOutput.trim().split('\n').at(-1)?.replace('"', '') ?? ''
+    assert.ok(output.includes(expectedString), `Expected output to contain "${expectedString}", but got "${output}"`)
+}
+
+describe('testSshConnection', function () {
+    let testWorkspace: WorkspaceFolder
+    let sshPath: string
+
+    before(async function () {
+        testWorkspace = await createTestWorkspaceFolder()
+        sshPath = path.join(testWorkspace.uri.fsPath, `fakeSSH${isWin() ? '.cmd' : ''}`)
+    })
+
+    after(async function () {
+        await fs.delete(testWorkspace.uri.fsPath, { recursive: true, force: true })
+        await fs.delete(sshPath, { force: true })
+    })
+
+    it('runs in bound process', async function () {
+        const envProvider = async () => ({ MY_VAR: 'yes' })
+        const process = createBoundProcess(envProvider)
+        const session = {
+            SessionId: 'testSession',
+            StreamUrl: 'testUrl',
+            TokenValue: 'testToken',
+        } as SSM.StartSessionResponse
+
+        await createExecutableFile(sshPath, echoEnvVarsCmd(['MY_VAR']))
+        const r = await testSshConnection(process, 'localhost', sshPath, 'test-user', session)
+        assertOutputContains(r.stdout, 'yes')
+    })
+
+    it('injects new session into env', async function () {
+        const oldSession = {
+            SessionId: 'testSession1',
+            StreamUrl: 'testUrl1',
+            TokenValue: 'testToken1',
+        } as SSM.StartSessionResponse
+        const newSession = {
+            SessionId: 'testSession2',
+            StreamUrl: 'testUrl2',
+            TokenValue: 'testToken2',
+        } as SSM.StartSessionResponse
+        const envProvider = async () => ({
+            SESSION_ID: oldSession.SessionId,
+            STREAM_URL: oldSession.StreamUrl,
+            TOKEN: oldSession.TokenValue,
+        })
+        const process = createBoundProcess(envProvider)
+
+        await createExecutableFile(sshPath, echoEnvVarsCmd(['SESSION_ID', 'STREAM_URL', 'TOKEN']))
+        const r = await testSshConnection(process, 'localhost', sshPath, 'test-user', newSession)
+        assertOutputContains(r.stdout, `${newSession.SessionId} ${newSession.StreamUrl} ${newSession.TokenValue}`)
+    })
+
+    it('passes proper args to the ssh invoke', async function () {
+        const executableFileContent = isWin() ? `echo "%1 %2"` : `echo "$1 $2"`
+        const process = createBoundProcess(async () => ({}))
+        await createExecutableFile(sshPath, executableFileContent)
+        const r = await testSshConnection(process, 'localhost', sshPath, 'test-user', {} as SSM.StartSessionResponse)
+        assertOutputContains(r.stdout, '-T')
+        assertOutputContains(r.stdout, 'test-user@localhost')
+    })
+})