diff --git a/packages/core/src/awsService/ec2/activation.ts b/packages/core/src/awsService/ec2/activation.ts index 45d0a5c369e..aa57b9ce79d 100644 --- a/packages/core/src/awsService/ec2/activation.ts +++ b/packages/core/src/awsService/ec2/activation.ts @@ -6,7 +6,7 @@ import * as vscode from 'vscode' import { ExtContext } from '../../shared/extensions' import { Commands } from '../../shared/vscode/commands2' import { telemetry } from '../../shared/telemetry/telemetry' -import { Ec2InstanceNode } from './explorer/ec2InstanceNode' +import { Ec2InstanceNode, tryRefreshNode } from './explorer/ec2InstanceNode' import { copyTextCommand } from '../../awsexplorer/commands/copyText' import { Ec2Node } from './explorer/ec2ParentNode' import { @@ -15,13 +15,15 @@ import { rebootInstance, startInstance, stopInstance, - refreshExplorer, - openLogDocument, linkToLaunchInstance, + openLogDocument, } from './commands' +import { Ec2ConnecterMap } from './connectionManagerMap' import { ec2LogsScheme } from '../../shared/constants' import { Ec2LogDocumentProvider } from './ec2LogDocumentProvider' +const connectionManagers = new Ec2ConnecterMap() + export async function activate(ctx: ExtContext): Promise { ctx.extensionContext.subscriptions.push( vscode.workspace.registerTextDocumentContentProvider(ec2LogsScheme, new Ec2LogDocumentProvider()) @@ -30,7 +32,7 @@ export async function activate(ctx: ExtContext): Promise { Commands.register('aws.ec2.openTerminal', async (node?: Ec2InstanceNode) => { await telemetry.ec2_connectToInstance.run(async (span) => { span.record({ ec2ConnectionType: 'ssm' }) - await openTerminal(node) + await openTerminal(connectionManagers, node) }) }), @@ -42,14 +44,14 @@ export async function activate(ctx: ExtContext): Promise { }), Commands.register('aws.ec2.openRemoteConnection', async (node?: Ec2Node) => { - await openRemoteConnection(node) + await openRemoteConnection(connectionManagers, node) }), Commands.register('aws.ec2.startInstance', async (node?: Ec2Node) => { await telemetry.ec2_changeState.run(async (span) => { span.record({ ec2InstanceState: 'start' }) await startInstance(node) - refreshExplorer(node) + await tryRefreshNode(node) }) }), @@ -57,7 +59,7 @@ export async function activate(ctx: ExtContext): Promise { await telemetry.ec2_changeState.run(async (span) => { span.record({ ec2InstanceState: 'stop' }) await stopInstance(node) - refreshExplorer(node) + await tryRefreshNode(node) }) }), @@ -65,7 +67,7 @@ export async function activate(ctx: ExtContext): Promise { await telemetry.ec2_changeState.run(async (span) => { span.record({ ec2InstanceState: 'reboot' }) await rebootInstance(node) - refreshExplorer(node) + await tryRefreshNode(node) }) }), @@ -76,3 +78,7 @@ export async function activate(ctx: ExtContext): Promise { }) ) } + +export async function deactivate(): Promise { + connectionManagers.forEach(async (manager) => await manager.dispose()) +} diff --git a/packages/core/src/awsService/ec2/commands.ts b/packages/core/src/awsService/ec2/commands.ts index 01cdb5b5f53..4f06f7876eb 100644 --- a/packages/core/src/awsService/ec2/commands.ts +++ b/packages/core/src/awsService/ec2/commands.ts @@ -4,36 +4,25 @@ */ import { Ec2InstanceNode } from './explorer/ec2InstanceNode' import { Ec2Node } from './explorer/ec2ParentNode' -import { Ec2ConnectionManager } from './model' -import { Ec2Prompter, instanceFilter, Ec2Selection } from './prompter' import { SafeEc2Instance, Ec2Client } from '../../shared/clients/ec2Client' import { copyToClipboard } from '../../shared/utilities/messages' -import { getLogger } from '../../shared/logger' import { ec2LogSchema } from './ec2LogDocumentProvider' import { getAwsConsoleUrl } from '../../shared/awsConsole' import { showRegionPrompter } from '../../auth/utils' import { openUrl } from '../../shared/utilities/vsCodeUtils' import { showFile } from '../../shared/utilities/textDocumentUtilities' +import { Ec2ConnecterMap } from './connectionManagerMap' +import { Ec2Prompter, Ec2Selection, instanceFilter } from './prompter' -export function refreshExplorer(node?: Ec2Node) { - if (node) { - const n = node instanceof Ec2InstanceNode ? node.parent : node - n.refreshNode().catch((e) => { - getLogger().error('refreshNode failed: %s', (e as Error).message) - }) - } -} - -export async function openTerminal(node?: Ec2Node) { +export async function openTerminal(connectionManagers: Ec2ConnecterMap, node?: Ec2Node) { const selection = await getSelection(node) - - const connectionManager = new Ec2ConnectionManager(selection.region) + const connectionManager = connectionManagers.getOrInit(selection.region) await connectionManager.attemptToOpenEc2Terminal(selection) } -export async function openRemoteConnection(node?: Ec2Node) { +export async function openRemoteConnection(connectionManagers: Ec2ConnecterMap, node?: Ec2Node) { const selection = await getSelection(node) - const connectionManager = new Ec2ConnectionManager(selection.region) + const connectionManager = connectionManagers.getOrInit(selection.region) await connectionManager.tryOpenRemoteConnection(selection) } diff --git a/packages/core/src/awsService/ec2/connectionManagerMap.ts b/packages/core/src/awsService/ec2/connectionManagerMap.ts new file mode 100644 index 00000000000..add39f05bea --- /dev/null +++ b/packages/core/src/awsService/ec2/connectionManagerMap.ts @@ -0,0 +1,26 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { getLogger } from '../../shared' +import { Ec2Connecter } from './model' + +export class Ec2ConnecterMap extends Map { + private static warnSize: number = 25 + + public getOrInit(regionCode: string) { + return this.has(regionCode) ? this.get(regionCode)! : this.initManager(regionCode) + } + + private initManager(regionCode: string): Ec2Connecter { + if (this.size >= Ec2ConnecterMap.warnSize) { + getLogger().warn( + `Connection manager exceeded threshold of ${Ec2ConnecterMap.warnSize} with ${this.size} active connections` + ) + } + const newConnectionManager = new Ec2Connecter(regionCode) + this.set(regionCode, newConnectionManager) + return newConnectionManager + } +} diff --git a/packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts b/packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts index 7078b6bdb02..d00fe55e9c2 100644 --- a/packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts +++ b/packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts @@ -10,8 +10,9 @@ import { SafeEc2Instance } from '../../../shared/clients/ec2Client' import globals from '../../../shared/extensionGlobals' import { getIconCode } from '../utils' import { Ec2Selection } from '../prompter' -import { Ec2ParentNode } from './ec2ParentNode' +import { Ec2Node, Ec2ParentNode } from './ec2ParentNode' import { EC2 } from 'aws-sdk' +import { getLogger } from '../../../shared' export const Ec2InstanceRunningContext = 'awsEc2RunningNode' export const Ec2InstanceStoppedContext = 'awsEc2StoppedNode' @@ -101,3 +102,14 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', this) } } + +export async function tryRefreshNode(node?: Ec2Node) { + if (node) { + const n = node instanceof Ec2InstanceNode ? node.parent : node + try { + await n.refreshNode() + } catch (e) { + getLogger().error('refreshNode failed: %s', (e as Error).message) + } + } +} diff --git a/packages/core/src/awsService/ec2/model.ts b/packages/core/src/awsService/ec2/model.ts index f47386ef47d..fa7bbee71b7 100644 --- a/packages/core/src/awsService/ec2/model.ts +++ b/packages/core/src/awsService/ec2/model.ts @@ -4,7 +4,7 @@ */ import * as vscode from 'vscode' import { Session } from 'aws-sdk/clients/ssm' -import { IAM, SSM } from 'aws-sdk' +import { EC2, IAM, SSM } from 'aws-sdk' import { Ec2Selection } from './prompter' import { getOrInstallCli } from '../../shared/utilities/cliUtils' import { isCloud9 } from '../../shared/extensionUtilities' @@ -25,20 +25,24 @@ import { createBoundProcess } from '../../codecatalyst/model' import { getLogger } from '../../shared/logger/logger' import { CancellationError, Timeout } from '../../shared/utilities/timeoutUtils' import { showMessageWithCancel } from '../../shared/utilities/messages' -import { SshConfig, sshLogFileLocation } from '../../shared/sshConfig' +import { SshConfig } from '../../shared/sshConfig' import { SshKeyPair } from './sshKeyPair' +import { Ec2SessionTracker } from './remoteSessionManager' +import { getEc2SsmEnv } from './utils' export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus' -interface Ec2RemoteEnv extends VscodeRemoteConnection { +export interface Ec2RemoteEnv extends VscodeRemoteConnection { selection: Ec2Selection keyPair: SshKeyPair + ssmSession: SSM.StartSessionResponse } -export class Ec2ConnectionManager { +export class Ec2Connecter implements vscode.Disposable { protected ssmClient: SsmClient protected ec2Client: Ec2Client protected iamClient: DefaultIamClient + protected sessionManager: Ec2SessionTracker private policyDocumentationUri = vscode.Uri.parse( 'https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-getting-started-instance-profile.html' @@ -52,6 +56,7 @@ export class Ec2ConnectionManager { this.ssmClient = this.createSsmSdkClient() this.ec2Client = this.createEc2SdkClient() this.iamClient = this.createIamSdkClient() + this.sessionManager = new Ec2SessionTracker(regionCode, this.ssmClient) } protected createSsmSdkClient(): SsmClient { @@ -66,6 +71,18 @@ export class Ec2ConnectionManager { return new DefaultIamClient(this.regionCode) } + public async addActiveSession(sessionId: SSM.SessionId, instanceId: EC2.InstanceId): Promise { + await this.sessionManager.addSession(instanceId, sessionId) + } + + public async dispose(): Promise { + await this.sessionManager.dispose() + } + + public isConnectedTo(instanceId: string): boolean { + return this.sessionManager.isConnectedTo(instanceId) + } + public async getAttachedIamRole(instanceId: string): Promise { const IamInstanceProfile = await this.ec2Client.getAttachedIamInstanceProfile(instanceId) if (IamInstanceProfile && IamInstanceProfile.Arn) { @@ -183,6 +200,7 @@ export class Ec2ConnectionManager { this.throwGeneralConnectionError(selection, err as Error) } } + public async prepareEc2RemoteEnvWithProgress(selection: Ec2Selection, remoteUser: string): Promise { const timeout = new Timeout(60000) await showMessageWithCancel('AWS: Opening remote connection...', timeout) @@ -204,8 +222,10 @@ export class Ec2ConnectionManager { throw err } - const session = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession') - const vars = getEc2SsmEnv(selection, ssm, session) + const ssmSession = await this.ssmClient.startSession(selection.instanceId, 'AWS-StartSSHSession') + await this.addActiveSession(selection.instanceId, ssmSession.SessionId!) + + const vars = getEc2SsmEnv(selection, ssm, ssmSession) const envProvider = async () => { return { [sshAgentSocketVariable]: await startSshAgent(), ...vars } } @@ -223,6 +243,7 @@ export class Ec2ConnectionManager { SessionProcess, selection, keyPair, + ssmSession, } } @@ -267,17 +288,3 @@ export class Ec2ConnectionManager { throw new ToolkitError(`Unrecognized OS name ${osName} on instance ${instanceId}`, { code: 'UnknownEc2OS' }) } } - -function getEc2SsmEnv(selection: Ec2Selection, ssmPath: string, session: SSM.StartSessionResponse): NodeJS.ProcessEnv { - return Object.assign( - { - AWS_REGION: selection.region, - AWS_SSM_CLI: ssmPath, - LOG_FILE_LOCATION: sshLogFileLocation('ec2', selection.instanceId), - STREAM_URL: session.StreamUrl, - SESSION_ID: session.SessionId, - TOKEN: session.TokenValue, - }, - process.env - ) -} diff --git a/packages/core/src/awsService/ec2/prompter.ts b/packages/core/src/awsService/ec2/prompter.ts index 6bbc301e515..638458942e2 100644 --- a/packages/core/src/awsService/ec2/prompter.ts +++ b/packages/core/src/awsService/ec2/prompter.ts @@ -10,6 +10,8 @@ import { isValidResponse } from '../../shared/wizards/wizard' import { CancellationError } from '../../shared/utilities/timeoutUtils' import { AsyncCollection } from '../../shared/utilities/asyncCollection' import { getIconCode } from './utils' +import { Ec2Node } from './explorer/ec2ParentNode' +import { Ec2InstanceNode } from './explorer/ec2InstanceNode' export type instanceFilter = (instance: SafeEc2Instance) => boolean export interface Ec2Selection { @@ -72,3 +74,9 @@ export class Ec2Prompter { ) } } + +export async function getSelection(node?: Ec2Node, filter?: instanceFilter): Promise { + const prompter = new Ec2Prompter(filter) + const selection = node && node instanceof Ec2InstanceNode ? node.toSelection() : await prompter.promptUser() + return selection +} diff --git a/packages/core/src/awsService/ec2/remoteSessionManager.ts b/packages/core/src/awsService/ec2/remoteSessionManager.ts new file mode 100644 index 00000000000..4c1843aabdb --- /dev/null +++ b/packages/core/src/awsService/ec2/remoteSessionManager.ts @@ -0,0 +1,40 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EC2, SSM } from 'aws-sdk' +import { SsmClient } from '../../shared/clients/ssmClient' +import { Disposable } from 'vscode' + +export class Ec2SessionTracker extends Map implements Disposable { + public constructor( + readonly regionCode: string, + protected ssmClient: SsmClient + ) { + super() + } + + public async addSession(instanceId: EC2.InstanceId, sessionId: SSM.SessionId): Promise { + if (this.isConnectedTo(instanceId)) { + const existingSessionId = this.get(instanceId)! + await this.ssmClient.terminateSessionFromId(existingSessionId) + this.set(instanceId, sessionId) + } else { + this.set(instanceId, sessionId) + } + } + + private async disconnectEnv(instanceId: EC2.InstanceId): Promise { + await this.ssmClient.terminateSessionFromId(this.get(instanceId)!) + this.delete(instanceId) + } + + public async dispose(): Promise { + this.forEach(async (_sessionId, instanceId) => await this.disconnectEnv(instanceId)) + } + + public isConnectedTo(instanceId: EC2.InstanceId): boolean { + return this.has(instanceId) + } +} diff --git a/packages/core/src/awsService/ec2/utils.ts b/packages/core/src/awsService/ec2/utils.ts index f9272f5cee1..cd0e374629e 100644 --- a/packages/core/src/awsService/ec2/utils.ts +++ b/packages/core/src/awsService/ec2/utils.ts @@ -4,6 +4,10 @@ */ import { SafeEc2Instance } from '../../shared/clients/ec2Client' +import { copyToClipboard } from '../../shared/utilities/messages' +import { Ec2Selection } from './prompter' +import { sshLogFileLocation } from '../../shared/sshConfig' +import { SSM } from 'aws-sdk' export function getIconCode(instance: SafeEc2Instance) { if (instance.LastSeenStatus === 'running') { @@ -16,3 +20,25 @@ export function getIconCode(instance: SafeEc2Instance) { return 'loading~spin' } + +export async function copyInstanceId(instanceId: string): Promise { + await copyToClipboard(instanceId, 'Id') +} + +export function getEc2SsmEnv( + selection: Ec2Selection, + ssmPath: string, + session: SSM.StartSessionResponse +): NodeJS.ProcessEnv { + return Object.assign( + { + AWS_REGION: selection.region, + AWS_SSM_CLI: ssmPath, + LOG_FILE_LOCATION: sshLogFileLocation('ec2', selection.instanceId), + STREAM_URL: session.StreamUrl, + SESSION_ID: session.SessionId, + TOKEN: session.TokenValue, + }, + process.env + ) +} diff --git a/packages/core/src/extensionNode.ts b/packages/core/src/extensionNode.ts index 06185f695cc..f5e580300e8 100644 --- a/packages/core/src/extensionNode.ts +++ b/packages/core/src/extensionNode.ts @@ -24,7 +24,7 @@ import { } from './shared/extensionUtilities' import { getLogger, Logger } from './shared/logger/logger' import { activate as activateEcr } from './awsService/ecr/activation' -import { activate as activateEc2 } from './awsService/ec2/activation' +import { activate as activateEc2, deactivate as deactivateEc2 } from './awsService/ec2/activation' import { activate as activateSam } from './shared/sam/activation' import { activate as activateS3 } from './awsService/s3/activation' import * as filetypes from './shared/filetypes' @@ -255,7 +255,7 @@ export async function activate(context: vscode.ExtensionContext) { export async function deactivate() { // Run concurrently to speed up execution. stop() does not throw so it is safe - await Promise.all([await (await CrashMonitoring.instance())?.shutdown(), deactivateCommon()]) + await Promise.all([await (await CrashMonitoring.instance())?.shutdown(), deactivateCommon(), deactivateEc2()]) await globals.resourceManager.dispose() } diff --git a/packages/core/src/shared/clients/ssmClient.ts b/packages/core/src/shared/clients/ssmClient.ts index b023b5bb0c0..cf5741928ed 100644 --- a/packages/core/src/shared/clients/ssmClient.ts +++ b/packages/core/src/shared/clients/ssmClient.ts @@ -19,6 +19,10 @@ export class SsmClient { public async terminateSession(session: SSM.Session): Promise { const sessionId = session.SessionId! + return await this.terminateSessionFromId(sessionId) + } + + public async terminateSessionFromId(sessionId: SSM.SessionId): Promise { const client = await this.createSdkClient() const termination = await client .terminateSession({ SessionId: sessionId }) @@ -101,4 +105,13 @@ export class SsmClient { const instanceInformation = await this.describeInstance(target) return instanceInformation ? instanceInformation.PingStatus! : 'Inactive' } + + public async describeSessions(state: SSM.SessionState) { + const client = await this.createSdkClient() + const requester = async (req: SSM.DescribeSessionsRequest) => client.describeSessions(req).promise() + + const response = await pageableToCollection(requester, { State: state }, 'NextToken', 'Sessions').promise() + + return response + } } diff --git a/packages/core/src/test/awsService/ec2/commands.test.ts b/packages/core/src/test/awsService/ec2/commands.test.ts new file mode 100644 index 00000000000..67974692bd2 --- /dev/null +++ b/packages/core/src/test/awsService/ec2/commands.test.ts @@ -0,0 +1,32 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import assert from 'assert' +import { Ec2Selection } from '../../../awsService/ec2/prompter' +import { Ec2ConnecterMap } from '../../../awsService/ec2/connectionManagerMap' + +describe('getConnectionManager', async function () { + let connectionManagers: Ec2ConnecterMap + + beforeEach(function () { + connectionManagers = new Ec2ConnecterMap() + }) + + it('only creates new connection managers once for each region ', async function () { + const fakeSelection: Ec2Selection = { + region: 'region-1', + instanceId: 'fake-id', + } + + const cm = connectionManagers.getOrInit(fakeSelection.region) + assert.strictEqual(connectionManagers.size, 1) + + await cm.addActiveSession('sessionId', 'instanceId') + + const cm2 = connectionManagers.getOrInit(fakeSelection.region) + + assert.strictEqual(cm2.isConnectedTo('instanceId'), true) + }) +}) diff --git a/packages/core/src/test/awsService/ec2/explorer/ec2ParentNode.test.ts b/packages/core/src/test/awsService/ec2/explorer/ec2ParentNode.test.ts index 788c97db046..67d801a4ea4 100644 --- a/packages/core/src/test/awsService/ec2/explorer/ec2ParentNode.test.ts +++ b/packages/core/src/test/awsService/ec2/explorer/ec2ParentNode.test.ts @@ -18,6 +18,19 @@ import { AsyncCollection } from '../../../../shared/utilities/asyncCollection' import * as FakeTimers from '@sinonjs/fake-timers' import { installFakeClock } from '../../../testUtil' +export const testInstance = { + InstanceId: 'testId', + Tags: [ + { + Key: 'Name', + Value: 'testName', + }, + ], + LastSeenStatus: 'running', +} +export const testClient = new Ec2Client('') +export const testParentNode = new Ec2ParentNode('fake-region', 'testPartition', testClient) + describe('ec2ParentNode', function () { let testNode: Ec2ParentNode let client: Ec2Client diff --git a/packages/core/src/test/awsService/ec2/model.test.ts b/packages/core/src/test/awsService/ec2/model.test.ts index f2d5c8728c0..6c7f6f1fb4a 100644 --- a/packages/core/src/test/awsService/ec2/model.test.ts +++ b/packages/core/src/test/awsService/ec2/model.test.ts @@ -5,7 +5,7 @@ import assert from 'assert' import * as sinon from 'sinon' -import { Ec2ConnectionManager } from '../../../awsService/ec2/model' +import { Ec2Connecter } from '../../../awsService/ec2/model' import { SsmClient } from '../../../shared/clients/ssmClient' import { Ec2Client } from '../../../shared/clients/ec2Client' import { Ec2Selection } from '../../../awsService/ec2/prompter' @@ -17,10 +17,10 @@ import { assertNoTelemetryMatch, createTestWorkspaceFolder } from '../../testUti import { fs } from '../../../shared' describe('Ec2ConnectClient', function () { - let client: Ec2ConnectionManager + let client: Ec2Connecter before(function () { - client = new Ec2ConnectionManager('test-region') + client = new Ec2Connecter('test-region') }) describe('getAttachedIamRole', async function () { @@ -84,7 +84,7 @@ describe('Ec2ConnectClient', function () { }) it('throws EC2SSMStatus error if instance is not running', async function () { - sinon.stub(Ec2ConnectionManager.prototype, 'isInstanceRunning').resolves(false) + sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(false) try { await client.checkForStartSessionError(instanceSelection) @@ -95,8 +95,8 @@ describe('Ec2ConnectClient', function () { }) it('throws EC2SSMPermission error if instance is running but has no role', async function () { - sinon.stub(Ec2ConnectionManager.prototype, 'isInstanceRunning').resolves(true) - sinon.stub(Ec2ConnectionManager.prototype, 'getAttachedIamRole').resolves(undefined) + sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(true) + sinon.stub(Ec2Connecter.prototype, 'getAttachedIamRole').resolves(undefined) try { await client.checkForStartSessionError(instanceSelection) @@ -107,9 +107,9 @@ describe('Ec2ConnectClient', function () { }) it('throws EC2SSMAgent error if instance is running and has IAM Role, but agent is not running', async function () { - sinon.stub(Ec2ConnectionManager.prototype, 'isInstanceRunning').resolves(true) - sinon.stub(Ec2ConnectionManager.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role) - sinon.stub(Ec2ConnectionManager.prototype, 'hasProperPermissions').resolves(true) + sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(true) + sinon.stub(Ec2Connecter.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role) + sinon.stub(Ec2Connecter.prototype, 'hasProperPermissions').resolves(true) sinon.stub(SsmClient.prototype, 'getInstanceAgentPingStatus').resolves('offline') try { @@ -121,9 +121,9 @@ describe('Ec2ConnectClient', function () { }) it('does not throw an error if all checks pass', async function () { - sinon.stub(Ec2ConnectionManager.prototype, 'isInstanceRunning').resolves(true) - sinon.stub(Ec2ConnectionManager.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role) - sinon.stub(Ec2ConnectionManager.prototype, 'hasProperPermissions').resolves(true) + sinon.stub(Ec2Connecter.prototype, 'isInstanceRunning').resolves(true) + sinon.stub(Ec2Connecter.prototype, 'getAttachedIamRole').resolves({ Arn: 'testRole' } as IAM.Role) + sinon.stub(Ec2Connecter.prototype, 'hasProperPermissions').resolves(true) sinon.stub(SsmClient.prototype, 'getInstanceAgentPingStatus').resolves('Online') assert.doesNotThrow(async () => await client.checkForStartSessionError(instanceSelection)) diff --git a/packages/core/src/test/awsService/ec2/prompter.test.ts b/packages/core/src/test/awsService/ec2/prompter.test.ts index b6de134d182..9eece4b44c9 100644 --- a/packages/core/src/test/awsService/ec2/prompter.test.ts +++ b/packages/core/src/test/awsService/ec2/prompter.test.ts @@ -3,13 +3,16 @@ * SPDX-License-Identifier: Apache-2.0 */ import assert from 'assert' -import { Ec2Prompter, instanceFilter } from '../../../awsService/ec2/prompter' +import * as sinon from 'sinon' +import { Ec2Prompter, getSelection, instanceFilter } from '../../../awsService/ec2/prompter' import { SafeEc2Instance } from '../../../shared/clients/ec2Client' import { RegionSubmenuResponse } from '../../../shared/ui/common/regionSubmenu' import { Ec2Selection } from '../../../awsService/ec2/prompter' import { AsyncCollection } from '../../../shared/utilities/asyncCollection' import { intoCollection } from '../../../shared/utilities/collectionUtils' import { DataQuickPickItem } from '../../../shared/ui/pickerPrompter' +import { Ec2InstanceNode } from '../../../awsService/ec2/explorer/ec2InstanceNode' +import { testClient, testInstance, testParentNode } from './explorer/ec2ParentNode.test' describe('Ec2Prompter', async function () { class MockEc2Prompter extends Ec2Prompter { @@ -183,4 +186,30 @@ describe('Ec2Prompter', async function () { assert.deepStrictEqual(items, expected) }) }) + + describe('getSelection', async function () { + it('uses node when passed', async function () { + const prompterStub = sinon.stub(Ec2Prompter.prototype, 'promptUser') + const testNode = new Ec2InstanceNode( + testParentNode, + testClient, + 'testRegion', + 'testPartition', + testInstance + ) + const result = await getSelection(testNode) + + assert.strictEqual(result.instanceId, testNode.toSelection().instanceId) + assert.strictEqual(result.region, testNode.toSelection().region) + sinon.assert.notCalled(prompterStub) + prompterStub.restore() + }) + + it('prompts user when no node is passed', async function () { + const prompterStub = sinon.stub(Ec2Prompter.prototype, 'promptUser') + await getSelection() + sinon.assert.calledOnce(prompterStub) + prompterStub.restore() + }) + }) }) diff --git a/packages/core/src/test/awsService/ec2/remoteSessionManager.test.ts b/packages/core/src/test/awsService/ec2/remoteSessionManager.test.ts new file mode 100644 index 00000000000..d125d4918dc --- /dev/null +++ b/packages/core/src/test/awsService/ec2/remoteSessionManager.test.ts @@ -0,0 +1,56 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import assert from 'assert' +import * as sinon from 'sinon' +import { Ec2SessionTracker } from '../../../awsService/ec2/remoteSessionManager' +import { SsmClient } from '../../../shared/clients/ssmClient' + +describe('Ec2RemoteSessionManager', async function () { + it('maintains connections to instances', async function () { + const envManager = new Ec2SessionTracker('test-region', new SsmClient('test-region')) + await envManager.addSession('test-instance', 'test-env') + await envManager.addSession('test-instance2', 'test-env2') + await envManager.addSession('test-instance3', 'test-env3') + + assert(envManager.isConnectedTo('test-instance')) + assert(envManager.isConnectedTo('test-instance2')) + assert(envManager.isConnectedTo('test-instance3')) + assert(!envManager.isConnectedTo('test-instance4')) + }) + + it('only allows for single connection with any given instance', async function () { + const envManager = new Ec2SessionTracker('test-region', new SsmClient('test-region')) + const terminateStub = sinon.stub(SsmClient.prototype, 'terminateSessionFromId') + + await envManager.addSession('test-instance', 'test-env') + sinon.assert.notCalled(terminateStub) + await envManager.addSession('test-instance', 'test-env2') + + sinon.assert.calledWith(terminateStub, 'test-env') + + assert(envManager.isConnectedTo('test-instance')) + + terminateStub.restore() + }) + + it('closes all active connections', async function () { + const envManager = new Ec2SessionTracker('test-region', new SsmClient('test-region')) + const terminateStub = sinon.stub(SsmClient.prototype, 'terminateSessionFromId') + + await envManager.addSession('test-instance', 'test-env') + await envManager.addSession('test-instance2', 'test-env2') + await envManager.addSession('test-instance3', 'test-env3') + + await envManager.dispose() + + sinon.assert.calledThrice(terminateStub) + assert(!envManager.isConnectedTo('test-instance')) + assert(!envManager.isConnectedTo('test-instance2')) + assert(!envManager.isConnectedTo('test-instance3')) + + terminateStub.restore() + }) +}) diff --git a/packages/core/src/test/awsService/ec2/utils.test.ts b/packages/core/src/test/awsService/ec2/utils.test.ts new file mode 100644 index 00000000000..6ee6ee65868 --- /dev/null +++ b/packages/core/src/test/awsService/ec2/utils.test.ts @@ -0,0 +1,50 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import assert from 'assert' +import * as sinon from 'sinon' +import { SafeEc2Instance } from '../../../shared/clients/ec2Client' +import { getIconCode } from '../../../awsService/ec2/utils' +import { DefaultAwsContext } from '../../../shared' + +describe('utils', async function () { + before(function () { + sinon.stub(DefaultAwsContext.prototype, 'getCredentialAccountId') + }) + + after(function () { + sinon.restore() + }) + + describe('getIconCode', function () { + it('gives code based on status', function () { + const runningInstance: SafeEc2Instance = { + InstanceId: 'X', + LastSeenStatus: 'running', + } + const stoppedInstance: SafeEc2Instance = { + InstanceId: 'XX', + LastSeenStatus: 'stopped', + } + + assert.strictEqual(getIconCode(runningInstance), 'pass') + assert.strictEqual(getIconCode(stoppedInstance), 'circle-slash') + }) + + it('defaults to loading~spin', function () { + const pendingInstance: SafeEc2Instance = { + InstanceId: 'X', + LastSeenStatus: 'pending', + } + const stoppingInstance: SafeEc2Instance = { + InstanceId: 'XX', + LastSeenStatus: 'shutting-down', + } + + assert.strictEqual(getIconCode(pendingInstance), 'loading~spin') + assert.strictEqual(getIconCode(stoppingInstance), 'loading~spin') + }) + }) +}) diff --git a/packages/toolkit/.changes/next-release/Bug Fix-1e058d1a-002a-4c7f-95eb-e8b676949d66.json b/packages/toolkit/.changes/next-release/Bug Fix-1e058d1a-002a-4c7f-95eb-e8b676949d66.json new file mode 100644 index 00000000000..414762bd8c8 --- /dev/null +++ b/packages/toolkit/.changes/next-release/Bug Fix-1e058d1a-002a-4c7f-95eb-e8b676949d66.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "EC2 connect: improve management of SSM sessions to minimize chance they are left active." +}