refactor(core): migrate sts + ecr + ec2 + lambda

chungjac · chungjac · commit d92f77c47bbd · 2025-10-13T20:20:22.000-07:00
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/core/package.json b/packages/core/package.json
@@ -593,6 +593,7 @@
         "@aws-sdk/client-docdb": "<3.731.0",
         "@aws-sdk/client-docdb-elastic": "<3.731.0",
         "@aws-sdk/client-ec2": "<3.731.0",
+        "@aws-sdk/client-ecr": "~3.693.0",
         "@aws-sdk/client-ecs": "~3.693.0",
         "@aws-sdk/client-glue": "^3.852.0",
         "@aws-sdk/client-iam": "<3.731.0",
diff --git a/packages/core/src/shared/clients/ec2MetadataClient.ts b/packages/core/src/shared/clients/ec2MetadataClient.ts
@@ -5,8 +5,8 @@
 
 import { getLogger } from '../logger/logger'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
-import { MetadataService } from 'aws-sdk'
-import { ServiceException } from '@smithy/smithy-client'
+import { httpRequest } from '@smithy/credential-provider-imds'
+import { RequestOptions } from 'http'
 
 export interface IamInfo {
     Code: string
@@ -22,8 +22,12 @@ export interface InstanceIdentity {
 export type Ec2MetadataClient = ClassToInterfaceType<DefaultEc2MetadataClient>
 export class DefaultEc2MetadataClient {
     private static readonly metadataServiceTimeout: number = 500
+    // AWS EC2 Instance Metadata Service (IMDS) constants
+    // https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-v2-how-it-works.html
+    private static readonly metadataServiceHost: string = '169.254.169.254'
+    private static readonly tokenPath: string = '/latest/api/token'
 
-    public constructor(private metadata: MetadataService = DefaultEc2MetadataClient.getMetadataService()) {}
+    public constructor() {}
 
     public getInstanceIdentity(): Promise<InstanceIdentity> {
         return this.invoke<InstanceIdentity>('/latest/dynamic/instance-identity/document')
@@ -33,52 +37,61 @@ export class DefaultEc2MetadataClient {
         return this.invoke<IamInfo>('/latest/meta-data/iam/info')
     }
 
-    public invoke<T>(path: string): Promise<T> {
-        return new Promise((resolve, reject) => {
-            // fetchMetadataToken is private for some reason, but has the exact token functionality
-            // that we want out of the metadata service.
-            // https://github.com/aws/aws-sdk-js/blob/3333f8b49283f5bbff823ab8a8469acedb7fe3d5/lib/metadata_service.js#L116-L136
-            ;(this.metadata as any).fetchMetadataToken((tokenErr: ServiceException, token: string) => {
-                let options
-                if (tokenErr) {
-                    getLogger().warn(
-                        'Ec2MetadataClient failed to fetch token. If this is an EC2 environment, then Toolkit will fall back to IMDSv1: %s',
-                        tokenErr
-                    )
+    public async invoke<T>(path: string): Promise<T> {
+        try {
+            // Try to get IMDSv2 token first
+            const token = await this.fetchMetadataToken()
+            const headers: Record<string, string> = {}
+            if (token) {
+                headers['x-aws-ec2-metadata-token'] = token
+            }
 
-                    // Fall back to IMDSv1 for legacy instances.
-                    options = {}
-                } else {
-                    options = {
-                        // By attaching the token we force the use of IMDSv2.
-                        // https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-v2-how-it-works.html
-                        headers: { 'x-aws-ec2-metadata-token': token },
-                    }
-                }
+            const response = await this.makeRequest(path, headers)
+            return JSON.parse(response.toString())
+        } catch (tokenErr) {
+            getLogger().warn(
+                'Ec2MetadataClient failed to fetch token. If this is an EC2 environment, then Toolkit will fall back to IMDSv1: %s',
+                tokenErr
+            )
 
-                this.metadata.request(path, options, (err, response) => {
-                    if (err) {
-                        reject(err)
-                        return
-                    }
-                    try {
-                        const jsonResponse: T = JSON.parse(response)
-                        resolve(jsonResponse)
-                    } catch (e) {
-                        reject(`Ec2MetadataClient: invalid response from "${path}": ${response}\nerror: ${e}`)
-                    }
-                })
-            })
-        })
+            // Fall back to IMDSv1 for legacy instances
+            try {
+                const response = await this.makeRequest(path, {})
+                return JSON.parse(response.toString())
+            } catch (err) {
+                throw new Error(`Ec2MetadataClient: failed to fetch "${path}": ${err}`)
+            }
+        }
     }
 
-    private static getMetadataService() {
-        return new MetadataService({
-            httpOptions: {
+    private async fetchMetadataToken(): Promise<string | undefined> {
+        try {
+            const options: RequestOptions = {
+                host: DefaultEc2MetadataClient.metadataServiceHost,
+                path: DefaultEc2MetadataClient.tokenPath,
+                method: 'PUT',
+                headers: {
+                    'x-aws-ec2-metadata-token-ttl-seconds': '21600',
+                },
                 timeout: DefaultEc2MetadataClient.metadataServiceTimeout,
-                connectTimeout: DefaultEc2MetadataClient.metadataServiceTimeout,
-            } as any,
-            // workaround for known bug: https://github.com/aws/aws-sdk-js/issues/3029
-        })
+            }
+
+            const response = await httpRequest(options)
+            return response.toString()
+        } catch (err) {
+            return undefined
+        }
+    }
+
+    private async makeRequest(path: string, headers: Record<string, string>): Promise<Buffer> {
+        const options: RequestOptions = {
+            host: DefaultEc2MetadataClient.metadataServiceHost,
+            path,
+            method: 'GET',
+            headers,
+            timeout: DefaultEc2MetadataClient.metadataServiceTimeout,
+        }
+
+        return httpRequest(options)
     }
 }
diff --git a/packages/core/src/shared/clients/ecrClient.ts b/packages/core/src/shared/clients/ecrClient.ts
@@ -3,23 +3,31 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ECR } from 'aws-sdk'
+import {
+    ECRClient,
+    DescribeImagesCommand,
+    DescribeRepositoriesCommand,
+    CreateRepositoryCommand,
+    DeleteRepositoryCommand,
+    BatchDeleteImageCommand,
+} from '@aws-sdk/client-ecr'
+import type { DescribeImagesRequest, DescribeRepositoriesRequest, Repository } from '@aws-sdk/client-ecr'
 import globals from '../extensionGlobals'
 import { AsyncCollection } from '../utilities/asyncCollection'
 import { pageableToCollection } from '../utilities/collectionUtils'
 import { assertHasProps, ClassToInterfaceType, isNonNullable, RequiredProps } from '../utilities/tsUtils'
 
-export type EcrRepository = RequiredProps<ECR.Repository, 'repositoryName' | 'repositoryUri' | 'repositoryArn'>
+export type EcrRepository = RequiredProps<Repository, 'repositoryName' | 'repositoryUri' | 'repositoryArn'>
 
 export type EcrClient = ClassToInterfaceType<DefaultEcrClient>
 export class DefaultEcrClient {
     public constructor(public readonly regionCode: string) {}
 
     public async *describeTags(repositoryName: string): AsyncIterableIterator<string> {
-        const sdkClient = await this.createSdkClient()
-        const request: ECR.DescribeImagesRequest = { repositoryName: repositoryName }
+        const sdkClient = this.createSdkClient()
+        const request: DescribeImagesRequest = { repositoryName: repositoryName }
         do {
-            const response = await sdkClient.describeImages(request).promise()
+            const response = await sdkClient.send(new DescribeImagesCommand(request))
             if (response.imageDetails) {
                 for (const item of response.imageDetails) {
                     if (item.imageTags !== undefined) {
@@ -34,13 +42,13 @@ export class DefaultEcrClient {
     }
 
     public async *describeRepositories(): AsyncIterableIterator<EcrRepository> {
-        const sdkClient = await this.createSdkClient()
-        const request: ECR.DescribeRepositoriesRequest = {}
+        const sdkClient = this.createSdkClient()
+        const request: DescribeRepositoriesRequest = {}
         do {
-            const response = await sdkClient.describeRepositories(request).promise()
+            const response = await sdkClient.send(new DescribeRepositoriesCommand(request))
             if (response.repositories) {
                 yield* response.repositories
-                    .map((repo) => {
+                    .map((repo: Repository) => {
                         // If any of these are not present, the repo returned is not valid. repositoryUri/Arn
                         // are both based on name, and it's not possible to not have a name
                         if (!repo.repositoryArn || !repo.repositoryName || !repo.repositoryUri) {
@@ -53,36 +61,43 @@ export class DefaultEcrClient {
                             }
                         }
                     })
-                    .filter((item) => item !== undefined) as EcrRepository[]
+                    .filter((item: EcrRepository | undefined) => item !== undefined) as EcrRepository[]
             }
             request.nextToken = response.nextToken
         } while (request.nextToken)
     }
 
     public listAllRepositories(): AsyncCollection<EcrRepository[]> {
-        const requester = async (req: ECR.DescribeRepositoriesRequest) =>
-            (await this.createSdkClient()).describeRepositories(req).promise()
+        const requester = async (req: DescribeRepositoriesRequest) =>
+            this.createSdkClient().send(new DescribeRepositoriesCommand(req))
         const collection = pageableToCollection(requester, {}, 'nextToken', 'repositories')
 
-        return collection.filter(isNonNullable).map((list) => list.map((repo) => (assertHasProps(repo), repo)))
+        return collection
+            .filter(isNonNullable)
+            .map((list: Repository[]) => list.map((repo: Repository) => (assertHasProps(repo), repo)))
     }
 
     public async createRepository(repositoryName: string) {
-        const sdkClient = await this.createSdkClient()
-        return sdkClient.createRepository({ repositoryName: repositoryName }).promise()
+        const sdkClient = this.createSdkClient()
+        return sdkClient.send(new CreateRepositoryCommand({ repositoryName: repositoryName }))
     }
 
     public async deleteRepository(repositoryName: string): Promise<void> {
-        const sdkClient = await this.createSdkClient()
-        await sdkClient.deleteRepository({ repositoryName: repositoryName }).promise()
+        const sdkClient = this.createSdkClient()
+        await sdkClient.send(new DeleteRepositoryCommand({ repositoryName: repositoryName }))
     }
 
     public async deleteTag(repositoryName: string, tag: string): Promise<void> {
-        const sdkClient = await this.createSdkClient()
-        await sdkClient.batchDeleteImage({ repositoryName: repositoryName, imageIds: [{ imageTag: tag }] }).promise()
+        const sdkClient = this.createSdkClient()
+        await sdkClient.send(
+            new BatchDeleteImageCommand({ repositoryName: repositoryName, imageIds: [{ imageTag: tag }] })
+        )
     }
 
-    protected async createSdkClient(): Promise<ECR> {
-        return await globals.sdkClientBuilder.createAwsService(ECR, undefined, this.regionCode)
+    protected createSdkClient(): ECRClient {
+        return globals.sdkClientBuilderV3.createAwsService({
+            serviceClient: ECRClient,
+            clientOptions: { region: this.regionCode },
+        })
     }
 }
diff --git a/packages/core/src/shared/clients/lambdaClient.ts b/packages/core/src/shared/clients/lambdaClient.ts
@@ -40,7 +40,6 @@ import {
 import { CancellationError } from '../utilities/timeoutUtils'
 import { fromSSO } from '@aws-sdk/credential-provider-sso'
 import { getIAMConnection } from '../../auth/utils'
-import { WaiterConfiguration } from 'aws-sdk/lib/service'
 import { NodeHttpHandler } from '@smithy/node-http-handler'
 
 export type LambdaClient = ClassToInterfaceType<DefaultLambdaClient>
@@ -301,11 +300,19 @@ export class DefaultLambdaClient {
         )
     }
 
-    public async waitForActive(functionName: string, waiter?: WaiterConfiguration): Promise<void> {
+    public async waitForActive(
+        functionName: string,
+        waiter?: { maxWaitTime?: number; minDelay?: number; maxDelay?: number }
+    ): Promise<void> {
         const sdkClient = await this.createSdkClient()
 
         await waitUntilFunctionActiveV2(
-            { client: sdkClient, maxWaitTime: (waiter?.maxAttempts ?? 600) * (waiter?.delay ?? 1) },
+            {
+                client: sdkClient,
+                maxWaitTime: waiter?.maxWaitTime ?? 600,
+                minDelay: waiter?.minDelay ?? 1,
+                maxDelay: waiter?.maxDelay ?? 120,
+            },
             { FunctionName: functionName }
         )
     }
diff --git a/packages/core/src/shared/clients/stsClient.ts b/packages/core/src/shared/clients/stsClient.ts
@@ -3,41 +3,60 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { STS } from 'aws-sdk'
+import { STSClient, AssumeRoleCommand, GetCallerIdentityCommand } from '@aws-sdk/client-sts'
+import type { AssumeRoleRequest, AssumeRoleResponse, GetCallerIdentityResponse } from '@aws-sdk/client-sts'
+import { AwsCredentialIdentityProvider } from '@smithy/types'
 import { Credentials } from '@aws-sdk/types'
 import globals from '../extensionGlobals'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
 
-export type GetCallerIdentityResponse = STS.GetCallerIdentityResponse
+export type { GetCallerIdentityResponse }
 export type StsClient = ClassToInterfaceType<DefaultStsClient>
+
+// Helper function to convert Credentials to AwsCredentialIdentityProvider
+function toCredentialProvider(credentials: Credentials | AwsCredentialIdentityProvider): AwsCredentialIdentityProvider {
+    if (typeof credentials === 'function') {
+        return credentials
+    }
+    // Convert static credentials to provider function
+    return async () => credentials
+}
+
 export class DefaultStsClient {
     public constructor(
         public readonly regionCode: string,
-        private readonly credentials?: Credentials,
+        private readonly credentials?: Credentials | AwsCredentialIdentityProvider,
         private readonly endpointUrl?: string
     ) {}
 
-    public async assumeRole(request: STS.AssumeRoleRequest): Promise<STS.AssumeRoleResponse> {
-        const sdkClient = await this.createSdkClient()
-        const response = await sdkClient.assumeRole(request).promise()
+    public async assumeRole(request: AssumeRoleRequest): Promise<AssumeRoleResponse> {
+        const sdkClient = this.createSdkClient()
+        const response = await sdkClient.send(new AssumeRoleCommand(request))
         return response
     }
 
-    public async getCallerIdentity(): Promise<STS.GetCallerIdentityResponse> {
-        const sdkClient = await this.createSdkClient()
-        const response = await sdkClient.getCallerIdentity().promise()
+    public async getCallerIdentity(): Promise<GetCallerIdentityResponse> {
+        const sdkClient = this.createSdkClient()
+        const response = await sdkClient.send(new GetCallerIdentityCommand({}))
         return response
     }
 
-    private async createSdkClient(): Promise<STS> {
-        return await globals.sdkClientBuilder.createAwsService(
-            STS,
-            {
-                credentials: this.credentials,
-                stsRegionalEndpoints: 'regional',
-                endpoint: this.endpointUrl,
-            },
-            this.regionCode
-        )
+    private createSdkClient(): STSClient {
+        const clientOptions: { region: string; endpoint?: string; credentials?: AwsCredentialIdentityProvider } = {
+            region: this.regionCode,
+        }
+
+        if (this.endpointUrl) {
+            clientOptions.endpoint = this.endpointUrl
+        }
+
+        if (this.credentials) {
+            clientOptions.credentials = toCredentialProvider(this.credentials)
+        }
+
+        return globals.sdkClientBuilderV3.createAwsService({
+            serviceClient: STSClient,
+            clientOptions,
+        })
     }
 }
diff --git a/packages/core/src/test/setupUtil.ts b/packages/core/src/test/setupUtil.ts
