From ced6a5bee0969e625a6243102668ea2966049c96 Mon Sep 17 00:00:00 2001
From: Yan Cui <theburningmonk@gmail.com>
Date: Thu, 23 Aug 2018 23:17:18 +0100
Subject: [PATCH] - renamed IAM role template to .txt so to avoid it getting
 automatically parsed as JSON (which then has to be JSON stringified) -
 updated the IAM role template so that resource is tempalted - collect
 function ARNs from state machines, and use them to set the resource of the
 IAM role

---
 lib/deploy/stepFunctions/compileIamRole.js    | 21 ++++++++--
 .../stepFunctions/compileIamRole.test.js      | 42 +++++++++++++++++++
 ...-role-statemachine-execution-template.txt} |  2 +-
 3 files changed, 60 insertions(+), 5 deletions(-)
 rename lib/{iam-role-statemachine-execution-template.json => iam-role-statemachine-execution-template.txt} (94%)

diff --git a/lib/deploy/stepFunctions/compileIamRole.js b/lib/deploy/stepFunctions/compileIamRole.js
index ec7de664..7927cc52 100644
--- a/lib/deploy/stepFunctions/compileIamRole.js
+++ b/lib/deploy/stepFunctions/compileIamRole.js
@@ -6,23 +6,36 @@ const path = require('path');
 module.exports = {
   compileIamRole() {
     const customRolesProvided = [];
+    let functionArns = [];
     this.getAllStateMachines().forEach((stateMachineName) => {
       const stateMachineObj = this.getStateMachine(stateMachineName);
       customRolesProvided.push('role' in stateMachineObj);
+
+      const stateMachineJson = JSON.stringify(stateMachineObj);
+      const regex = new RegExp(/"Resource":"([a-z:#{}_\-.]*)"/gi);
+      let match = regex.exec(stateMachineJson);
+      while (match !== null) {
+        functionArns.push(match[1]);
+        match = regex.exec(stateMachineJson);
+      }
     });
     if (_.isEqual(_.uniq(customRolesProvided), [true])) {
       return BbPromise.resolve();
     }
+    functionArns = _.uniq(functionArns);
 
-    let iamRoleStateMachineExecutionTemplate = JSON.stringify(this.serverless.utils.readFileSync(
+    let iamRoleStateMachineExecutionTemplate = this.serverless.utils.readFileSync(
       path.join(__dirname,
         '..',
         '..',
-        'iam-role-statemachine-execution-template.json'))
+        'iam-role-statemachine-execution-template.txt')
     );
+
     iamRoleStateMachineExecutionTemplate =
-      iamRoleStateMachineExecutionTemplate.replace('[region]', this.options.region)
-      .replace('[PolicyName]', this.getStateMachinePolicyName());
+      iamRoleStateMachineExecutionTemplate
+        .replace('[region]', this.options.region)
+        .replace('[PolicyName]', this.getStateMachinePolicyName())
+        .replace('[functions]', JSON.stringify(functionArns));
 
     const iamRoleStateMachineLogicalId = this.getiamRoleStateMachineLogicalId();
     const newIamRoleStateMachineExecutionObject = {
diff --git a/lib/deploy/stepFunctions/compileIamRole.test.js b/lib/deploy/stepFunctions/compileIamRole.test.js
index ebf4d464..58be01a8 100644
--- a/lib/deploy/stepFunctions/compileIamRole.test.js
+++ b/lib/deploy/stepFunctions/compileIamRole.test.js
@@ -85,4 +85,46 @@ describe('#compileIamRole', () => {
       .provider.compiledCloudFormationTemplate.Resources.IamRoleStateMachineExecution.Type
     ).to.equal('AWS::IAM::Role');
   });
+
+  it('should give invokeFunction permission for only functions referenced by state machine', () => {
+    const helloLambda = 'arn:aws:lambda:#{AWS::Region}:#{AWS::AccountId}:function:hello';
+    const worldLambda = 'arn:aws:lambda:#{AWS::Region}:#{AWS::AccountId}:function:world';
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine1: {
+          name: 'stateMachineBeta1',
+          definition: {
+            StartAt: 'Hello',
+            States: {
+              Hello: {
+                Type: 'Task',
+                Resource: helloLambda,
+                End: true,
+              },
+            },
+          },
+        },
+        myStateMachine2: {
+          name: 'stateMachineBeta2',
+          definition: {
+            StartAt: 'World',
+            States: {
+              World: {
+                Type: 'Task',
+                Resource: worldLambda,
+                End: true,
+              },
+            },
+          },
+        },
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+    const policy = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources.IamRoleStateMachineExecution
+      .Properties.Policies[0];
+    expect(policy.PolicyDocument.Statement[0].Resource)
+      .to.be.deep.equal([helloLambda, worldLambda]);
+  });
 });
diff --git a/lib/iam-role-statemachine-execution-template.json b/lib/iam-role-statemachine-execution-template.txt
similarity index 94%
rename from lib/iam-role-statemachine-execution-template.json
rename to lib/iam-role-statemachine-execution-template.txt
index 958161b1..49b01c02 100644
--- a/lib/iam-role-statemachine-execution-template.json
+++ b/lib/iam-role-statemachine-execution-template.txt
@@ -24,7 +24,7 @@
               "Action": [
                 "lambda:InvokeFunction"
               ],
-              "Resource": "*"
+              "Resource": [functions]
             }
           ]
         }