fix(iam): Merge multiple principals correctly

packages/@aws-cdk/aws-iam/lib/policy-document.ts

@@ -65,7 +65,7 @@ export abstract class PolicyPrincipal {
  */
 export class PrincipalPolicyFragment {
   constructor(
-    public readonly principalJson: any,
+    public readonly principalJson: { [key: string]: any },
     public readonly conditions: {[key: string]: any} = {}) {
   }
 }
@@ -144,18 +144,9 @@ export class AccountRootPrincipal extends AccountPrincipal {
 /**
  * A principal representing all identities in all accounts
  */
-export class Anyone extends PolicyPrincipal {
-  /**
-   * Interface compatibility with AccountPrincipal for the purposes of the Lambda library
-   *
-   * The Lambda's addPermission() call works differently from regular
-   * statements, and will use the value of this property directly if present
-   * (which leads to the correct statement ultimately).
-   */
-  public readonly accountId = '*';
-
-  public policyFragment(): PrincipalPolicyFragment {
-    return new PrincipalPolicyFragment('*');
+export class Anyone extends ArnPrincipal {
+  constructor() {
+    super('*');
   }
 }
 
@@ -164,7 +155,7 @@ export class Anyone extends PolicyPrincipal {
  */
 export class PolicyStatement extends Token {
   private action = new Array<any>();
-  private principal = new Array<any>();
+  private principal: { [key: string]: any[] } = {};
   private resource = new Array<any>();
   private condition: { [key: string]: any } = { };
   private effect?: PolicyStatementEffect;
@@ -197,12 +188,23 @@ export class PolicyStatement extends Token {
    * Indicates if this permission has a "Principal" section.
    */
   public get hasPrincipal() {
-    return this.principal && this.principal.length > 0;
+    return Object.keys(this.principal).length > 0;
   }
 
   public addPrincipal(principal: PolicyPrincipal): PolicyStatement {
     const fragment = principal.policyFragment();
-    this.principal.push(fragment.principalJson);
+    for (const key of Object.keys(fragment.principalJson)) {
+      if (Object.keys(this.principal).length > 0 && !(key in this.principal)) {
+        throw new Error(`Attempted to add principal key ${key} in principal of type ${Object.keys(this.principal)[0]}`);
+      }
+      this.principal[key] = this.principal[key] || [];
+      const value = fragment.principalJson[key];
+      if (Array.isArray(value)) {
+        this.principal[key].push(...value);
+      } else {
+        this.principal[key].push(value);
+      }
+    }
     this.addConditions(fragment.conditions);
     return this;
   }
@@ -330,7 +332,7 @@ export class PolicyStatement extends Token {
       Action: _norm(this.action),
       Condition: _norm(this.condition),
       Effect: _norm(this.effect),
-      Principal: _norm(this.principal),
+      Principal: _normPrincipal(this.principal),
       Resource: _norm(this.resource),
       Sid: _norm(this.sid),
     };
@@ -361,6 +363,22 @@ export class PolicyStatement extends Token {
 
       return values;
     }
+
+    function _normPrincipal(principal: { [key: string]: any[] }) {
+      const keys = Object.keys(principal);
+      if (keys.length === 0) { return undefined; }
+      const result: any = {};
+      for (const key of keys) {
+        const normVal = _norm(principal[key]);
+        if (normVal) {
+          result[key] = normVal;
+        }
+      }
+      if (Object.keys(result).length === 1 && result.AWS === '*') {
+        return '*';
+      }
+      return result;
+    }
   }
 }
 

packages/@aws-cdk/aws-iam/test/test.policy-document.ts

@@ -1,6 +1,6 @@
 import { FnConcat, resolve } from '@aws-cdk/cdk';
 import { Test } from 'nodeunit';
-import { CanonicalUserPrincipal, PolicyDocument, PolicyStatement } from '../lib';
+import { Anyone, CanonicalUserPrincipal, PolicyDocument, PolicyPrincipal, PolicyStatement, PrincipalPolicyFragment } from '../lib';
 
 export = {
   'the Permission class is a programming model for iam'(test: Test) {
@@ -143,6 +143,22 @@ export = {
     test.done();
   },
 
+  'addAwsAccountPrincipal can be used multiple times'(test: Test) {
+    const p = new PolicyStatement();
+    p.addAwsAccountPrincipal('1234');
+    p.addAwsAccountPrincipal('5678');
+    test.deepEqual(resolve(p), {
+      Effect: 'Allow',
+      Principal: {
+        AWS: [
+          { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::1234:root']] },
+          { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::5678:root']] }
+        ]
+      }
+    });
+    test.done();
+  },
+
   'hasResource': {
     'false if there are no resources'(test: Test) {
       test.equal(new PolicyStatement().hasResource, false, 'hasResource should be false for an empty permission');
@@ -188,5 +204,39 @@ export = {
     p.addStatement(new PolicyStatement());
     test.equal(p.statementCount, 2);
     test.done();
-  }
+  },
+
+  'the { AWS: "*" } principal is represented as "*"'(test: Test) {
+    const p = new PolicyDocument().addStatement(new PolicyStatement().addPrincipal(new Anyone()));
+    test.deepEqual(resolve(p), { Statement: [{ Effect: 'Allow', Principal: '*' }], Version: '2012-10-17' });
+    test.done();
+  },
+
+  'addPrincipal prohibits mixing principal types'(test: Test) {
+    const s = new PolicyStatement().addAccountRootPrincipal();
+    test.throws(() => { s.addServicePrincipal('rds.amazonaws.com'); },
+                /Attempted to add principal key Service/);
+    test.throws(() => { s.addFederatedPrincipal('federation', { ConditionOp: { ConditionKey: 'ConditionValue' } }); },
+                /Attempted to add principal key Federated/);
+    test.done();
+  },
+
+  'addPrincipal correctly merges array in'(test: Test) {
+    const arrayPrincipal: PolicyPrincipal = {
+      assumeRoleAction: 'sts:AssumeRole',
+      policyFragment: () => new PrincipalPolicyFragment({ AWS: ['foo', 'bar'] }),
+    };
+    const s = new PolicyStatement().addAccountRootPrincipal()
+                                   .addPrincipal(arrayPrincipal);
+    test.deepEqual(resolve(s), {
+      Effect: 'Allow',
+      Principal: {
+        AWS: [
+          { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::', { Ref: 'AWS::AccountId' }, ':root']] },
+          'foo', 'bar'
+        ]
+      }
+    });
+    test.done();
+  },
 };