feat(stepfunctions): support cross-account task invocations

packages/@aws-cdk/aws-stepfunctions/README.md

@@ -577,6 +577,28 @@ const definition = sfn.Chain
   // ...
 ```
 
+## Task Credentials
+
+Tasks are executed using the State Machine's execution role. In some cases, e.g. cross-account access, an IAM role can be assumed by the State Machine's execution role to provide access to the resource.
+This can be achieved by providing the optional `credentials` property which allows using a literal `roleArn` or a json expression to resolve the `roleArn` at runtime.
+
+```ts
+import * as lambda from '@aws-cdk/aws-lambda';
+
+declare const submitLambda: lambda.Function;
+
+const submitJob = new tasks.LambdaInvoke(this, 'Submit Job', {
+  lambdaFunction: submitLambda,
+  outputPath: '$.Payload',
+  credentials: {
+    // literal role
+    roleArn: 'arn:aws:iam::123456789012:role/role-to-invoke-lambda',
+    // or use a json expression role
+    // roleArn: sfn.JsonPath.stringAt('$.RoleArn'),
+  },
+});
+```
+
 ## State Machine Fragments
 
 It is possible to define reusable (or abstracted) mini-state machines by

packages/@aws-cdk/aws-stepfunctions/lib/states/task-base.ts

@@ -3,6 +3,7 @@ import * as iam from '@aws-cdk/aws-iam';
 import * as cdk from '@aws-cdk/core';
 import { Construct } from 'constructs';
 import { Chain } from '../chain';
+import { FieldUtils, JsonPath } from '../fields';
 import { StateGraph } from '../state-graph';
 import { CatchProps, IChainable, INextable, RetryProps } from '../types';
 import { renderJsonPath, State } from './state';
@@ -91,6 +92,16 @@ export interface TaskStateBaseProps {
    *
    */
   readonly integrationPattern?: IntegrationPattern;
+
+  /**
+   * Credentials for an IAM Role that the State Machine assumes for executing the task.
+   * This enables cross-account resource invocations.
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/concepts-access-cross-acct-resources.html
+   *
+   * @default - None (Task is executed using the State Machine's execution role)
+   */
+  readonly credentials?: Credentials;
 }
 
 /**
@@ -112,12 +123,14 @@ export abstract class TaskStateBase extends State implements INextable {
 
   private readonly timeout?: cdk.Duration;
   private readonly heartbeat?: cdk.Duration;
+  private readonly credentials?: Credentials;
 
   constructor(scope: Construct, id: string, props: TaskStateBaseProps) {
     super(scope, id, props);
     this.endStates = [this];
     this.timeout = props.timeout;
     this.heartbeat = props.heartbeat;
+    this.credentials = props.credentials;
   }
 
   /**
@@ -263,6 +276,14 @@ export abstract class TaskStateBase extends State implements INextable {
     for (const policyStatement of this.taskPolicies || []) {
       graph.registerPolicyStatement(policyStatement);
     }
+    if (this.credentials) {
+      const resource = JsonPath.isEncodedJsonPath(this.credentials.roleArn) ? '*' : this.credentials.roleArn;
+      graph.registerPolicyStatement(new iam.PolicyStatement({
+        effect: iam.Effect.ALLOW,
+        actions: ['sts:AssumeRole'],
+        resources: [resource],
+      }));
+    }
   }
 
   /**
@@ -277,6 +298,10 @@ export abstract class TaskStateBase extends State implements INextable {
     return this.metric(prefix + suffix, props);
   }
 
+  private renderCredentials() {
+    return this.credentials ? FieldUtils.renderObject({ Credentials: { RoleArn: this.credentials.roleArn } }) : undefined;
+  }
+
   private renderTaskBase() {
     return {
       Type: 'Task',
@@ -287,6 +312,7 @@ export abstract class TaskStateBase extends State implements INextable {
       OutputPath: renderJsonPath(this.outputPath),
       ResultPath: renderJsonPath(this.resultPath),
       ...this.renderResultSelector(),
+      ...this.renderCredentials(),
     };
   }
 }
@@ -347,4 +373,16 @@ export enum IntegrationPattern {
    * @see https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#connect-wait-token
    */
   WAIT_FOR_TASK_TOKEN = 'WAIT_FOR_TASK_TOKEN'
+}
+
+/**
+ * Specifies a target role assumed by the State Machine's execution role for invoking the task's resource.
+ */
+export interface Credentials {
+
+  /**
+   * The ARN of the IAM role to be assumed.
+   * Either a fixed value or a JSONPath expression can be used.
+   */
+  readonly roleArn: string;
 }

packages/@aws-cdk/aws-stepfunctions/test/private/fake-task.ts

@@ -0,0 +1,29 @@
+import * as iam from '@aws-cdk/aws-iam';
+import * as constructs from 'constructs';
+import * as sfn from '../../lib';
+
+export interface FakeTaskProps extends sfn.TaskStateBaseProps {
+  readonly metrics?: sfn.TaskMetricsConfig;
+}
+
+export class FakeTask extends sfn.TaskStateBase {
+  protected readonly taskMetrics?: sfn.TaskMetricsConfig;
+  protected readonly taskPolicies?: iam.PolicyStatement[];
+
+  constructor(scope: constructs.Construct, id: string, props: FakeTaskProps = {}) {
+    super(scope, id, props);
+    this.taskMetrics = props.metrics;
+  }
+
+  /**
+     * @internal
+     */
+  protected _renderTask(): any {
+    return {
+      Resource: 'my-resource',
+      Parameters: sfn.FieldUtils.renderObject({
+        MyParameter: 'myParameter',
+      }),
+    };
+  }
+}

packages/@aws-cdk/aws-stepfunctions/test/private/render-util.ts

@@ -9,4 +9,8 @@ import * as sfn from '../../lib';
  */
 export function render(stack: cdk.Stack, definition: sfn.IChainable) {
   return stack.resolve(new sfn.StateGraph(definition.startState, 'Test Graph').toGraphJson());
-}
+}
+
+export function renderGraph(definition: sfn.IChainable) {
+  return render(new cdk.Stack(), definition);
+}

packages/@aws-cdk/aws-stepfunctions/test/state-machine.test.ts

@@ -3,6 +3,7 @@ import * as logs from '@aws-cdk/aws-logs';
 import * as s3 from '@aws-cdk/aws-s3';
 import * as cdk from '@aws-cdk/core';
 import * as stepfunctions from '../lib';
+import { FakeTask } from './private/fake-task';
 
 describe('State Machine', () => {
   test('Instantiate Default State Machine', () => {
@@ -278,6 +279,74 @@ describe('State Machine', () => {
     });
   });
 
+  test('Instantiate a State Machine with a task assuming a literal roleArn', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+
+    // WHEN
+    new stepfunctions.StateMachine(stack, 'MyStateMachine', {
+      definition: new FakeTask(stack, 'fakeTask', { credentials: { roleArn: 'arn:aws:iam::123456789012:role/example-role' } }),
+    });
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::StepFunctions::StateMachine', {
+      DefinitionString: '{"StartAt":"fakeTask","States":{"fakeTask":{"End":true,"Type":"Task","Credentials":{"RoleArn":"arn:aws:iam::123456789012:role/example-role"},"Resource":"my-resource","Parameters":{"MyParameter":"myParameter"}}}}',
+    });
+
+    Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: [
+          {
+            Effect: 'Allow',
+            Action: 'sts:AssumeRole',
+            Resource: 'arn:aws:iam::123456789012:role/example-role',
+          },
+        ],
+        Version: '2012-10-17',
+      },
+      PolicyName: 'MyStateMachineRoleDefaultPolicyE468EB18',
+      Roles: [
+        {
+          Ref: 'MyStateMachineRoleD59FFEBC',
+        },
+      ],
+    });
+  });
+
+  test('Instantiate a State Machine with a task assuming a JSONPath roleArn', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+
+    // WHEN
+    new stepfunctions.StateMachine(stack, 'MyStateMachine', {
+      definition: new FakeTask(stack, 'fakeTask', { credentials: { roleArn: stepfunctions.JsonPath.stringAt('$.RoleArn') } }),
+    });
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::StepFunctions::StateMachine', {
+      DefinitionString: '{"StartAt":"fakeTask","States":{"fakeTask":{"End":true,"Type":"Task","Credentials":{"RoleArn.$":"$.RoleArn"},"Resource":"my-resource","Parameters":{"MyParameter":"myParameter"}}}}',
+    });
+
+    Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: [
+          {
+            Effect: 'Allow',
+            Action: 'sts:AssumeRole',
+            Resource: '*',
+          },
+        ],
+        Version: '2012-10-17',
+      },
+      PolicyName: 'MyStateMachineRoleDefaultPolicyE468EB18',
+      Roles: [
+        {
+          Ref: 'MyStateMachineRoleD59FFEBC',
+        },
+      ],
+    });
+  });
+
   describe('StateMachine.fromStateMachineArn()', () => {
     let stack: cdk.Stack;