feat(iam): Synthesize a compact composite principal in assume role po…

…licy
aws · May 8, 2023 · 1c07e87 · 1c07e87
1 parent 8ef0ba2
commit 1c07e87
Show file tree

Hide file tree

Showing 2 changed files with 115 additions and 6 deletions.
diff --git a/packages/aws-cdk-lib/aws-iam/lib/principals.ts b/packages/aws-cdk-lib/aws-iam/lib/principals.ts
@@ -830,11 +830,30 @@ export class CompositePrincipal extends PrincipalBase {
   /**
    * Adds IAM principals to the composite principal. Composite principals cannot have
    * conditions.
+   * Principals with same role action will be merged into a single principal to generate more compact Policy.
    *
    * @param principals IAM principals that will be added to the composite principal
    */
   public addPrincipals(...principals: IPrincipal[]): this {
-    this.principals.push(...principals);
+    const buffer = [...this.principals, ...principals];
+    this.principals.length = 0;
+
+    const groupByAction: { [key: string]: { [key: string]: string[] } } = {};
+    for (const p of buffer) {
+      let noConditions = Object.keys(p.policyFragment.conditions).length === 0;
+      if (noConditions) {
+        groupByAction[p.assumeRoleAction] = groupByAction[p.assumeRoleAction] ?? {};
+        mergePrincipal(groupByAction[p.assumeRoleAction], p.policyFragment.principalJson);
+      } else {
+        // Keep a principal with condition as-is. See `policyFragment` for detail.
+        this.principals.push(p);
+      }
+    }
+    // Creates a policy-merged principal and push it back
+    Object.entries(groupByAction).forEach(([action, fragmentJson]) => {
+      const policyFragment = new PrincipalPolicyFragment(fragmentJson);
+      this.principals.push(new PrincipalFromFragment(action, policyFragment));
+    });
     return this;
   }
 
@@ -876,6 +895,21 @@ export class CompositePrincipal extends PrincipalBase {
   }
 }
 
+class PrincipalFromFragment extends PrincipalBase {
+  public readonly assumeRoleAction: string;
+  readonly policyFragment: PrincipalPolicyFragment;
+
+  constructor(assumeRoleAction: string, policyFragment: PrincipalPolicyFragment) {
+    super();
+    this.assumeRoleAction = assumeRoleAction;
+    this.policyFragment = policyFragment;
+  }
+
+  dedupeString(): string | undefined {
+    return undefined;
+  }
+}
+
 /**
  * A lazy token that requires an instance of Stack to evaluate
  */

diff --git a/packages/aws-cdk-lib/aws-iam/test/policy-document.test.ts b/packages/aws-cdk-lib/aws-iam/test/policy-document.test.ts
@@ -495,11 +495,6 @@ describe('IAM policy document', () => {
       Template.fromStack(stack).hasResourceProperties('AWS::IAM::Role', {
         AssumeRolePolicyDocument: {
           Statement: [
-            {
-              Action: 'sts:AssumeRole',
-              Effect: 'Allow',
-              Principal: { AWS: 'i:am' },
-            },
             {
               Action: 'sts:AssumeRole',
               Condition: {
@@ -508,6 +503,11 @@ describe('IAM policy document', () => {
               Effect: 'Allow',
               Principal: { Federated: 'federated' },
             },
+            {
+              Action: 'sts:AssumeRole',
+              Effect: 'Allow',
+              Principal: { AWS: 'i:am' },
+            },
           ],
         },
       });
@@ -555,6 +555,25 @@ describe('IAM policy document', () => {
           Service: ['amazon.com', 'another.service'],
         },
       });
+
+      new Role(stack, 'Role', {
+        assumedBy: p,
+      });
+
+      Template.fromStack(stack).hasResourceProperties('AWS::IAM::Role', {
+        AssumeRolePolicyDocument: {
+          Statement: [
+            {
+              Action: 'sts:AssumeRole',
+              Effect: 'Allow',
+              Principal: {
+                AWS: ['i:am:an:arn', '*'],
+                Service: ['amazon.com', 'another.service'],
+              },
+            },
+          ],
+        },
+      });
     });
 
     test('can mix types of assumeRoleAction in a single composite', () => {
@@ -585,6 +604,62 @@ describe('IAM policy document', () => {
         },
       });
     });
+
+    test('can merge principals with same action into a single principal', () => {
+      const stack = new Stack();
+
+      // WHEN
+      const compositePrincipal = new CompositePrincipal(
+        new ArnPrincipal('arn:aws:iam::123456789012:root'),
+        new ArnPrincipal('arn:aws:iam::234567890123:root'),
+        new ArnPrincipal('arn:aws:iam::345678901234:root'),
+        new FederatedPrincipal('fed', {}, 'sts:Boom'));
+      new Role(stack, 'Role', {
+        assumedBy: compositePrincipal,
+        inlinePolicies: {
+          readOnlyAccess: new PolicyDocument({
+            assignSids: true,
+            statements: [
+              new PolicyStatement({
+                effect: Effect.ALLOW,
+                actions: [
+                  's3:Get*',
+                  's3:List*',
+                ],
+                resources: [
+                  'arn:aws:s3:::test',
+                  'arn:aws:s3:::test/*',
+                ],
+              }),
+            ],
+          }),
+        },
+      });
+
+      // THEN
+      Template.fromStack(stack).hasResourceProperties('AWS::IAM::Role', {
+        AssumeRolePolicyDocument: {
+          Statement: [
+            {
+              Action: 'sts:AssumeRole',
+              Effect: 'Allow',
+              Principal: {
+                AWS: [
+                  'arn:aws:iam::123456789012:root',
+                  'arn:aws:iam::234567890123:root',
+                  'arn:aws:iam::345678901234:root',
+                ],
+              },
+            },
+            {
+              Action: 'sts:Boom',
+              Effect: 'Allow',
+              Principal: { Federated: 'fed' },
+            },
+          ],
+        },
+      });
+    });
   });
 
   describe('PrincipalWithConditions can be used to add a principal with conditions', () => {