Merge master into feature/emr

Author: aws-toolkit-automation

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

@@ -29,14 +29,14 @@ import {
     testSshConnection,
 } from '../../shared/extensions/ssh'
 import { getLogger } from '../../shared/logger/logger'
-import { CancellationError, Timeout } from '../../shared/utilities/timeoutUtils'
+import { CancellationError, Timeout, waitUntil } from '../../shared/utilities/timeoutUtils'
 import { showMessageWithCancel } from '../../shared/utilities/messages'
 import { SshConfig } from '../../shared/sshConfig'
 import { SshKeyPair } from './sshKeyPair'
 import { Ec2SessionTracker } from './remoteSessionManager'
 import { getEc2SsmEnv } from './utils'
 
-export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus'
+export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMTestConnect' | 'EC2SSMAgentStatus'
 
 export interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
@@ -150,8 +150,14 @@ export class Ec2Connecter implements vscode.Disposable {
         }
     }
 
-    private async checkForInstanceSsmError(selection: Ec2Selection): Promise<void> {
-        const isSsmAgentRunning = (await this.ssmClient.getInstanceAgentPingStatus(selection.instanceId)) === 'Online'
+    public async checkForInstanceSsmError(
+        selection: Ec2Selection,
+        options?: Partial<{ interval: number; timeout: number }>
+    ): Promise<void> {
+        const isSsmAgentRunning = await waitUntil(
+            async () => (await this.ssmClient.getInstanceAgentPingStatus(selection.instanceId)) === 'Online',
+            { interval: options?.interval ?? 500, timeout: options?.timeout ?? 5000 }
+        )
 
         if (!isSsmAgentRunning) {
             this.throwConnectionError('Is SSM Agent running on the target instance?', selection, {
@@ -215,8 +221,8 @@ export class Ec2Connecter implements vscode.Disposable {
                 remoteUser.name
             )
         } catch (err) {
-            const message = err instanceof SshError ? 'Testing SSH connection to instance failed' : ''
-            this.throwConnectionError(message, selection, err as Error)
+            const message = err instanceof SshError ? `Testing SSM connection to instance failed: ${err.message}` : ''
+            this.throwConnectionError(message, selection, { ...(err as Error), code: 'EC2SSMTestConnect' })
         } finally {
             await this.ssmClient.terminateSession(testSession)
         }

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

@@ -130,27 +130,34 @@ export class RemoteSshSettings extends Settings.define('remote.SSH', remoteSshTy
     }
 }
 
+/**
+ * Test a SSH connection over SSM.
+ * @param ProcessClass given process to test the connection within.
+ * @param hostname
+ * @param sshPath
+ * @param user
+ * @param session SSM session credentials. These cannot be reused, so it may be required to create a seperate session for the test connection.
+ * @returns
+ */
 export async function testSshConnection(
     ProcessClass: typeof ChildProcess,
     hostname: string,
     sshPath: string,
     user: string,
     session: SSM.StartSessionResponse
 ): Promise<ChildProcessResult | never> {
+    const env = { SESSION_ID: session.SessionId, STREAM_URL: session.StreamUrl, TOKEN: session.TokenValue }
+    const process = new ProcessClass(sshPath, ['-T', `${user}@${hostname}`, 'echo "test connection succeeded" && exit'])
     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({
+        return await process.run({
             spawnOptions: {
                 env,
             },
         })
-        return result
     } catch (error) {
-        throw new SshError('SSH connection test failed', { cause: error as Error })
+        throw new SshError(process.result()?.stderr ?? 'An unknown error occurred when testing the connection', {
+            cause: error as Error,
+        })
     }
 }
 

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

@@ -125,6 +125,17 @@ describe('Ec2ConnectClient', function () {
             }
         })
 
+        it('retries if agent status is not online', async function () {
+            const instanceAgentStatus = sinon.stub(SsmClient.prototype, 'getInstanceAgentPingStatus')
+            instanceAgentStatus.onFirstCall().resolves('Offline')
+            instanceAgentStatus.onSecondCall().resolves('Online')
+            try {
+                await client.checkForInstanceSsmError(instanceSelection, { interval: 10, timeout: 100 })
+            } catch (err) {
+                assert.ok(false, `checkForInstanceSsmError failed with error '${err}'`)
+            }
+        })
+
         it('does not throw an error if all checks pass', async function () {
             sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(true)
             sinon.stub(Ec2Connecter.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role)