feat(iam): session tagging

packages/@aws-cdk/aws-ecs-patterns/test/fargate/integ.executionrole.expected.json

@@ -367,10 +367,14 @@
               "Action": "sts:AssumeRole",
               "Effect": "Allow",
               "Principal": {
-                "Service": [
-                  "ecs.amazonaws.com",
-                  "ecs-tasks.amazonaws.com"
-                ]
+                "Service": "ecs.amazonaws.com"
+              }
+            },
+            {
+              "Action": "sts:AssumeRole",
+              "Effect": "Allow",
+              "Principal": {
+                "Service": "ecs-tasks.amazonaws.com"
               }
             }
           ],

packages/@aws-cdk/aws-iam/README.md

@@ -209,13 +209,31 @@ The `WebIdentityPrincipal` class can be used as a principal for web identities l
 Cognito, Amazon, Google or Facebook, for example:
 
 ```ts
-const principal = new iam.WebIdentityPrincipal('cognito-identity.amazonaws.com')
-  .withConditions({
-    "StringEquals": { "cognito-identity.amazonaws.com:aud": "us-east-2:12345678-abcd-abcd-abcd-123456" },
-    "ForAnyValue:StringLike": {"cognito-identity.amazonaws.com:amr": "unauthenticated" },
-  });
+const principal = new iam.WebIdentityPrincipal('cognito-identity.amazonaws.com', {
+  'StringEquals': { 'cognito-identity.amazonaws.com:aud': 'us-east-2:12345678-abcd-abcd-abcd-123456' },
+  'ForAnyValue:StringLike': {'cognito-identity.amazonaws.com:amr': 'unauthenticated' },
+});
 ```
 
+If your identity provider is configured to assume a Role with [session
+tags](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_session-tags.html), you
+need to call `.withSessionTags()` to add the required permissions to the Role's
+policy document:
+
+```ts
+new iam.Role(this, 'Role', {
+  assumedBy: new iam.WebIdentityPrincipal('cognito-identity.amazonaws.com', {
+    'StringEquals': {
+      'cognito-identity.amazonaws.com:aud': 'us-east-2:12345678-abcd-abcd-abcd-123456',
+     },
+    'ForAnyValue:StringLike': {
+      'cognito-identity.amazonaws.com:amr': 'unauthenticated',
+    },
+  }).withSessionTags(),
+});
+```
+
+
 ## Parsing JSON Policy Documents
 
 The `PolicyDocument.fromJson` and `PolicyStatement.fromJson` static methods can be used to parse JSON objects. For example:

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

@@ -1,7 +1,9 @@
 import * as cdk from '@aws-cdk/core';
 import { Default, FactName, RegionInfo } from '@aws-cdk/region-info';
 import { IOpenIdConnectProvider } from './oidc-provider';
+import { PolicyDocument } from './policy-document';
 import { Condition, Conditions, PolicyStatement } from './policy-statement';
+import { defaultAddPrincipalToAssumeRole } from './private/assume-role-policy';
 import { ISamlProvider } from './saml-provider';
 import { LITERAL_STRING_KEY, mergePrincipal } from './util';
 
@@ -68,6 +70,25 @@ export interface IPrincipal extends IGrantable {
   addToPrincipalPolicy(statement: PolicyStatement): AddToPrincipalPolicyResult;
 }
 
+/**
+ * A type of principal that has more control over its own representation in AssumeRolePolicyDocuments
+ *
+ * More complex types of identity providers need more control over Role's policy documents
+ * than simply `{ Effect: 'Allow', Action: 'AssumeRole', Principal: <Whatever> }`.
+ *
+ * If that control is necessary, they can implement `IAssumeRolePrincipal` to get full
+ * access to a Role's AssumeRolePolicyDocument.
+ */
+export interface IAssumeRolePrincipal extends IPrincipal {
+  /**
+   * Add the princpial to the AssumeRolePolicyDocument
+   *
+   * Add the statements to the AssumeRolePolicyDocument necessary to give this principal
+   * permissions to assume the given role.
+   */
+  addToAssumeRolePolicy(document: PolicyDocument): void;
+}
+
 /**
  * Result of calling `addToPrincipalPolicy`
  */
@@ -89,7 +110,7 @@ export interface AddToPrincipalPolicyResult {
 /**
  * Base class for policy principals
  */
-export abstract class PrincipalBase implements IPrincipal {
+export abstract class PrincipalBase implements IAssumeRolePrincipal {
   public readonly grantPrincipal: IPrincipal = this;
   public readonly principalAccount: string | undefined = undefined;
 
@@ -113,6 +134,14 @@ export abstract class PrincipalBase implements IPrincipal {
     return { statementAdded: false };
   }
 
+  public addToAssumeRolePolicy(document: PolicyDocument): void {
+    // Default implementation of this protocol, compatible with the legacy behavior
+    document.addStatements(new PolicyStatement({
+      actions: [this.assumeRoleAction],
+      principals: [this],
+    }));
+  }
+
   public toString() {
     // This is a first pass to make the object readable. Descendant principals
     // should return something nicer.
@@ -138,9 +167,39 @@ export abstract class PrincipalBase implements IPrincipal {
    *
    * @returns a new PrincipalWithConditions object.
    */
-  public withConditions(conditions: Conditions): IPrincipal {
+  public withConditions(conditions: Conditions): PrincipalBase {
     return new PrincipalWithConditions(this, conditions);
   }
+
+  /**
+   * Returns a new principal using this principal as the base, with session tags enabled.
+   *
+   * @returns a new SessionTagsPrincipal object.
+   */
+  public withSessionTags(): PrincipalBase {
+    return new SessionTagsPrincipal(this);
+  }
+}
+
+/**
+ * Base class for Principals that wrap other principals
+ */
+class PrincipalAdapter extends PrincipalBase {
+  public readonly assumeRoleAction = this.wrapped.assumeRoleAction;
+  public readonly principalAccount = this.wrapped.principalAccount;
+
+  constructor(protected readonly wrapped: IPrincipal) {
+    super();
+  }
+
+  public get policyFragment(): PrincipalPolicyFragment { return this.wrapped.policyFragment; }
+
+  addToPolicy(statement: PolicyStatement): boolean {
+    return this.wrapped.addToPolicy(statement);
+  }
+  addToPrincipalPolicy(statement: PolicyStatement): AddToPrincipalPolicyResult {
+    return this.wrapped.addToPrincipalPolicy(statement);
+  }
 }
 
 /**
@@ -149,15 +208,11 @@ export abstract class PrincipalBase implements IPrincipal {
  * For more information about conditions, see:
  * https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_condition.html
  */
-export class PrincipalWithConditions implements IPrincipal {
-  public readonly grantPrincipal: IPrincipal = this;
-  public readonly assumeRoleAction: string = this.principal.assumeRoleAction;
+export class PrincipalWithConditions extends PrincipalAdapter {
   private additionalConditions: Conditions;
 
-  constructor(
-    private readonly principal: IPrincipal,
-    conditions: Conditions,
-  ) {
+  constructor(principal: IPrincipal, conditions: Conditions) {
+    super(principal);
     this.additionalConditions = conditions;
   }
 
@@ -186,27 +241,15 @@ export class PrincipalWithConditions implements IPrincipal {
    * See [the IAM documentation](https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_condition.html).
    */
   public get conditions() {
-    return this.mergeConditions(this.principal.policyFragment.conditions, this.additionalConditions);
+    return this.mergeConditions(this.wrapped.policyFragment.conditions, this.additionalConditions);
   }
 
   public get policyFragment(): PrincipalPolicyFragment {
-    return new PrincipalPolicyFragment(this.principal.policyFragment.principalJson, this.conditions);
-  }
-
-  public get principalAccount(): string | undefined {
-    return this.principal.principalAccount;
-  }
-
-  public addToPolicy(statement: PolicyStatement): boolean {
-    return this.addToPrincipalPolicy(statement).statementAdded;
-  }
-
-  public addToPrincipalPolicy(statement: PolicyStatement): AddToPrincipalPolicyResult {
-    return this.principal.addToPrincipalPolicy(statement);
+    return new PrincipalPolicyFragment(this.wrapped.policyFragment.principalJson, this.conditions);
   }
 
   public toString() {
-    return this.principal.toString();
+    return this.wrapped.toString();
   }
 
   /**
@@ -247,6 +290,30 @@ export class PrincipalWithConditions implements IPrincipal {
   }
 }
 
+/**
+ * Enables session tags on role assumptions from a principal
+ *
+ * For more information on session tags, see:
+ * https://docs.aws.amazon.com/IAM/latest/UserGuide/id_session-tags.html
+ */
+export class SessionTagsPrincipal extends PrincipalAdapter {
+  constructor(principal: IPrincipal) {
+    super(principal);
+  }
+
+  public addToAssumeRolePolicy(doc: PolicyDocument) {
+    // Lazy import to avoid circular import dependencies during startup
+
+    // eslint-disable-next-line @typescript-eslint/no-require-imports
+    const adapter: typeof import('./private/policydoc-adapter') = require('./private/policydoc-adapter');
+
+    defaultAddPrincipalToAssumeRole(this.wrapped, new adapter.MutatingPolicyDocumentAdapter(doc, (statement) => {
+      statement.addActions('sts:TagSession');
+      return statement;
+    }));
+  }
+}
+
 /**
  * A collection of the fields in a PolicyStatement that can be used to identify a principal.
  *
@@ -441,6 +508,7 @@ export class FederatedPrincipal extends PrincipalBase {
    * @param federated federated identity provider (i.e. 'cognito-identity.amazonaws.com' for users authenticated through Cognito)
    * @param conditions The conditions under which the policy is in effect.
    *   See [the IAM documentation](https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_condition.html).
+   * @param sessionTags Whether to enable session tagging (see https://docs.aws.amazon.com/IAM/latest/UserGuide/id_session-tags.html)
    */
   constructor(
     public readonly federated: string,
@@ -471,6 +539,7 @@ export class WebIdentityPrincipal extends FederatedPrincipal {
    * @param identityProvider identity provider (i.e. 'cognito-identity.amazonaws.com' for users authenticated through Cognito)
    * @param conditions The conditions under which the policy is in effect.
    *   See [the IAM documentation](https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_condition.html).
+   * @param sessionTags Whether to enable session tagging (see https://docs.aws.amazon.com/IAM/latest/UserGuide/id_session-tags.html)
    */
   constructor(identityProvider: string, conditions: Conditions = {}) {
     super(identityProvider, conditions ?? {}, 'sts:AssumeRoleWithWebIdentity');
@@ -606,9 +675,9 @@ export class StarPrincipal extends PrincipalBase {
  */
 export class CompositePrincipal extends PrincipalBase {
   public readonly assumeRoleAction: string;
-  private readonly principals = new Array<PrincipalBase>();
+  private readonly principals = new Array<IPrincipal>();
 
-  constructor(...principals: PrincipalBase[]) {
+  constructor(...principals: IPrincipal[]) {
     super();
     if (principals.length === 0) {
       throw new Error('CompositePrincipals must be constructed with at least 1 Principal but none were passed.');
@@ -623,28 +692,29 @@ export class CompositePrincipal extends PrincipalBase {
    *
    * @param principals IAM principals that will be added to the composite principal
    */
-  public addPrincipals(...principals: PrincipalBase[]): this {
-    for (const p of principals) {
-      if (p.assumeRoleAction !== this.assumeRoleAction) {
-        throw new Error(
-          'Cannot add multiple principals with different "assumeRoleAction". ' +
-          `Expecting "${this.assumeRoleAction}", got "${p.assumeRoleAction}"`);
-      }
+  public addPrincipals(...principals: IPrincipal[]): this {
+    this.principals.push(...principals);
+    return this;
+  }
+
+  public addToAssumeRolePolicy(doc: PolicyDocument) {
+    for (const p of this.principals) {
+      defaultAddPrincipalToAssumeRole(p, doc);
+    }
+  }
 
+  public get policyFragment(): PrincipalPolicyFragment {
+    // We only have a problem with conditions if we are trying to render composite
+    // princpals into a single statement (which is when `policyFragment` would get called)
+    for (const p of this.principals) {
       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)}`);
       }
-
-      this.principals.push(p);
     }
 
-    return this;
-  }
-
-  public get policyFragment(): PrincipalPolicyFragment {
     const principalJson: { [key: string]: string[] } = {};
 
     for (const p of this.principals) {

packages/@aws-cdk/aws-iam/lib/private/assume-role-policy.ts

@@ -0,0 +1,25 @@
+import { PolicyDocument } from '../policy-document';
+import { PolicyStatement } from '../policy-statement';
+import { IPrincipal, IAssumeRolePrincipal } from '../principals';
+
+/**
+ * Add a principal to an AssumeRolePolicyDocument in the right way
+ *
+ * Delegate to the principal if it can do the job itself, do a default job if it can't.
+ */
+export function defaultAddPrincipalToAssumeRole(principal: IPrincipal, doc: PolicyDocument) {
+  if (isAssumeRolePrincipal(principal)) {
+    // Principal knows how to add itself
+    principal.addToAssumeRolePolicy(doc);
+  } else {
+    // Principal can't add itself, we do it for them
+    doc.addStatements(new PolicyStatement({
+      actions: [principal.assumeRoleAction],
+      principals: [principal],
+    }));
+  }
+}
+
+function isAssumeRolePrincipal(principal: IPrincipal): principal is IAssumeRolePrincipal {
+  return !!(principal as IAssumeRolePrincipal).addToAssumeRolePolicy;
+}

packages/@aws-cdk/aws-iam/lib/private/policydoc-adapter.ts

@@ -0,0 +1,17 @@
+import { PolicyDocument } from '../policy-document';
+import { PolicyStatement } from '../policy-statement';
+
+/**
+ * A PolicyDocument adapter that can modify statements flowing through it
+ */
+export class MutatingPolicyDocumentAdapter extends PolicyDocument {
+  constructor(private readonly wrapped: PolicyDocument, private readonly mutator: (s: PolicyStatement) => PolicyStatement) {
+    super();
+  }
+
+  public addStatements(...statements: PolicyStatement[]): void {
+    for (const st of statements) {
+      this.wrapped.addStatements(this.mutator(st));
+    }
+  }
+}