fix(stepfunctions): distributed maps under branches

distributed maps under branch states (i.e., Parallel) do not apply the
necessary permissions to run the state.

this moves the bind functionality into state and calls it on both state
and all child states.

rather than relying on the single purpose that it is now (add
distributed map perms) and fast returning all the way out, this instead
just checks if the policy it is trying to add is in place before
proceeding and uses that condition to return immediately.

packages/aws-cdk-lib/aws-stepfunctions/lib/state-graph.ts

@@ -166,24 +166,9 @@ export class StateGraph {
    */
   public bind(stateMachine: StateMachine) {
     for (const state of this.allStates) {
-      if (DistributedMap.isDistributedMap(state)) {
-        stateMachine.role.attachInlinePolicy(new iam.Policy(stateMachine, 'DistributedMapPolicy', {
-          document: new iam.PolicyDocument({
-            statements: [
-              new iam.PolicyStatement({
-                actions: ['states:StartExecution'],
-                resources: [stateMachine.stateMachineArn],
-              }),
-              new iam.PolicyStatement({
-                actions: ['states:DescribeExecution', 'states:StopExecution'],
-                resources: [`${stateMachine.stateMachineArn}:*`],
-              }),
-            ],
-          }),
-        }));
-
-        break;
-      }
+      state.bindToStateMachine(stateMachine, {
+        isDistributedMap: DistributedMap.isDistributedMap,
+      });
     }
   }
 }

packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts

@@ -1,10 +1,18 @@
 import { IConstruct, Construct, Node } from 'constructs';
+import * as iam from '../../../aws-iam';
 import { Token } from '../../../core';
 import { Condition } from '../condition';
 import { FieldUtils } from '../fields';
 import { StateGraph } from '../state-graph';
+import { StateMachine } from '../state-machine';
 import { CatchProps, Errors, IChainable, INextable, ProcessorConfig, ProcessorMode, RetryProps } from '../types';
 
+interface BindToStateMachineProps {
+  // We create a cyclic dependency if we include a call to DistributedMap.isDistributedMap in State. Instead, accept
+  // it as a property for bindToStateMachine.
+  isDistributedMap: (x: any) => boolean;
+}
+
 /**
  * Properties shared by all states
  */
@@ -277,6 +285,36 @@ export abstract class State extends Construct implements IChainable {
     }
   }
 
+  public bindToStateMachine(stateMachine: StateMachine, props: BindToStateMachineProps) {
+    if (stateMachine.role.node.tryFindChild('DistributedMapPolicy')) {
+      return;
+    }
+
+    if (props.isDistributedMap(this)) {
+      stateMachine.role.attachInlinePolicy(new iam.Policy(stateMachine, 'DistributedMapPolicy', {
+        document: new iam.PolicyDocument({
+          statements: [
+            new iam.PolicyStatement({
+              actions: ['states:StartExecution'],
+              resources: [stateMachine.stateMachineArn],
+            }),
+            new iam.PolicyStatement({
+              actions: ['states:DescribeExecution', 'states:StopExecution'],
+              resources: [`${stateMachine.stateMachineArn}:*`],
+            }),
+          ],
+        }),
+      }));
+
+      return;
+    }
+
+    // nothing here, let's check the branches.
+    for (const graph of this.branches) {
+      graph.bind(stateMachine);
+    }
+  }
+
   /**
    * Render the state as JSON
    */

packages/aws-cdk-lib/aws-stepfunctions/test/state-graph.test.ts

@@ -0,0 +1,76 @@
+import * as assertions from '../../assertions';
+import * as cdk from '../../core';
+import * as stepfunctions from '../lib';
+
+describe('State Graph', () => {
+  test('bind adds execution permissions to state machine when distributed map is used within the primary graph', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+
+    // WHEN
+    const map = createMap(stack);
+    const stateMachine = new stepfunctions.StateMachine(stack, 'StateMachine', {
+      definitionBody: stepfunctions.DefinitionBody.fromChainable(map),
+    });
+    const stateMachineLogicalId = stack.getLogicalId(stateMachine.node.defaultChild as stepfunctions.CfnStateMachine);
+    const template = assertions.Template.fromStack(stack);
+
+    // THEN
+    template.hasResource('AWS::IAM::Policy', createPolicyProps(stateMachineLogicalId));
+  });
+
+  test('bind adds execution permissions to state machine when distributed map is used within a child graph', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+
+    // WHEN
+    const map = createMap(stack);
+    const stateMachine = new stepfunctions.StateMachine(stack, 'StateMachine', {
+      definitionBody: stepfunctions.DefinitionBody.fromChainable(new stepfunctions.Parallel(stack, 'Parallel', {
+        resultPath: '$.result',
+      }).branch(
+        map,
+      )),
+    });
+    const stateMachineLogicalId = stack.getLogicalId(stateMachine.node.defaultChild as stepfunctions.CfnStateMachine);
+    const template = assertions.Template.fromStack(stack);
+
+    // THEN
+    template.hasResource('AWS::IAM::Policy', createPolicyProps(stateMachineLogicalId));
+  });
+});
+
+function createMap(stack: cdk.Stack) {
+  return new stepfunctions.DistributedMap(stack, 'Map', {
+    maxConcurrency: 1,
+    itemsPath: stepfunctions.JsonPath.stringAt('$.inputForMap'),
+    itemSelector: {
+      foo: 'foo',
+      bar: stepfunctions.JsonPath.stringAt('$.bar'),
+    },
+  }).itemProcessor(new stepfunctions.Pass(this, 'Pass State'));
+}
+
+function createPolicyProps(stateMachineLogicalId: string) {
+  return {
+    Properties: {
+      PolicyDocument: {
+        // ensure that self-starting permission is added which is necessary for distributed maps
+        Statement: [
+          {
+            Action: 'states:StartExecution',
+            Resource: {
+              Ref: stateMachineLogicalId,
+            },
+          },
+          {
+            Action: ['states:DescribeExecution', 'states:StopExecution'],
+            Resource: {
+              'Fn::Join': ['', [{ Ref: stateMachineLogicalId }, ':*']],
+            },
+          },
+        ],
+      },
+    },
+  };
+}