feat(iam): Assume role actions

packages/@aws-cdk/aws-iam/lib/principals.ts

@@ -38,6 +38,11 @@ export interface IPrincipal extends IGrantable {
    */
   readonly assumeRoleAction: string;
 
+  /**
+   * When this Principal is used in an AssumeRole policy, the actions to use.
+   */
+  readonly assumeRoleActions?: string[];
+
   /**
    * Return the policy fragment that identifies this principal in a Policy.
    */
@@ -103,6 +108,11 @@ export abstract class PrincipalBase implements IPrincipal {
    */
   public readonly assumeRoleAction: string = 'sts:AssumeRole';
 
+  /**
+   * When this Principal is used in an AssumeRole policy, the actions to use.
+   */
+  public readonly assumeRoleActions?: string[] = [];
+
   public addToPolicy(statement: PolicyStatement): boolean {
     return this.addToPrincipalPolicy(statement).statementAdded;
   }
@@ -141,6 +151,25 @@ export abstract class PrincipalBase implements IPrincipal {
   public withConditions(conditions: Conditions): IPrincipal {
     return new PrincipalWithConditions(this, conditions);
   }
+
+  /**
+   * Given we have a principal with trust relationship, we can extend the assume
+   * role actions with additional actions
+   *
+   * @param additionalAssumeRoleAction String assume role action
+   * @returns void
+   */
+  public addAssumeRoleAction(additionalAssumeRoleAction: string) {
+    // Ensure the default assumeRoleAction is included
+    if (!this.assumeRoleActions?.includes(this.assumeRoleAction)) {
+      this.assumeRoleActions?.push(this.assumeRoleAction);
+    }
+
+    // Insert the additional assume role action
+    if (!this.assumeRoleActions?.includes(additionalAssumeRoleAction)) {
+      this.assumeRoleActions?.push(additionalAssumeRoleAction);
+    }
+  }
 }
 
 /**
@@ -627,14 +656,26 @@ export class CompositePrincipal extends PrincipalBase {
       if (p.assumeRoleAction !== this.assumeRoleAction) {
         throw new Error(
           'Cannot add multiple principals with different "assumeRoleAction". ' +
-          `Expecting "${this.assumeRoleAction}", got "${p.assumeRoleAction}"`);
+            `Expecting "${this.assumeRoleAction}", got "${p.assumeRoleAction}"`);
+      }
+
+      if (!(
+        Array.isArray(this.assumeRoleActions) &&
+        Array.isArray(p.assumeRoleActions) &&
+        this.assumeRoleActions.length === p.assumeRoleActions.length &&
+        this.assumeRoleActions.every((action) => p.assumeRoleActions?.includes(action))
+      )) {
+        throw new Error(
+          'Cannot add multiple principals with different "assumeRoleActions". ' +
+          `Expecting "${JSON.stringify(this.assumeRoleActions)}", got "${JSON.stringify(p.assumeRoleActions)}"`,
+        );
       }
 
       const fragment = p.policyFragment;
       if (fragment.conditions && Object.keys(fragment.conditions).length > 0) {
         throw new Error(
           'Components of a CompositePrincipal must not have conditions. ' +
-          `Tried to add the following fragment: ${JSON.stringify(fragment)}`);
+            `Tried to add the following fragment: ${JSON.stringify(fragment)}`);
       }
 
       this.principals.push(p);

packages/@aws-cdk/aws-iam/lib/role.ts

@@ -486,7 +486,12 @@ export interface IRole extends IIdentity {
 function createAssumeRolePolicy(principal: IPrincipal, externalIds: string[]) {
   const statement = new AwsStarStatement();
   statement.addPrincipals(principal);
-  statement.addActions(principal.assumeRoleAction);
+
+  if (Array.isArray(principal.assumeRoleActions) && principal.assumeRoleActions.length > 1) {
+    statement.addActions(...principal.assumeRoleActions);
+  } else {
+    statement.addActions(principal.assumeRoleAction);
+  }
 
   if (externalIds.length) {
     statement.addCondition('StringEquals', { 'sts:ExternalId': externalIds.length === 1 ? externalIds[0] : externalIds });
@@ -540,4 +545,4 @@ export interface WithoutPolicyUpdatesOptions {
    * @default false
    */
   readonly addGrantsToResources?: boolean;
-}
+}

packages/@aws-cdk/aws-iam/test/principals.test.ts

@@ -243,4 +243,42 @@ test('PrincipalWithConditions inherits principalAccount from AccountPrincipal ',
   // THEN
   expect(accountPrincipal.principalAccount).toStrictEqual('123456789012');
   expect(principalWithConditions.principalAccount).toStrictEqual('123456789012');
-});
+});
+
+test('Principal assume Role actions can be extended with additional actions', () => {
+  const stack = new Stack();
+  const federatedPrincipal = new iam.FederatedPrincipal(
+    'cognito-identity.amazonaws.com',
+    { StringEquals: { hairColor: 'blond' } },
+  );
+
+  federatedPrincipal.addAssumeRoleAction('sts:TagSession');
+
+  new iam.Role(stack, 'Role', {
+    assumedBy: federatedPrincipal,
+  });
+
+  // THEN
+  expect(stack).toHaveResource('AWS::IAM::Role', {
+    AssumeRolePolicyDocument: {
+      Statement: [
+        {
+          Action: [
+            'sts:AssumeRole',
+            'sts:TagSession',
+          ],
+          Condition: {
+            StringEquals: {
+              hairColor: 'blond',
+            },
+          },
+          Effect: 'Allow',
+          Principal: {
+            Federated: 'cognito-identity.amazonaws.com',
+          },
+        },
+      ],
+      Version: '2012-10-17',
+    },
+  });
+});