feat(iam): session tagging (#17689)

To allow session tagging, the `sts:TagSession` permission needs to
be added to the role's AssumeRolePolicyDocument.

Introduce a new principal which enables this, and add a convenience
method `.withSessionTags()` to the `PrincipalBase` class so all
built-in principals will have this convenience method by default.

To build this, we had to get rid of some cruft and assumptions around
policy documents and statements, and defer more power to the
`IPrincipal` objects themselves. In order not to break existing
implementors, introduce a new interface `IAssumeRolePrincipal` which
knows how to add itself to an AssumeRolePolicyDocument and gets complete
freedom doing so.

That same new interface could be used to lift some old limitations on
`CompositePrincipal` so did that as well.

Fixes #15908, closes #16725, fixes #2041, fixes #1578.


----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: rix0rrr

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));
+    }
+  }
+}