From 4c38aab31ebf420c6423a61b953c607f53fa8b6d Mon Sep 17 00:00:00 2001 From: Hweinstock <42325418+Hweinstock@users.noreply.github.com> Date: Thu, 5 Dec 2024 10:24:14 -0500 Subject: [PATCH] feat(ec2): dry run connection script to surface errors earlier. (#6037) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem Follow up to https://github.com/aws/aws-toolkit-vscode/pull/6018#discussion_r1844284580 ## 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: Screenshot 2024-11-15 at 5 53 51 PM --- License: I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: Justin M. Keyes --- packages/core/src/awsService/ec2/model.ts | 50 +++++++---- packages/core/src/codecatalyst/model.ts | 25 +----- packages/core/src/shared/extensions/ssh.ts | 37 ++++++++- packages/core/src/shared/remoteSession.ts | 22 ++++- .../src/test/shared/extensions/ssh.test.ts | 83 ++++++++++++++++++- 5 files changed, 173 insertions(+), 44 deletions(-) diff --git a/packages/core/src/awsService/ec2/model.ts b/packages/core/src/awsService/ec2/model.ts index fa7bbee71b7..085bfa0674b 100644 --- a/packages/core/src/awsService/ec2/model.ts +++ b/packages/core/src/awsService/ec2/model.ts @@ -13,6 +13,7 @@ import { SsmClient } from '../../shared/clients/ssmClient' import { Ec2Client } from '../../shared/clients/ec2Client' import { VscodeRemoteConnection, + createBoundProcess, ensureDependencies, getDeniedSsmActions, openRemoteTerminal, @@ -20,8 +21,13 @@ import { } 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 { 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 { + 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 { 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, diff --git a/packages/core/src/codecatalyst/model.ts b/packages/core/src/codecatalyst/model.ts index b2ba6912106..768a97890ee 100644 --- a/packages/core/src/codecatalyst/model.ts +++ b/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 - -/** - * 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): ReturnType { - 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 { await fs.writeFile(bearerTokenCacheLocation(devenvId), `${bearerToken}`, 'utf8') } diff --git a/packages/core/src/shared/extensions/ssh.ts b/packages/core/src/shared/extensions/ssh.ts index ff9046b3225..1e75f9921aa 100644 --- a/packages/core/src/shared/extensions/ssh.ts +++ b/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 { + 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, diff --git a/packages/core/src/shared/remoteSession.ts b/packages/core/src/shared/remoteSession.ts index 95c45832fa8..9f51c747de7 100644 --- a/packages/core/src/shared/remoteSession.ts +++ b/packages/core/src/shared/remoteSession.ts @@ -77,7 +77,7 @@ interface DependencyPaths { readonly ssh: string } -type EnvProvider = () => Promise +export type EnvProvider = () => Promise 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): ReturnType { + const options = args[0] + const envVars = await envProvider() + const spawnOptions = { + ...options?.spawnOptions, + env: { ...envVars, ...options?.spawnOptions?.env }, + } + + return super.run({ ...options, spawnOptions }) + } + } +} diff --git a/packages/core/src/test/shared/extensions/ssh.test.ts b/packages/core/src/test/shared/extensions/ssh.test.ts index c7abc7095cd..38874e2df68 100644 --- a/packages/core/src/test/shared/extensions/ssh.test.ts +++ b/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') + }) +})