feat: add additional options for assuming roles (#115)

This change was made to v2 but was not made to v3. This PR duplicates the changes in #40. Fixes #
cdklabs · Oct 3, 2024 · b2d1f04 · b2d1f04
1 parent a40effc
commit b2d1f04
Show file tree

Hide file tree

Showing 9 changed files with 115 additions and 20 deletions.
diff --git a/lib/aws.ts b/lib/aws.ts
@@ -31,7 +31,12 @@ import {
   GetSecretValueCommandOutput,
   SecretsManagerClient,
 } from '@aws-sdk/client-secrets-manager';
-import { GetCallerIdentityCommand, STSClient, STSClientConfig } from '@aws-sdk/client-sts';
+import {
+  AssumeRoleCommandInput,
+  GetCallerIdentityCommand,
+  STSClient,
+  STSClientConfig,
+} from '@aws-sdk/client-sts';
 import { fromNodeProviderChain, fromTemporaryCredentials } from '@aws-sdk/credential-providers';
 import { Upload } from '@aws-sdk/lib-storage';
 import {
@@ -41,6 +46,10 @@ import {
 import { loadConfig } from '@smithy/node-config-provider';
 import type { AwsCredentialIdentityProvider } from '@smithy/types';
 
+export type AssumeRoleAdditionalOptions = Partial<
+  Omit<AssumeRoleCommandInput, 'ExternalId' | 'RoleArn'>
+>;
+
 export interface IS3Client {
   getBucketEncryption(
     input: GetBucketEncryptionCommandInput
@@ -82,6 +91,7 @@ export interface ClientOptions {
   region?: string;
   assumeRoleArn?: string;
   assumeRoleExternalId?: string;
+  assumeRoleAdditionalOptions?: AssumeRoleAdditionalOptions;
   quiet?: boolean;
 }
 
@@ -228,6 +238,10 @@ export class DefaultAwsClient implements IAws {
             RoleArn: options.assumeRoleArn,
             ExternalId: options.assumeRoleExternalId,
             RoleSessionName: `${USER_AGENT}-${safeUsername()}`,
+            TransitiveTagKeys: options.assumeRoleAdditionalOptions?.Tags
+              ? options.assumeRoleAdditionalOptions.Tags.map((t) => t.Key!)
+              : undefined,
+            ...options.assumeRoleAdditionalOptions,
           },
           clientConfig: this.config.clientConfig,
         });

diff --git a/lib/private/handlers/container-images.ts b/lib/private/handlers/container-images.ts
@@ -1,5 +1,6 @@
 import * as path from 'path';
 import { DockerImageDestination } from '@aws-cdk/cloud-assembly-schema';
+import { destinationToClientOptions } from '.';
 import { DockerImageManifestEntry } from '../../asset-manifest';
 import type { IECRClient } from '../../aws';
 import { EventType } from '../../progress';
@@ -105,7 +106,7 @@ export class ContainerImageAssetHandler implements IAssetHandler {
 
     const destination = await replaceAwsPlaceholders(this.asset.destination, this.host.aws);
     const ecr = await this.host.aws.ecrClient({
-      ...destination,
+      ...destinationToClientOptions(destination),
       quiet: options.quiet,
     });
     const account = async () => (await this.host.aws.discoverCurrentAccount())?.accountId;

diff --git a/lib/private/handlers/files.ts b/lib/private/handlers/files.ts
@@ -2,6 +2,7 @@ import { createReadStream, promises as fs } from 'fs';
 import * as path from 'path';
 import { FileAssetPackaging, FileSource } from '@aws-cdk/cloud-assembly-schema';
 import * as mime from 'mime';
+import { destinationToClientOptions } from '.';
 import { FileManifestEntry } from '../../asset-manifest';
 import { IS3Client } from '../../aws';
 import { EventType } from '../../progress';
@@ -36,7 +37,7 @@ export class FileAssetHandler implements IAssetHandler {
     const s3Url = `s3://${destination.bucketName}/${destination.objectKey}`;
     try {
       const s3 = await this.host.aws.s3Client({
-        ...destination,
+        ...destinationToClientOptions(destination),
         quiet: true,
       });
       this.host.emitMessage(EventType.CHECK, `Check ${s3Url}`);
@@ -54,7 +55,7 @@ export class FileAssetHandler implements IAssetHandler {
   public async publish(): Promise<void> {
     const destination = await replaceAwsPlaceholders(this.asset.destination, this.host.aws);
     const s3Url = `s3://${destination.bucketName}/${destination.objectKey}`;
-    const s3 = await this.host.aws.s3Client(destination);
+    const s3 = await this.host.aws.s3Client(destinationToClientOptions(destination));
     this.host.emitMessage(EventType.CHECK, `Check ${s3Url}`);
 
     const bucketInfo = BucketInformation.for(this.host);

diff --git a/lib/private/handlers/index.ts b/lib/private/handlers/index.ts
@@ -1,12 +1,14 @@
+import type { AwsDestination } from '@aws-cdk/cloud-assembly-schema';
 import { ContainerImageAssetHandler } from './container-images';
 import { FileAssetHandler } from './files';
 import {
-  AssetManifest,
+  type AssetManifest,
   DockerImageManifestEntry,
   FileManifestEntry,
-  IManifestEntry,
+  type IManifestEntry,
 } from '../../asset-manifest';
-import { IAssetHandler, IHandlerHost, IHandlerOptions } from '../asset-handler';
+import type { ClientOptions } from '../../aws';
+import type { IAssetHandler, IHandlerHost, IHandlerOptions } from '../asset-handler';
 
 export function makeAssetHandler(
   manifest: AssetManifest,
@@ -23,3 +25,12 @@ export function makeAssetHandler(
 
   throw new Error(`Unrecognized asset type: '${asset}'`);
 }
+
+export function destinationToClientOptions(destination: AwsDestination): ClientOptions {
+  return {
+    assumeRoleArn: destination.assumeRoleArn,
+    assumeRoleExternalId: destination.assumeRoleExternalId,
+    assumeRoleAdditionalOptions: destination.assumeRoleAdditionalOptions,
+    region: destination.region,
+  };
+}
diff --git a/package.json b/package.json
diff --git a/test/aws.test.ts b/test/aws.test.ts
@@ -0,0 +1,76 @@
+import 'aws-sdk-client-mock-jest';
+
+import { GetCallerIdentityCommand } from '@aws-sdk/client-sts';
+import { fromTemporaryCredentials } from '@aws-sdk/credential-providers';
+import { mockSTS } from './mock-aws';
+import { DefaultAwsClient } from '../lib';
+
+jest.mock('@aws-sdk/credential-providers');
+
+const { fromNodeProviderChain } = jest.requireActual('@aws-sdk/credential-providers');
+
+const roleArn = 'arn:aws:iam:123456789012:role/the-role-of-a-lifetime';
+
+mockSTS.on(GetCallerIdentityCommand).resolves({
+  Account: '123456789012',
+  Arn: roleArn,
+});
+
+test('the correct credentials are passed to fromTemporaryCredentials in awsOptions', async () => {
+  const aws = new DefaultAwsClient();
+
+  await aws.discoverTargetAccount({
+    region: 'far-far-away',
+    assumeRoleArn: roleArn,
+    assumeRoleExternalId: 'external-id',
+    assumeRoleAdditionalOptions: {
+      DurationSeconds: 3600,
+      RoleSessionName: 'definitely-me',
+    },
+  });
+
+  expect(fromTemporaryCredentials).toHaveBeenCalledWith({
+    clientConfig: {
+      customUserAgent: 'cdk-assets',
+    },
+    params: {
+      ExternalId: 'external-id',
+      RoleArn: roleArn,
+      RoleSessionName: 'definitely-me',
+      DurationSeconds: 3600,
+    },
+  });
+});
+
+test('session tags are passed to fromTemporaryCredentials in awsOptions', async () => {
+  const aws = new DefaultAwsClient();
+
+  await aws.discoverTargetAccount({
+    region: 'far-far-away',
+    assumeRoleArn: roleArn,
+    assumeRoleExternalId: 'external-id',
+    assumeRoleAdditionalOptions: {
+      RoleSessionName: 'definitely-me',
+      Tags: [
+        { Key: 'this', Value: 'one' },
+        { Key: 'that', Value: 'one' },
+      ],
+    },
+  });
+
+  expect(fromTemporaryCredentials).toHaveBeenCalledWith({
+    clientConfig: {
+      customUserAgent: 'cdk-assets',
+    },
+    params: {
+      ExternalId: 'external-id',
+      RoleArn: roleArn,
+      RoleSessionName: 'definitely-me',
+      Tags: [
+        { Key: 'this', Value: 'one' },
+        { Key: 'that', Value: 'one' },
+      ],
+      TransitiveTagKeys: ['this', 'that'],
+    },
+  });
+});
diff --git a/test/docker-images.test.ts b/test/docker-images.test.ts
@@ -359,10 +359,8 @@ test('pass destination properties to AWS client', async () => {
   await pub.publish();
 
   expect(ecrClient).toHaveBeenCalledWith({
-    imageTag: 'abcdef',
     region: 'us-north-50',
     assumeRoleArn: 'arn:aws:role',
-    repositoryName: 'repo',
   });
 });
 
@@ -698,10 +696,8 @@ describe('external assets', () => {
     await pub.publish();
 
     expect(ecrClient).toHaveBeenCalledWith({
-      imageTag: 'ghijkl',
       region: 'us-north-50',
       assumeRoleArn: 'arn:aws:role',
-      repositoryName: 'repo',
     });
 
     expectAllSpawns();

diff --git a/test/placeholders.test.ts b/test/placeholders.test.ts
@@ -69,8 +69,6 @@ test('correct calls are made', async () => {
 
   expect(s3Client).toHaveBeenCalledWith({
     assumeRoleArn: 'arn:aws:role-current_account',
-    bucketName: 'some_bucket-current_account-current_region',
-    objectKey: 'some_key-current_account-current_region',
   });
 
   expect(mockS3).toHaveReceivedCommandWith(ListObjectsV2Command, {
@@ -81,10 +79,8 @@ test('correct calls are made', async () => {
 
   expect(ecrClient).toHaveBeenCalledWith({
     assumeRoleArn: 'arn:aws:role-current_account',
-    imageTag: 'abcdef',
     quiet: undefined,
     region: 'explicit_region',
-    repositoryName: 'repo-current_account-explicit_region',
   });
 
   expect(mockEcr).toHaveReceivedCommandWith(DescribeImagesCommand, {

diff --git a/yarn.lock b/yarn.lock