fix(iam): Merge multiple principals correctly (#983)

When adidng multiple principals, an array of principals was created, but
IAM documents expect an object with different keys, for which values may
be arrays. This corrects the merging process so it produces valid
statements.

This restores the fix for #924 that was introduced in #916 and reverted
in #958 due to an other feature part of the same commit.

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