feat: additional options when assuming publishing roles (#40)

Required to implement [session tags](aws/aws-cdk#26157) and a prerequisite to aws/aws-cdk#31089.

### Notes

- Requires cdklabs/cloud-assembly-schema#33 to be merged first.

.projenrc.ts

@@ -32,6 +32,7 @@ const project = new typescript.TypeScriptProject({
     '@types/glob',
     '@types/mime',
     '@types/yargs',
+    '@types/mock-fs',
     'fs-extra',
     'graceful-fs',
     'jszip',

lib/aws.ts

@@ -14,10 +14,17 @@ export interface IAws {
   secretsManagerClient(options: ClientOptions): Promise<AWS.SecretsManager>;
 }
 
+// Partial because `RoleSessionName` is required in STS, but we have a default value for it.
+export type AssumeRoleAdditionalOptions = Partial<
+  // cloud-assembly-schema validates that `ExternalId` and `RoleArn` are not configured
+  Omit<AWS.STS.Types.AssumeRoleRequest, 'ExternalId' | 'RoleArn'>
+>;
+
 export interface ClientOptions {
   region?: string;
   assumeRoleArn?: string;
   assumeRoleExternalId?: string;
+  assumeRoleAdditionalOptions?: AssumeRoleAdditionalOptions;
   quiet?: boolean;
 }
 
@@ -118,7 +125,8 @@ export class DefaultAwsClient implements IAws {
       credentials = await this.assumeRole(
         options.region,
         options.assumeRoleArn,
-        options.assumeRoleExternalId
+        options.assumeRoleExternalId,
+        options.assumeRoleAdditionalOptions
       );
     }
 
@@ -140,13 +148,18 @@ export class DefaultAwsClient implements IAws {
   private async assumeRole(
     region: string | undefined,
     roleArn: string,
-    externalId?: string
+    externalId?: string,
+    additionalOptions?: AssumeRoleAdditionalOptions
   ): Promise<AWS.Credentials> {
     return new this.AWS.ChainableTemporaryCredentials({
       params: {
         RoleArn: roleArn,
         ExternalId: externalId,
         RoleSessionName: `cdk-assets-${safeUsername()}`,
+        TransitiveTagKeys: additionalOptions?.Tags
+          ? additionalOptions.Tags.map((t) => t.Key)
+          : undefined,
+        ...(additionalOptions ?? {}),
       },
       stsConfig: {
         region,

lib/private/handlers/container-images.ts

@@ -1,6 +1,7 @@
 import * as path from 'path';
 import { DockerImageDestination } from '@aws-cdk/cloud-assembly-schema';
 import type * as AWS from 'aws-sdk';
+import { destinationToClientOptions } from '.';
 import { DockerImageManifestEntry } from '../../asset-manifest';
 import { EventType } from '../../progress';
 import { IAssetHandler, IHandlerHost, IHandlerOptions } from '../asset-handler';
@@ -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;

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 { EventType } from '../../progress';
 import { zipDirectory } from '../archive';
@@ -35,7 +36,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}`);
@@ -53,14 +54,17 @@ 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 clientOptions = destinationToClientOptions(destination);
+    const s3 = await this.host.aws.s3Client(clientOptions);
     this.host.emitMessage(EventType.CHECK, `Check ${s3Url}`);
 
     const bucketInfo = BucketInformation.for(this.host);
 
     // A thunk for describing the current account. Used when we need to format an error
     // message, not in the success case.
-    const account = async () => (await this.host.aws.discoverTargetAccount(destination))?.accountId;
+    const account = async () =>
+      (await this.host.aws.discoverTargetAccount(clientOptions))?.accountId;
     switch (await bucketInfo.bucketOwnership(s3, destination.bucketName)) {
       case BucketOwnership.MINE:
         break;

lib/private/handlers/index.ts

@@ -1,3 +1,4 @@
+import { AwsDestination } from '@aws-cdk/cloud-assembly-schema';
 import { ContainerImageAssetHandler } from './container-images';
 import { FileAssetHandler } from './files';
 import {
@@ -6,6 +7,7 @@ import {
   FileManifestEntry,
   IManifestEntry,
 } from '../../asset-manifest';
+import type { ClientOptions } from '../../aws';
 import { IAssetHandler, IHandlerHost, IHandlerOptions } from '../asset-handler';
 
 export function makeAssetHandler(
@@ -23,3 +25,15 @@ export function makeAssetHandler(
 
   throw new Error(`Unrecognized asset type: '${asset}'`);
 }
+
+export function destinationToClientOptions(destination: AwsDestination): ClientOptions {
+  // Explicitly build ClientOptions from AwsDestination. The fact they are structurally compatible is coincidental.
+  // This also enforces better type checking that cdk-assets depends on the appropriate version of
+  // @aws-cdk/cloud-assembly-schema.
+  return {
+    assumeRoleArn: destination.assumeRoleArn,
+    assumeRoleExternalId: destination.assumeRoleExternalId,
+    assumeRoleAdditionalOptions: destination.assumeRoleAdditionalOptions,
+    region: destination.region,
+  };
+}

test/aws.test.ts

@@ -0,0 +1,136 @@
+import * as os from 'os';
+import { DefaultAwsClient } from '../lib';
+
+beforeEach(() => {
+  jest.requireActual('aws-sdk');
+});
+
+test('assumeRole passes the right parameters to STS', async () => {
+  // eslint-disable-next-line @typescript-eslint/no-require-imports
+  const AWS = require('aws-sdk');
+
+  jest.mock('aws-sdk', () => {
+    return {
+      STS: jest.fn().mockReturnValue({
+        getCallerIdentity: jest.fn().mockReturnValue({
+          promise: jest.fn().mockResolvedValue({
+            Account: '123456789012',
+            Arn: 'arn:aws:iam::123456789012:role/my-role',
+          }),
+        }),
+      }),
+      ChainableTemporaryCredentials: jest.fn(),
+    };
+  });
+  const aws = new DefaultAwsClient();
+  await withMocked(os, 'userInfo', async (userInfo) => {
+    userInfo.mockReturnValue({
+      username: 'foo',
+      uid: 1,
+      gid: 1,
+      homedir: '/here',
+      shell: '/bin/sh',
+    });
+    await aws.discoverTargetAccount({
+      region: 'us-east-1',
+      assumeRoleArn: 'arn:aws:iam::123456789012:role/my-role',
+      assumeRoleExternalId: 'external-id',
+      assumeRoleAdditionalOptions: {
+        DurationSeconds: 3600,
+      },
+    });
+    expect(AWS.ChainableTemporaryCredentials).toHaveBeenCalledWith({
+      params: {
+        ExternalId: 'external-id',
+        RoleArn: 'arn:aws:iam::123456789012:role/my-role',
+        DurationSeconds: 3600,
+        RoleSessionName: `cdk-assets-foo`,
+      },
+      stsConfig: {
+        customUserAgent: 'cdk-assets',
+        region: 'us-east-1',
+      },
+    });
+  });
+});
+
+test('assumeRole defaults session tags to all', async () => {
+  // eslint-disable-next-line @typescript-eslint/no-require-imports
+  const AWS = require('aws-sdk');
+
+  jest.mock('aws-sdk', () => {
+    return {
+      STS: jest.fn().mockReturnValue({
+        getCallerIdentity: jest.fn().mockReturnValue({
+          promise: jest.fn().mockResolvedValue({
+            Account: '123456789012',
+            Arn: 'arn:aws:iam::123456789012:role/my-role',
+          }),
+        }),
+      }),
+      ChainableTemporaryCredentials: jest.fn(),
+    };
+  });
+  const aws = new DefaultAwsClient();
+  await withMocked(os, 'userInfo', async (userInfo) => {
+    userInfo.mockReturnValue({
+      username: 'foo',
+      uid: 1,
+      gid: 1,
+      homedir: '/here',
+      shell: '/bin/sh',
+    });
+    await aws.discoverTargetAccount({
+      region: 'us-east-1',
+      assumeRoleArn: 'arn:aws:iam::123456789012:role/my-role',
+      assumeRoleExternalId: 'external-id',
+      assumeRoleAdditionalOptions: {
+        Tags: [{ Key: 'Departement', Value: 'Engineering' }],
+      },
+    });
+    expect(AWS.ChainableTemporaryCredentials).toHaveBeenCalledWith({
+      params: {
+        ExternalId: 'external-id',
+        RoleArn: 'arn:aws:iam::123456789012:role/my-role',
+        Tags: [{ Key: 'Departement', Value: 'Engineering' }],
+        TransitiveTagKeys: ['Departement'],
+        RoleSessionName: `cdk-assets-foo`,
+      },
+      stsConfig: {
+        customUserAgent: 'cdk-assets',
+        region: 'us-east-1',
+      },
+    });
+  });
+});
+
+export function withMocked<A extends object, K extends keyof A, B>(
+  obj: A,
+  key: K,
+  block: (fn: jest.Mocked<A>[K]) => B
+): B {
+  const original = obj[key];
+  const mockFn = jest.fn();
+  (obj as any)[key] = mockFn;
+
+  let asyncFinally: boolean = false;
+  try {
+    const ret = block(mockFn as any);
+    if (!isPromise(ret)) {
+      return ret;
+    }
+
+    asyncFinally = true;
+    return ret.finally(() => {
+      obj[key] = original;
+    }) as any;
+  } finally {
+    if (!asyncFinally) {
+      obj[key] = original;
+    }
+  }
+}
+
+function isPromise<A>(object: any): object is Promise<A> {
+  return Promise.resolve(object) === object;
+}

test/docker-images.test.ts

@@ -29,6 +29,10 @@ beforeEach(() => {
             theDestination: {
               region: 'us-north-50',
               assumeRoleArn: 'arn:aws:role',
+              assumeRoleExternalId: 'external-id',
+              assumeRoleAdditionalOptions: {
+                Tags: [{ Key: 'Departement', Value: 'Engineering' }],
+              },
               repositoryName: 'repo',
               imageTag: 'abcdef',
             },
@@ -249,6 +253,10 @@ test('pass destination properties to AWS client', async () => {
     expect.objectContaining({
       region: 'us-north-50',
       assumeRoleArn: 'arn:aws:role',
+      assumeRoleExternalId: 'external-id',
+      assumeRoleAdditionalOptions: {
+        Tags: [{ Key: 'Departement', Value: 'Engineering' }],
+      },
     })
   );
 });

test/files.test.ts

@@ -1,6 +1,6 @@
 jest.mock('child_process');
 
-import { Manifest } from '@aws-cdk/cloud-assembly-schema';
+import { FileDestination, Manifest } from '@aws-cdk/cloud-assembly-schema';
 import * as mockfs from 'mock-fs';
 import { FakeListener } from './fake-listener';
 import { mockAws, mockedApiFailure, mockedApiResult, mockUpload } from './mock-aws';
@@ -9,9 +9,13 @@ import { AssetPublishing, AssetManifest } from '../lib';
 
 const ABS_PATH = '/simple/cdk.out/some_external_file';
 
-const DEFAULT_DESTINATION = {
+const DEFAULT_DESTINATION: FileDestination = {
   region: 'us-north-50',
   assumeRoleArn: 'arn:aws:role',
+  assumeRoleExternalId: 'external-id',
+  assumeRoleAdditionalOptions: {
+    Tags: [{ Key: 'Departement', Value: 'Engineering' }],
+  },
   bucketName: 'some_bucket',
   objectKey: 'some_key',
 };
@@ -114,6 +118,10 @@ test('pass destination properties to AWS client', async () => {
     expect.objectContaining({
       region: 'us-north-50',
       assumeRoleArn: 'arn:aws:role',
+      assumeRoleExternalId: 'external-id',
+      assumeRoleAdditionalOptions: {
+        Tags: [{ Key: 'Departement', Value: 'Engineering' }],
+      },
     })
   );
 });