refactor(sdkv3): migrate ssm client (#6137)

## Problem
There is a lot of duplicate code across the sdk clients. Additionally,
we want to only construct the clients once (per wrapper), and have these
clients call destroy.

## Solution
Note: the line changes are mostly the added dependency of an `@aws-sdk`
client and its hundreds of dependencies.

- create abstract class `ClientWrapper` to minimize code dupe. 
- cache sdk clients on creation within the wrappers. 
- migrate ssm client to sdkv3, using new common wrapper class. 

## Alternative Solution 
- Make all wrappers singletons so clients are created once for the
toolkit lifetime.
- Pro: Makes it impossible for a client to be re-created. 
- Con: Zero flexibility when it comes to client lifetime. All sdk
clients would be created whenever their wrapper is referenced. Without
being very careful, we could create many SDK clients up front that we
never use and can't delete.

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: Hweinstock

package-lock.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/package.json

@@ -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/model.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/awsService/ec2/remoteSessionManager.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/extension.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/awsClientBuilderV3.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()
+    }
+}