ec2: connect to terminal w/o error handling. (#3575)

* add empty connect command

* create sample QuickPick on click

* hacky prompt + func to list instanceId

* very hacky way to create prompt with instanceIds

* move extractInstanceIds to own function

* utilize lastTouchedRegion in prompt

* throw cancellationError on user cancel

* add a title to quickpick

* refactor: increase testability

* move test utility func to shared utility file

* add general test cases for extractInstanceIdsFromReservations

* add basic test for prompter

* configure command to devMode

* refactor to utilize wizard + integrate regionSubmenu

* add new testing file

* refactor to utilize regionSubmenu in isolation, rather than wrapped in Wizard class

* fix awkward indent

* add test file from feature/cwl branch

* remove old prompter old + tests

* delete tests that rely on feature/cwl changes

* refactor to avoid circular dependency

* fix improper headers + imports

* introduce Ec2ConnectClient to handle connection

* remove dead parameter

* close connection on terminal close

* add a little more to error msg

* remove dead imports

* log error to console

* rename function

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

* fix header

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

* fix headers and imports

* remove year from header

* delete outdated test case

* delete outdated test file

* remove year from copyright header

* fix formatting of files

* remove alias in Ec2Instance

* utilize interface for object-like shape

* move code to general Ec2Client

* change Log Groups to selection in region submenu

* refactor to avoid circular dependency

* fix formatting

* fix formatting

* comment out the ec2Client we are not using

* style fix

* generalize some code to a remoteSession file

* add space between pieces

* remove dead import

* refactor to avoid circular dependency

* change callback to async/await

* remove onError parameter and utilize .catch instead

* make error handling a try-catch

* avoid unnecessary default prefix on client name

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

* fix old import

* refactor defaultEc2client -> ec2Client

* rename Ec2ConnectClient

---------

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

Author: Hweinstock

src/ec2/activation.ts

@@ -9,7 +9,6 @@ import { tryConnect } from './commands'
 export async function activate(ctx: ExtContext): Promise<void> {
     ctx.extensionContext.subscriptions.push(
         Commands.register('aws.ec2.connectToInstance', async (param?: unknown) => {
-            console.log('You just ran the aws.ec2.connectToInstance command!')
             tryConnect()
         })
     )

src/ec2/commands.ts

@@ -3,15 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { createEC2ConnectPrompter, handleEc2ConnectPrompterResponse } from './prompter'
+import { createEc2ConnectPrompter, handleEc2ConnectPrompterResponse } from './prompter'
 import { isValidResponse } from '../shared/wizards/wizard'
+import { Ec2ConnectionManager } from './model'
 
 export async function tryConnect(): Promise<void> {
-    const prompter = createEC2ConnectPrompter()
+    const prompter = createEc2ConnectPrompter()
     const response = await prompter.prompt()
 
     if (isValidResponse(response)) {
         const selection = handleEc2ConnectPrompterResponse(response)
-        console.log(selection)
+        const ec2Client = new Ec2ConnectionManager(selection.region)
+        await ec2Client.attemptEc2Connection(selection)
     }
 }

src/ec2/model.ts

@@ -0,0 +1,58 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import * as vscode from 'vscode'
+import { Session } from 'aws-sdk/clients/ssm'
+import { Ec2Selection } from './utils'
+import { getOrInstallCli } from '../shared/utilities/cliUtils'
+import { isCloud9 } from '../shared/extensionUtilities'
+import { ToolkitError, isAwsError } from '../shared/errors'
+import { SsmClient } from '../shared/clients/ssmClient'
+import { openRemoteTerminal } from '../shared/remoteSession'
+
+export class Ec2ConnectionManager {
+    // Will need the ec2Client for probing errors,
+    private ssmClient: SsmClient
+    //private ec2Client: DefaultEc2Client
+
+    public constructor(readonly regionCode: string) {
+        this.ssmClient = new SsmClient(this.regionCode)
+        //this.ec2Client = new DefaultEc2Client(this.regionCode)
+    }
+
+    private async handleStartSessionError(err: AWS.AWSError): Promise<void> {
+        const failureMessage =
+            "SSM: Failed to start session with target instance. Common reasons include: \n 1. SSM Agent not installed on instance. \n 2. The required IAM instance profile isn't attached to the instance.  \n 3. Session manager setup is incomplete."
+        await vscode.window.showErrorMessage(failureMessage)
+
+        throw new ToolkitError('Start Session Failed. ')
+    }
+
+    private async openSessionInTerminal(session: Session, selection: Ec2Selection) {
+        const ssmPlugin = await getOrInstallCli('session-manager-plugin', !isCloud9)
+        const shellArgs = [JSON.stringify(session), selection.region, 'StartSession']
+        const terminalOptions = {
+            name: selection.region + '/' + selection.instanceId,
+            shellPath: ssmPlugin,
+            shellArgs: shellArgs,
+        }
+
+        await openRemoteTerminal(terminalOptions, () => this.ssmClient.terminateSession(session)).catch(err => {
+            throw ToolkitError.chain(err, 'Failed to open ec2 instance.')
+        })
+    }
+
+    public async attemptEc2Connection(selection: Ec2Selection): Promise<void> {
+        try {
+            const response = await this.ssmClient.startSession(selection.instanceId)
+            await this.openSessionInTerminal(response, selection)
+        } catch (err) {
+            if (isAwsError(err)) {
+                await this.handleStartSessionError(err)
+            } else {
+                throw err
+            }
+        }
+    }
+}

src/ec2/prompter.ts

@@ -4,29 +4,24 @@
  */
 
 import { RegionSubmenu, RegionSubmenuResponse } from '../shared/ui/common/regionSubmenu'
-import { getInstanceIdsFromRegion } from './utils'
+import { Ec2Selection, getInstanceIdsFromRegion } from './utils'
 import { DataQuickPickItem } from '../shared/ui/pickerPrompter'
 
-interface EC2Selection {
-    instanceId: string
-    region: string
-}
-
 function asQuickpickItem(instanceId: string): DataQuickPickItem<string> {
     return {
         label: instanceId,
         data: instanceId,
     }
 }
 
-export function handleEc2ConnectPrompterResponse(response: RegionSubmenuResponse<string>): EC2Selection {
+export function handleEc2ConnectPrompterResponse(response: RegionSubmenuResponse<string>): Ec2Selection {
     return {
         instanceId: response.data,
         region: response.region,
     }
 }
 
-export function createEC2ConnectPrompter(): RegionSubmenu<string> {
+export function createEc2ConnectPrompter(): RegionSubmenu<string> {
     return new RegionSubmenu(
         async region => (await getInstanceIdsFromRegion(region)).map(asQuickpickItem).promise(),
         { title: 'Select EC2 Instance Id' },

src/ec2/utils.ts

@@ -4,9 +4,14 @@
  */
 
 import { AsyncCollection } from '../shared/utilities/asyncCollection'
-import { DefaultEc2Client } from '../shared/clients/ec2Client'
+import { Ec2Client } from '../shared/clients/ec2Client'
+
+export interface Ec2Selection {
+    instanceId: string
+    region: string
+}
 
 export async function getInstanceIdsFromRegion(regionCode: string): Promise<AsyncCollection<string>> {
-    const client = new DefaultEc2Client(regionCode)
+    const client = new Ec2Client(regionCode)
     return client.getInstanceIds()
 }

src/ecs/commands.ts

@@ -11,7 +11,7 @@ import * as vscode from 'vscode'
 import { DefaultIamClient } from '../shared/clients/iamClient'
 import { INSIGHTS_TIMESTAMP_FORMAT } from '../shared/constants'
 import globals from '../shared/extensionGlobals'
-import { PromptSettings, Settings } from '../shared/settings'
+import { PromptSettings } from '../shared/settings'
 import { ChildProcess } from '../shared/utilities/childProcess'
 import { showMessageWithCancel, showOutputMessage } from '../shared/utilities/messages'
 import { removeAnsi } from '../shared/utilities/textUtilities'
@@ -24,6 +24,7 @@ import { getResourceFromTreeNode } from '../shared/treeview/utils'
 import { Container, Service } from './model'
 import { Instance } from '../shared/utilities/typeConstructors'
 import { telemetry } from '../shared/telemetry/telemetry'
+import { openRemoteTerminal } from '../shared/remoteSession'
 
 async function runCommandWizard(
     param?: unknown,
@@ -136,44 +137,22 @@ export const runCommandInContainer = Commands.register('aws.ecs.runCommandInCont
     })
 })
 
-// VSC is logging args to the PTY host log file if shell integration is enabled :(
-async function withoutShellIntegration<T>(cb: () => T | Promise<T>): Promise<T> {
-    const userValue = Settings.instance.get('terminal.integrated.shellIntegration.enabled', Boolean)
-
-    try {
-        await Settings.instance.update('terminal.integrated.shellIntegration.enabled', false)
-        return await cb()
-    } finally {
-        Settings.instance.update('terminal.integrated.shellIntegration.enabled', userValue)
-    }
-}
-
 export const openTaskInTerminal = Commands.register('aws.ecs.openTaskInTerminal', (obj?: unknown) => {
     return telemetry.ecs_runExecuteCommand.run(async span => {
         span.record({ ecsExecuteCommandType: 'shell' })
 
         const startCommand = new EcsSettings().get('openTerminalCommand')
         const { container, task, command } = await runCommandWizard(obj, startCommand)
+        const session = await container.prepareCommandForTask(command, task)
 
-        try {
-            const session = await container.prepareCommandForTask(command, task)
-            await withoutShellIntegration(() => {
-                const terminal = vscode.window.createTerminal({
-                    name: `${container.description.name}/${task}`,
-                    shellPath: session.path,
-                    shellArgs: session.args,
-                })
-
-                const listener = vscode.window.onDidCloseTerminal(t => {
-                    if (t.processId === terminal.processId) {
-                        vscode.Disposable.from(listener, session).dispose()
-                    }
-                })
+        const terminalOptions = {
+            name: `${container.description.name}/${task}`,
+            shellPath: session.path,
+            shellArgs: session.args,
+        }
 
-                terminal.show()
-            })
-        } catch (err) {
+        await openRemoteTerminal(terminalOptions, session.dispose).catch(err => {
             throw ToolkitError.chain(err, localize('AWS.ecs.openTaskInTerminal.error', 'Failed to open terminal.'))
-        }
+        })
     })
 })

src/shared/clients/ec2Client.ts

@@ -8,7 +8,7 @@ import globals from '../extensionGlobals'
 import { AsyncCollection } from '../utilities/asyncCollection'
 import { pageableToCollection } from '../utilities/collectionUtils'
 
-export class DefaultEc2Client {
+export class Ec2Client {
     public constructor(public readonly regionCode: string) {}
 
     private async createSdkClient(): Promise<EC2> {

src/shared/clients/ssmClient.ts

@@ -0,0 +1,37 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { AWSError, SSM } from 'aws-sdk'
+import globals from '../extensionGlobals'
+import { PromiseResult } from 'aws-sdk/lib/request'
+import { getLogger } from '../logger/logger'
+
+export class SsmClient {
+    public constructor(public readonly regionCode: string) {}
+
+    private async createSdkClient(): Promise<SSM> {
+        return await globals.sdkClientBuilder.createAwsService(SSM, undefined, this.regionCode)
+    }
+
+    public async terminateSession(
+        session: SSM.Session
+    ): Promise<void | PromiseResult<SSM.TerminateSessionResponse, AWSError>> {
+        const sessionId = session.SessionId!
+        const client = await this.createSdkClient()
+        const termination = await client
+            .terminateSession({ SessionId: sessionId })
+            .promise()
+            .catch(err => {
+                getLogger().warn(`ssm: failed to terminate session "${sessionId}": %s`, err)
+            })
+        return termination
+    }
+
+    public async startSession(target: string): Promise<PromiseResult<SSM.StartSessionResponse, AWSError>> {
+        const client = await this.createSdkClient()
+        const response = await client.startSession({ Target: target }).promise()
+        return response
+    }
+}

src/shared/remoteSession.ts

@@ -0,0 +1,33 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { Settings } from '../shared/settings'
+
+export async function openRemoteTerminal(options: vscode.TerminalOptions, onClose: () => void) {
+    await withoutShellIntegration(() => {
+        const terminal = vscode.window.createTerminal(options)
+
+        const listener = vscode.window.onDidCloseTerminal(t => {
+            if (t.processId === terminal.processId) {
+                vscode.Disposable.from(listener, { dispose: onClose }).dispose()
+            }
+        })
+
+        terminal.show()
+    })
+}
+
+// VSC is logging args to the PTY host log file if shell integration is enabled :(
+async function withoutShellIntegration<T>(cb: () => T | Promise<T>): Promise<T> {
+    const userValue = Settings.instance.get('terminal.integrated.shellIntegration.enabled', Boolean)
+
+    try {
+        await Settings.instance.update('terminal.integrated.shellIntegration.enabled', false)
+        return await cb()
+    } finally {
+        Settings.instance.update('terminal.integrated.shellIntegration.enabled', userValue)
+    }
+}

src/test/shared/clients/defaultEc2Client.test.ts

@@ -8,10 +8,10 @@ import { AsyncCollection } from '../../../shared/utilities/asyncCollection'
 import { EC2 } from 'aws-sdk'
 import { toCollection } from '../../../shared/utilities/asyncCollection'
 import { intoCollection } from '../../../shared/utilities/collectionUtils'
-import { DefaultEc2Client } from '../../../shared/clients/ec2Client'
+import { Ec2Client } from '../../../shared/clients/ec2Client'
 
 describe('extractInstanceIdsFromReservations', function () {
-    const client = new DefaultEc2Client('')
+    const client = new Ec2Client('')
     it('returns empty when given empty collection', async function () {
         const actualResult = await client
             .extractInstanceIdsFromReservations(