From 87af50c9bc9ac3bd4f11890a66dafdc024b2b7d9 Mon Sep 17 00:00:00 2001
From: corymhall <43035978+corymhall@users.noreply.github.com>
Date: Thu, 8 Sep 2022 18:12:56 +0000
Subject: [PATCH] fix(lambda-event-sources): cannot add sqs event source to an
 imported function

If an SQS event sources is added to an imported function it will throw
an error if the function is not imported with an IAM role.

This PR updates the logic to only attempt to add permissions to the
principal if the role exists, otherwise it will add a warning indicating
that permissions were not added.

fixes #12607
---
 .../aws-lambda-event-sources/lib/sqs.ts       |  11 +-
 .../aws-lambda-event-sources/test/sqs.test.ts | 110 ++++++++++++++++++
 2 files changed, 119 insertions(+), 2 deletions(-)

diff --git a/packages/@aws-cdk/aws-lambda-event-sources/lib/sqs.ts b/packages/@aws-cdk/aws-lambda-event-sources/lib/sqs.ts
index c6cba08abd13d..0e4f39c7e3858 100644
--- a/packages/@aws-cdk/aws-lambda-event-sources/lib/sqs.ts
+++ b/packages/@aws-cdk/aws-lambda-event-sources/lib/sqs.ts
@@ -1,6 +1,6 @@
 import * as lambda from '@aws-cdk/aws-lambda';
 import * as sqs from '@aws-cdk/aws-sqs';
-import { Duration, Names, Token } from '@aws-cdk/core';
+import { Duration, Names, Token, Annotations } from '@aws-cdk/core';
 
 export interface SqsEventSourceProps {
   /**
@@ -76,7 +76,14 @@ export class SqsEventSource implements lambda.IEventSource {
     });
     this._eventSourceMappingId = eventSourceMapping.eventSourceMappingId;
 
-    this.queue.grantConsumeMessages(target);
+    // only grant access if the lambda function has an IAM role
+    // otherwise the IAM module will throw an error
+    if (target.role) {
+      this.queue.grantConsumeMessages(target);
+    } else {
+      Annotations.of(target).addWarning(`Function '${target.node.path}' was imported without an IAM role `+
+        `so it was not granted access to consume messages from '${this.queue.node.path}'`);
+    }
   }
 
   /**
diff --git a/packages/@aws-cdk/aws-lambda-event-sources/test/sqs.test.ts b/packages/@aws-cdk/aws-lambda-event-sources/test/sqs.test.ts
index 0f019671460af..06a834f2f8c36 100644
--- a/packages/@aws-cdk/aws-lambda-event-sources/test/sqs.test.ts
+++ b/packages/@aws-cdk/aws-lambda-event-sources/test/sqs.test.ts
@@ -1,6 +1,9 @@
 import { Template } from '@aws-cdk/assertions';
+import * as iam from '@aws-cdk/aws-iam';
+import * as lambda from '@aws-cdk/aws-lambda';
 import * as sqs from '@aws-cdk/aws-sqs';
 import * as cdk from '@aws-cdk/core';
+import { App } from '@aws-cdk/core';
 import * as sources from '../lib';
 import { TestFunction } from './test-function';
 
@@ -281,7 +284,114 @@ describe('SQSEventSource', () => {
     Template.fromStack(stack).hasResourceProperties('AWS::Lambda::EventSourceMapping', {
       'FunctionResponseTypes': ['ReportBatchItemFailures'],
     });
+  });
+
+  test('warning added if lambda function imported without role', () => {
+    const app = new App();
+    const stack = new cdk.Stack(app);
+    const fn = lambda.Function.fromFunctionName(stack, 'Handler', 'testFunction');
+    const q = new sqs.Queue(stack, 'Q');
+
+    // WHEN
+    fn.addEventSource(new sources.SqsEventSource(q));
+    const assembly = app.synth();
+
+    const messages = assembly.getStackArtifact(stack.artifactId).messages;
+
+    // THEN
+    expect(messages.length).toEqual(1);
+    expect(messages[0]).toMatchObject({
+      level: 'warning',
+      id: '/Default/Handler',
+      entry: {
+        data: expect.stringMatching(/Function 'Default\/Handler' was imported without an IAM role/),
+      },
+    });
 
+    // THEN
+    Template.fromStack(stack).resourceCountIs('AWS::Lambda::EventSourceMapping', 1);
+    Template.fromStack(stack).resourceCountIs('AWS::IAM::Policy', 0);
+  });
 
+  test('policy added to imported function role', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const fn = lambda.Function.fromFunctionAttributes(stack, 'Handler', {
+      functionArn: stack.formatArn({
+        service: 'lambda',
+        resource: 'function',
+        resourceName: 'testFunction',
+      }),
+      role: iam.Role.fromRoleName(stack, 'Role', 'testFunctionRole'),
+    });
+    const q = new sqs.Queue(stack, 'Q');
+
+    // WHEN
+    fn.addEventSource(new sources.SqsEventSource(q));
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
+      'PolicyDocument': {
+        'Statement': [
+          {
+            'Action': [
+              'sqs:ReceiveMessage',
+              'sqs:ChangeMessageVisibility',
+              'sqs:GetQueueUrl',
+              'sqs:DeleteMessage',
+              'sqs:GetQueueAttributes',
+            ],
+            'Effect': 'Allow',
+            'Resource': {
+              'Fn::GetAtt': [
+                'Q63C6E3AB',
+                'Arn',
+              ],
+            },
+          },
+        ],
+        'Version': '2012-10-17',
+      },
+      'Roles': ['testFunctionRole'],
+    });
+
+    Template.fromStack(stack).hasResourceProperties('AWS::Lambda::EventSourceMapping', {
+      'EventSourceArn': {
+        'Fn::GetAtt': [
+          'Q63C6E3AB',
+          'Arn',
+        ],
+      },
+      'FunctionName': {
+        'Fn::Select': [
+          6,
+          {
+            'Fn::Split': [
+              ':',
+              {
+                'Fn::Join': [
+                  '',
+                  [
+                    'arn:',
+                    {
+                      'Ref': 'AWS::Partition',
+                    },
+                    ':lambda:',
+                    {
+                      'Ref': 'AWS::Region',
+                    },
+                    ':',
+                    {
+                      'Ref': 'AWS::AccountId',
+                    },
+                    ':function/testFunction',
+                  ],
+                ],
+              },
+            ],
+          },
+        ],
+      },
+    });
   });
 });