refactor(sdkv3): migrate ssm client

packages/core/package.json

@@ -509,6 +509,7 @@
         "@aws-sdk/property-provider": "3.46.0",
         "@aws-sdk/smithy-client": "^3.46.0",
         "@aws-sdk/util-arn-parser": "^3.46.0",
+        "@aws-sdk/client-ssm": "^3.699.0",
         "@aws/mynah-ui": "^4.21.5",
         "@gerhobbelt/gitignore-parser": "^0.2.0-9",
         "@iarna/toml": "^2.2.5",

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

@@ -3,13 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import { Session } from 'aws-sdk/clients/ssm'
-import { EC2, IAM, SSM } from 'aws-sdk'
+import { EC2, IAM } from 'aws-sdk'
 import { Ec2Selection } from './prompter'
 import { getOrInstallCli } from '../../shared/utilities/cliUtils'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { ToolkitError } from '../../shared/errors'
-import { SsmClient } from '../../shared/clients/ssmClient'
+import { SSMWrapper } from '../../shared/clients/ssm'
 import { Ec2Client } from '../../shared/clients/ec2Client'
 import {
     VscodeRemoteConnection,
@@ -35,13 +34,14 @@ import { SshConfig } from '../../shared/sshConfig'
 import { SshKeyPair } from './sshKeyPair'
 import { Ec2SessionTracker } from './remoteSessionManager'
 import { getEc2SsmEnv } from './utils'
+import { Session, StartSessionResponse } from '@aws-sdk/client-ssm'
 
 export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus'
 
 export interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
     keyPair: SshKeyPair
-    ssmSession: SSM.StartSessionResponse
+    ssmSession: StartSessionResponse
 }
 
 export type Ec2OS = 'Amazon Linux' | 'Ubuntu' | 'macOS'
@@ -51,7 +51,7 @@ interface RemoteUser {
 }
 
 export class Ec2Connecter implements vscode.Disposable {
-    protected ssmClient: SsmClient
+    protected ssm: SSMWrapper
     protected ec2Client: Ec2Client
     protected iamClient: DefaultIamClient
     protected sessionManager: Ec2SessionTracker
@@ -65,14 +65,14 @@ export class Ec2Connecter implements vscode.Disposable {
     )
 
     public constructor(readonly regionCode: string) {
-        this.ssmClient = this.createSsmSdkClient()
+        this.ssm = this.createSsmSdkClient()
         this.ec2Client = this.createEc2SdkClient()
         this.iamClient = this.createIamSdkClient()
-        this.sessionManager = new Ec2SessionTracker(regionCode, this.ssmClient)
+        this.sessionManager = new Ec2SessionTracker(regionCode, this.ssm)
     }
 
-    protected createSsmSdkClient(): SsmClient {
-        return new SsmClient(this.regionCode)
+    protected createSsmSdkClient(): SSMWrapper {
+        return new SSMWrapper(this.regionCode)
     }
 
     protected createEc2SdkClient(): Ec2Client {
@@ -83,7 +83,7 @@ export class Ec2Connecter implements vscode.Disposable {
         return new DefaultIamClient(this.regionCode)
     }
 
-    public async addActiveSession(sessionId: SSM.SessionId, instanceId: EC2.InstanceId): Promise<void> {
+    public async addActiveSession(sessionId: string, instanceId: EC2.InstanceId): Promise<void> {
         await this.sessionManager.addSession(instanceId, sessionId)
     }
 
@@ -151,7 +151,7 @@ export class Ec2Connecter implements vscode.Disposable {
     }
 
     private async checkForInstanceSsmError(selection: Ec2Selection): Promise<void> {
-        const isSsmAgentRunning = (await this.ssmClient.getInstanceAgentPingStatus(selection.instanceId)) === 'Online'
+        const isSsmAgentRunning = (await this.ssm.getInstanceAgentPingStatus(selection.instanceId)) === 'Online'
 
         if (!isSsmAgentRunning) {
             this.throwConnectionError('Is SSM Agent running on the target instance?', selection, {
@@ -178,15 +178,15 @@ export class Ec2Connecter implements vscode.Disposable {
             shellArgs: shellArgs,
         }
 
-        await openRemoteTerminal(terminalOptions, () => this.ssmClient.terminateSession(session)).catch((err) => {
+        await openRemoteTerminal(terminalOptions, () => this.ssm.terminateSession(session)).catch((err) => {
             throw ToolkitError.chain(err, 'Failed to open ec2 instance.')
         })
     }
 
     public async attemptToOpenEc2Terminal(selection: Ec2Selection): Promise<void> {
         await this.checkForStartSessionError(selection)
         try {
-            const response = await this.ssmClient.startSession(selection.instanceId)
+            const response = await this.ssm.startSession(selection.instanceId)
             await this.openSessionInTerminal(response, selection)
         } catch (err: unknown) {
             this.throwConnectionError('', selection, err as Error)
@@ -198,7 +198,7 @@ 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')
+        const testSession = await this.ssm.startSession(selection.instanceId, 'AWS-StartSSHSession')
         try {
             await testSshConnection(
                 remoteEnv.SessionProcess,
@@ -218,7 +218,7 @@ export class Ec2Connecter implements vscode.Disposable {
             const message = err instanceof SshError ? 'Testing SSH connection to instance failed' : ''
             this.throwConnectionError(message, selection, err as Error)
         } finally {
-            await this.ssmClient.terminateSession(testSession)
+            await this.ssm.terminateSession(testSession)
         }
     }
 
@@ -232,8 +232,8 @@ 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')
+    private async startSSMSession(instanceId: string): Promise<StartSessionResponse> {
+        const ssmSession = await this.ssm.startSession(instanceId, 'AWS-StartSSHSession')
         await this.addActiveSession(instanceId, ssmSession.SessionId!)
         return ssmSession
     }
@@ -308,7 +308,7 @@ export class Ec2Connecter implements vscode.Disposable {
     }
 
     private async sendCommandAndWait(instanceId: string, command: string) {
-        return await this.ssmClient.sendCommandAndWait(instanceId, 'AWS-RunShellScript', {
+        return await this.ssm.sendCommandAndWait(instanceId, 'AWS-RunShellScript', {
             commands: [command],
         })
     }
@@ -331,7 +331,7 @@ export class Ec2Connecter implements vscode.Disposable {
     }
 
     public async getRemoteUser(instanceId: string): Promise<RemoteUser> {
-        const os = await this.ssmClient.getTargetPlatformName(instanceId)
+        const os = await this.ssm.getTargetPlatformName(instanceId)
         if (os === 'Amazon Linux') {
             return { name: 'ec2-user', os }
         }

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

@@ -3,30 +3,30 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { EC2, SSM } from 'aws-sdk'
-import { SsmClient } from '../../shared/clients/ssmClient'
+import { EC2 } from 'aws-sdk'
+import { SSMWrapper } from '../../shared/clients/ssm'
 import { Disposable } from 'vscode'
 
-export class Ec2SessionTracker extends Map<EC2.InstanceId, SSM.SessionId> implements Disposable {
+export class Ec2SessionTracker extends Map<EC2.InstanceId, string> implements Disposable {
     public constructor(
         readonly regionCode: string,
-        protected ssmClient: SsmClient
+        protected ssm: SSMWrapper
     ) {
         super()
     }
 
-    public async addSession(instanceId: EC2.InstanceId, sessionId: SSM.SessionId): Promise<void> {
+    public async addSession(instanceId: EC2.InstanceId, sessionId: string): Promise<void> {
         if (this.isConnectedTo(instanceId)) {
             const existingSessionId = this.get(instanceId)!
-            await this.ssmClient.terminateSessionFromId(existingSessionId)
+            await this.ssm.terminateSessionFromId(existingSessionId)
             this.set(instanceId, sessionId)
         } else {
             this.set(instanceId, sessionId)
         }
     }
 
     private async disconnectEnv(instanceId: EC2.InstanceId): Promise<void> {
-        await this.ssmClient.terminateSessionFromId(this.get(instanceId)!)
+        await this.ssm.terminateSessionFromId(this.get(instanceId)!)
         this.delete(instanceId)
     }
 

packages/core/src/extension.ts

@@ -51,6 +51,7 @@ import { registerCommands } from './commands'
 import endpoints from '../resources/endpoints.json'
 import { getLogger, maybeShowMinVscodeWarning, setupUninstallHandler } from './shared'
 import { showViewLogsMessage } from './shared/utilities/messages'
+import { AWSClientBuilderV3 } from './shared/awsClientBuilderV3'
 
 disableAwsSdkWarning()
 
@@ -102,6 +103,7 @@ export async function activateCommon(
     globals.machineId = await getMachineId()
     globals.awsContext = new DefaultAwsContext()
     globals.sdkClientBuilder = new DefaultAWSClientBuilder(globals.awsContext)
+    globals.sdkClientBuilderV3 = new AWSClientBuilderV3(globals.awsContext)
     globals.loginManager = new LoginManager(globals.awsContext, new CredentialsStore())
 
     // order matters here

packages/core/src/shared/awsClientBuilderV3.ts

@@ -12,6 +12,7 @@ import {
     DeserializeHandler,
     DeserializeHandlerOptions,
     DeserializeMiddleware,
+    Handler,
     HandlerExecutionContext,
     MetadataBearer,
     MiddlewareStack,
@@ -33,10 +34,18 @@ export type AwsClientConstructor<C> = new (o: AwsClientOptions) => C
 
 // AWS-SDKv3 does not export generic types for clients so we need to build them as needed
 // https://github.com/aws/aws-sdk-js-v3/issues/5856#issuecomment-2096950979
-interface AwsClient {
+export interface AwsClient {
     middlewareStack: {
         add: MiddlewareStack<any, MetadataBearer>['add']
     }
+    send: (command: AwsCommand, options?: any) => Promise<any>
+    destroy: () => void
+}
+
+export interface AwsCommand {
+    input: object
+    middlewareStack: any
+    resolveMiddleware: (stack: any, configuration: any, options: any) => Handler<any, any>
 }
 
 interface AwsClientOptions {

packages/core/src/shared/clients/ssm.ts

@@ -0,0 +1,120 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+    SSMClient,
+    Session,
+    StartSessionCommand,
+    TerminateSessionCommand,
+    TerminateSessionResponse,
+    StartSessionCommandOutput,
+    DescribeInstanceInformationCommand,
+    DescribeInstanceInformationCommandInput,
+    InstanceInformation,
+    SendCommandCommand,
+    SendCommandCommandOutput,
+    waitUntilCommandExecuted,
+    SessionState,
+    DescribeSessionsCommand,
+} from '@aws-sdk/client-ssm'
+import { WaiterState } from '@smithy/util-waiter'
+import { ToolkitError } from '../errors'
+import { ClientWrapper } from './wrapper'
+
+export class SSMWrapper extends ClientWrapper<SSMClient> {
+    public constructor(public override readonly regionCode: string) {
+        super(regionCode, SSMClient)
+    }
+
+    public async terminateSession(session: Session): Promise<TerminateSessionResponse> {
+        const sessionId = session.SessionId!
+        return await this.terminateSessionFromId(sessionId)
+    }
+
+    public async terminateSessionFromId(sessionId: string): Promise<TerminateSessionResponse> {
+        return await this.makeRequest(TerminateSessionCommand, { SessionId: sessionId })
+    }
+
+    public async startSession(
+        target: string,
+        document?: string,
+        reason?: string,
+        parameters?: Record<string, string[]>
+    ): Promise<StartSessionCommandOutput> {
+        return await this.makeRequest(StartSessionCommand, {
+            Target: target,
+            DocumentName: document,
+            Reason: reason,
+            Parameters: parameters,
+        })
+    }
+
+    public async describeInstance(target: string): Promise<InstanceInformation> {
+        const response = this.makePaginatedRequest(
+            DescribeInstanceInformationCommand,
+            {
+                InstanceInformationFilterList: [
+                    {
+                        key: 'InstanceIds',
+                        valueSet: [target],
+                    },
+                ],
+            } as DescribeInstanceInformationCommandInput,
+            'InstanceIds'
+        )
+        const resolvedResponse = await response.flatten().flatten().promise()
+        return resolvedResponse[0]!
+    }
+
+    public async getTargetPlatformName(target: string): Promise<string> {
+        const instanceInformation = await this.describeInstance(target)
+        return instanceInformation.PlatformName!
+    }
+
+    public async sendCommand(
+        target: string,
+        documentName: string,
+        parameters: Record<string, string[]>
+    ): Promise<SendCommandCommandOutput> {
+        return await this.makeRequest(SendCommandCommand, {
+            InstanceIds: [target],
+            DocumentName: documentName,
+            Parameters: parameters,
+        })
+    }
+
+    private async waitForCommand(commandId: string, target: string) {
+        const result = await waitUntilCommandExecuted(
+            { client: await this.getClient(), maxWaitTime: 30 },
+            { CommandId: commandId, InstanceId: target }
+        )
+        if (result.state !== WaiterState.SUCCESS) {
+            throw new ToolkitError(`Command ${commandId} failed to execute on target ${target}`)
+        }
+    }
+
+    public async sendCommandAndWait(
+        target: string,
+        documentName: string,
+        parameters: Record<string, string[]>
+    ): Promise<SendCommandCommandOutput> {
+        const response = await this.sendCommand(target, documentName, parameters)
+        try {
+            await this.waitForCommand(response.Command!.CommandId!, target)
+            return response
+        } catch (err) {
+            throw new ToolkitError(`Failed in sending command to target ${target}`, { cause: err as Error })
+        }
+    }
+
+    public async getInstanceAgentPingStatus(target: string): Promise<string> {
+        const instanceInformation = await this.describeInstance(target)
+        return instanceInformation ? instanceInformation.PingStatus! : 'Inactive'
+    }
+
+    public async describeSessions(state: SessionState) {
+        return await this.makePaginatedRequest(DescribeSessionsCommand, { State: state }, 'Sessions').promise()
+    }
+}