refactor(sdkv3): migrate ssm client

packages/core/package.json

@@ -511,6 +511,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.22.1",
         "@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 { SsmClient } 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: SsmClient
     protected ec2Client: Ec2Client
     protected iamClient: DefaultIamClient
     protected sessionManager: Ec2SessionTracker
@@ -65,10 +65,10 @@ 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 {
@@ -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 { SsmClient } 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: SsmClient
     ) {
         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

@@ -50,6 +50,7 @@ import { registerCommands } from './commands'
 // The following imports the endpoints file, which causes webpack to bundle it in the final output file
 import endpoints from '../resources/endpoints.json'
 import { showViewLogsMessage } from './shared/utilities/messages'
+import { AWSClientBuilderV3 } from './shared/awsClientBuilderV3'
 import { setupUninstallHandler } from './shared/handleUninstall'
 import { maybeShowMinVscodeWarning } from './shared/extensionStartup'
 import { getLogger } from './shared/logger/logger'
@@ -104,6 +105,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

@@ -13,6 +13,7 @@ import {
     BuildMiddleware,
     DeserializeHandler,
     DeserializeMiddleware,
+    Handler,
     FinalizeHandler,
     FinalizeRequestMiddleware,
     HandlerExecutionContext,
@@ -37,10 +38,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/clientWrapper.ts

@@ -0,0 +1,57 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import * as vscode from 'vscode'
+import globals from '../extensionGlobals'
+import { AwsClient, AwsClientConstructor, AwsCommand } from '../awsClientBuilderV3'
+import { pageableToCollection } from '../utilities/collectionUtils'
+
+export abstract class ClientWrapper<C extends AwsClient> implements vscode.Disposable {
+    protected client?: C
+
+    public constructor(
+        public readonly regionCode: string,
+        private readonly clientType: AwsClientConstructor<C>
+    ) {}
+
+    protected async getClient() {
+        if (this.client) {
+            return this.client
+        }
+        this.client = await globals.sdkClientBuilderV3.createAwsService(this.clientType, undefined, this.regionCode)
+        return this.client!
+    }
+
+    protected async makeRequest<CommandInput extends object, Command extends AwsCommand>(
+        command: new (o: CommandInput) => Command,
+        commandOptions: CommandInput
+    ) {
+        const client = await this.getClient()
+        return await client.send(new command(commandOptions))
+    }
+
+    protected makePaginatedRequest<
+        CommandInput extends object,
+        CommandOutput extends object,
+        Command extends AwsCommand,
+    >(
+        command: new (o: CommandInput) => Command,
+        commandOptions: CommandInput,
+        collectKey: keyof CommandOutput & string,
+        nextTokenKey?: keyof CommandOutput & keyof CommandInput & string
+    ) {
+        const requester = async (req: CommandInput) => await this.makeRequest(command, req)
+        const response = pageableToCollection(
+            requester,
+            commandOptions,
+            nextTokenKey ?? ('NextToken' as never),
+            collectKey
+        )
+        return response
+    }
+
+    public dispose() {
+        this.client?.destroy()
+    }
+}