From e222e87e31873b0dae4ee9ede5be2bccbd5db853 Mon Sep 17 00:00:00 2001
From: Iain Cole <iaincole1@gmail.com>
Date: Mon, 5 Aug 2019 11:13:25 -0700
Subject: [PATCH] fix(lambda): allow ArnPrincipal in grantInvoke (#3501)

Fixes #3264

I'm trying to allow a lambda function in another account to be able to invoke my CDK generated lambda function. This works through the CLI like so:

    aws lambda add-permission --function-name=myFunction --statement-id=ABoldStatement --action=lambda:InvokeFunction --principal=arn:aws:iam::{account_id}:role/a_lambda_execution_role

But CDK doesn't seem to allow me to add an ArnPrincipal doing something like this:

    myFunction.grantInvoke(new iam.ArnPrincipal(props.myARN))

With the error:

    Invalid principal type for Lambda permission statement: ArnPrincipal. Supported: AccountPrincipal, ServicePrincipal

This PR allows ArnPrincipal to be passed to lambda.grantInvoke.

There might be some additional validation required on the exact ARN as I believe only some ARNs are supported by lambda add-permission
---
 .../@aws-cdk/aws-lambda/lib/function-base.ts  |  6 +++-
 .../@aws-cdk/aws-lambda/test/test.lambda.ts   | 33 +++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/packages/@aws-cdk/aws-lambda/lib/function-base.ts b/packages/@aws-cdk/aws-lambda/lib/function-base.ts
index 623fabb68820a..3bbda22f1b465 100644
--- a/packages/@aws-cdk/aws-lambda/lib/function-base.ts
+++ b/packages/@aws-cdk/aws-lambda/lib/function-base.ts
@@ -285,8 +285,12 @@ export abstract class FunctionBase extends Resource implements IFunction {
       return (principal as iam.ServicePrincipal).service;
     }
 
+    if (`arn` in principal) {
+      return (principal as iam.ArnPrincipal).arn;
+    }
+
     throw new Error(`Invalid principal type for Lambda permission statement: ${principal.constructor.name}. ` +
-      'Supported: AccountPrincipal, ServicePrincipal');
+      'Supported: AccountPrincipal, ArnPrincipal, ServicePrincipal');
   }
 }
 
diff --git a/packages/@aws-cdk/aws-lambda/test/test.lambda.ts b/packages/@aws-cdk/aws-lambda/test/test.lambda.ts
index 09a36697952a9..7cb3fee771356 100644
--- a/packages/@aws-cdk/aws-lambda/test/test.lambda.ts
+++ b/packages/@aws-cdk/aws-lambda/test/test.lambda.ts
@@ -186,15 +186,16 @@ export = {
       test.done();
     },
 
-    'fails if the principal is not a service or account principals'(test: Test) {
+    'fails if the principal is not a service, account or arn principal'(test: Test) {
       const stack = new cdk.Stack();
       const fn = newTestLambda(stack);
 
-      test.throws(() => fn.addPermission('F1', { principal: new iam.ArnPrincipal('just:arn') }),
+      test.throws(() => fn.addPermission('F1', { principal: new iam.OrganizationPrincipal('org') }),
         /Invalid principal type for Lambda permission statement/);
 
       fn.addPermission('S1', { principal: new iam.ServicePrincipal('my-service') });
       fn.addPermission('S2', { principal: new iam.AccountPrincipal('account') });
+      fn.addPermission('S3', { principal: new iam.ArnPrincipal('my:arn') });
 
       test.done();
     },
@@ -1074,6 +1075,34 @@ export = {
     test.done();
   },
 
+  'grantInvoke with an arn principal'(test: Test) {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const fn = new lambda.Function(stack, 'Function', {
+      code: lambda.Code.inline('xxx'),
+      handler: 'index.handler',
+      runtime: lambda.Runtime.NODEJS_8_10,
+    });
+    const account = new iam.ArnPrincipal('arn:aws:iam::123456789012:role/someRole');
+
+    // WHEN
+    fn.grantInvoke(account);
+
+    // THEN
+    expect(stack).to(haveResource('AWS::Lambda::Permission', {
+      Action: 'lambda:InvokeFunction',
+      FunctionName: {
+        'Fn::GetAtt': [
+          'Function76856677',
+          'Arn'
+        ]
+      },
+      Principal: 'arn:aws:iam::123456789012:role/someRole'
+    }));
+
+    test.done();
+  },
+
   'Can use metricErrors on a lambda Function'(test: Test) {
     // GIVEN
     const stack = new cdk.Stack();