From 57e1723146d2a65ff4d6d52ff22d56152d99bd8b Mon Sep 17 00:00:00 2001
From: Todd Hainsworth <todd.hainsworth@aligent.com.au>
Date: Tue, 31 Jan 2023 08:31:38 +1030
Subject: [PATCH] feat: support custom state machine names when using a
 distributed mapA

---
 lib/deploy/stepFunctions/compileIamRole.js    |  9 ++--
 .../stepFunctions/compileIamRole.test.js      | 51 +++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/lib/deploy/stepFunctions/compileIamRole.js b/lib/deploy/stepFunctions/compileIamRole.js
index fcb2fa90..2c4aebb1 100644
--- a/lib/deploy/stepFunctions/compileIamRole.js
+++ b/lib/deploy/stepFunctions/compileIamRole.js
@@ -597,14 +597,15 @@ module.exports = {
     logger.config(this.serverless, this.v3Api);
     const service = this.serverless.service;
     const permissionsBoundary = service.provider.rolePermissionsBoundary;
-    this.getAllStateMachines().forEach((stateMachineName) => {
-      const stateMachineObj = this.getStateMachine(stateMachineName);
+    this.getAllStateMachines().forEach((stateMachineId) => {
+      const stateMachineObj = this.getStateMachine(stateMachineId);
+      const stateMachineName = stateMachineObj.name || stateMachineId;
       if (stateMachineObj.role) {
         return;
       }
 
       if (!stateMachineObj.definition) {
-        throw new Error(`Missing "definition" for state machine ${stateMachineName}`);
+        throw new Error(`Missing "definition" for state machine ${stateMachineId}`);
       }
 
       const taskStates = getTaskStates(stateMachineObj.definition.States, stateMachineName);
@@ -646,7 +647,7 @@ module.exports = {
       }
 
       const stateMachineLogicalId = this.getStateMachineLogicalId(
-        stateMachineName,
+        stateMachineId,
         stateMachineObj,
       );
       const iamRoleStateMachineLogicalId = `${stateMachineLogicalId}Role`;
diff --git a/lib/deploy/stepFunctions/compileIamRole.test.js b/lib/deploy/stepFunctions/compileIamRole.test.js
index 7c75fb83..f049d232 100644
--- a/lib/deploy/stepFunctions/compileIamRole.test.js
+++ b/lib/deploy/stepFunctions/compileIamRole.test.js
@@ -2344,6 +2344,57 @@ describe('#compileIamRole', () => {
     ]);
   });
 
+  it('should support custom state machine names in a Distributed Map', () => {
+    const getStateMachine = (id, lambdaArn) => ({
+      id,
+      name: 'DistributedMapper',
+      definition: {
+        StartAt: 'A',
+        States: {
+          A: {
+            Type: 'Map',
+            ItemProcessor: {
+              ProcessorConfig: {
+                Mode: 'DISTRIBUTED',
+              },
+              StartAt: 'B',
+              States: {
+                B: {
+                  Type: 'Task',
+                  Resource: lambdaArn,
+                  End: true,
+                },
+              },
+            },
+            End: true,
+          },
+        },
+      },
+    });
+
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine: getStateMachine('StateMachine1', 'arn:aws:lambda:us-west-2:1234567890:function:foo'),
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+
+    const statements = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources.StateMachine1Role
+      .Properties.Policies[0].PolicyDocument.Statement;
+
+    const stepFunctionPermission = statements.filter(s => _.isEqual(s.Action, ['states:StartExecution']));
+    expect(stepFunctionPermission).to.have.lengthOf(1);
+    expect(stepFunctionPermission[0].Resource).to.deep.eq([{
+      'Fn::Sub': [
+        'arn:aws:states:${AWS::Region}:${AWS::AccountId}:stateMachine:DistributedMapper',
+        {},
+      ],
+    },
+    ]);
+  });
+
   it('should support nested Map state type', () => {
     const getStateMachine = (id, lambdaArn1, lambdaArn2) => ({
       id,