feat(stepfunctions): support cross-account task invocations (#23012)

support configuring a role to be assumed for task invocations https://docs.aws.amazon.com/step-functions/latest/dg/concepts-access-cross-acct-resources.html

closes #22994

----

### All Submissions:

* [x] Have you followed the guidelines in our [Contributing guide?](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md)

### Adding new Unconventional Dependencies:

* [ ] This PR adds new unconventional dependencies following the process described [here](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md/#adding-new-unconventional-dependencies)

### New Features

* [x] Have you added the new feature to an [integration test](https://github.com/aws/aws-cdk/blob/main/INTEGRATION_TESTS.md)?
	* [x] Did you use `yarn integ` to deploy the infrastructure and generate the snapshot (i.e. `yarn integ` without `--dry-run`)?

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

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

@@ -577,6 +577,34 @@ 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 fixed role or a json expression to resolve the role at runtime from the task's inputs.
+
+```ts
+import * as iam from '@aws-cdk/aws-iam';
+import * as lambda from '@aws-cdk/aws-lambda';
+
+declare const submitLambda: lambda.Function;
+declare const iamRole: iam.Role;
+
+// use a fixed role for all task invocations
+const role = sfn.TaskRole.fromRole(iamRole);
+// or use a json expression to resolve the role at runtime based on task inputs
+//const role = sfn.TaskRole.fromRoleArnJsonPath('$.RoleArn');
+
+const submitJob = new tasks.LambdaInvoke(this, 'Submit Job', {
+  lambdaFunction: submitLambda,
+  outputPath: '$.Payload',
+  // use credentials
+  credentials: { role },
+});
+```
+
+See [the AWS documentation](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-access-cross-acct-resources.html)
+to learn more about AWS Step Functions support for accessing resources in other AWS accounts.
+
 ## State Machine Fragments
 
 It is possible to define reusable (or abstracted) mini-state machines by

packages/@aws-cdk/aws-stepfunctions/lib/index.ts

@@ -22,6 +22,7 @@ export * from './states/map';
 export * from './states/custom-state';
 
 export * from './states/task-base';
+export * from './task-credentials';
 
 // AWS::StepFunctions CloudFormation Resources:
 export * from './stepfunctions.generated';

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

@@ -3,7 +3,9 @@ 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 } from '../fields';
 import { StateGraph } from '../state-graph';
+import { Credentials } from '../task-credentials';
 import { CatchProps, IChainable, INextable, RetryProps } from '../types';
 import { renderJsonPath, State } from './state';
 
@@ -91,6 +93,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 +124,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 +277,13 @@ export abstract class TaskStateBase extends State implements INextable {
     for (const policyStatement of this.taskPolicies || []) {
       graph.registerPolicyStatement(policyStatement);
     }
+    if (this.credentials) {
+      graph.registerPolicyStatement(new iam.PolicyStatement({
+        effect: iam.Effect.ALLOW,
+        actions: ['sts:AssumeRole'],
+        resources: [this.credentials.role.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.role.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,4 @@ 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'
-}
+}

packages/@aws-cdk/aws-stepfunctions/lib/task-credentials.ts

@@ -0,0 +1,77 @@
+import * as iam from '@aws-cdk/aws-iam';
+import { JsonPath } from './fields';
+
+/**
+ * Specifies a target role assumed by the State Machine's execution role for invoking the task's resource.
+ *
+ * @see https://docs.aws.amazon.com/step-functions/latest/dg/concepts-access-cross-acct-resources.html
+ * @see https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-task-state.html#task-state-fields
+ */
+export interface Credentials {
+  /**
+   * The role to be assumed for executing the Task.
+   */
+  readonly role: TaskRole;
+}
+
+/**
+ * Role to be assumed by the State Machine's execution role for invoking a task's resource.
+ *
+ * @see https://docs.aws.amazon.com/step-functions/latest/dg/concepts-access-cross-acct-resources.html
+ * @see https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-task-state.html#task-state-fields
+ */
+export abstract class TaskRole {
+  /**
+   * Construct a task role retrieved from task inputs using a json expression
+   *
+   * @param expression json expression to roleArn
+   *
+   * @example
+   *
+   * TaskRole.fromRoleArnJsonPath('$.RoleArn');
+   */
+  public static fromRoleArnJsonPath(expression: string): TaskRole {
+    return new JsonExpressionTaskRole(expression);
+  }
+
+  /**
+   * Construct a task role based on the provided IAM Role
+   *
+   * @param role IAM Role
+   */
+  public static fromRole(role: iam.IRole): TaskRole {
+    return new IamRoleTaskRole(role);
+  }
+
+  /**
+   * Retrieves the roleArn for this TaskRole
+   */
+  public abstract readonly roleArn: string;
+
+  /**
+   * Retrieves the resource for use in IAM Policies for this TaskRole
+   */
+  public abstract readonly resource: string;
+}
+
+class JsonExpressionTaskRole extends TaskRole {
+  public readonly resource: string;
+  public readonly roleArn: string;
+
+  constructor(expression: string) {
+    super();
+    this.roleArn = JsonPath.stringAt(expression);
+    this.resource = '*';
+  }
+}
+
+class IamRoleTaskRole extends TaskRole {
+  public readonly resource: string;
+  public readonly roleArn: string;
+
+  constructor(role: iam.IRole) {
+    super();
+    this.roleArn = role.roleArn;
+    this.resource = role.roleArn;
+  }
+}

packages/@aws-cdk/aws-stepfunctions/test/integ.state-machine-credentials.js.snapshot/StateMachineCredentialsDefaultTestDeployAssert3F5E6D8D.assets.json

@@ -0,0 +1,19 @@
+{
+  "version": "21.0.0",
+  "files": {
+    "21fbb51d7b23f6a6c262b46a9caee79d744a3ac019fd45422d988b96d44b2a22": {
+      "source": {
+        "path": "StateMachineCredentialsDefaultTestDeployAssert3F5E6D8D.template.json",
+        "packaging": "file"
+      },
+      "destinations": {
+        "current_account-current_region": {
+          "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
+          "objectKey": "21fbb51d7b23f6a6c262b46a9caee79d744a3ac019fd45422d988b96d44b2a22.json",
+          "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
+        }
+      }
+    }
+  },
+  "dockerImages": {}
+}

packages/@aws-cdk/aws-stepfunctions/test/integ.state-machine-credentials.js.snapshot/StateMachineCredentialsDefaultTestDeployAssert3F5E6D8D.template.json

@@ -0,0 +1,36 @@
+{
+ "Parameters": {
+  "BootstrapVersion": {
+   "Type": "AWS::SSM::Parameter::Value<String>",
+   "Default": "/cdk-bootstrap/hnb659fds/version",
+   "Description": "Version of the CDK Bootstrap resources in this environment, automatically retrieved from SSM Parameter Store. [cdk:skip]"
+  }
+ },
+ "Rules": {
+  "CheckBootstrapVersion": {
+   "Assertions": [
+    {
+     "Assert": {
+      "Fn::Not": [
+       {
+        "Fn::Contains": [
+         [
+          "1",
+          "2",
+          "3",
+          "4",
+          "5"
+         ],
+         {
+          "Ref": "BootstrapVersion"
+         }
+        ]
+       }
+      ]
+     },
+     "AssertDescription": "CDK bootstrap stack version 6 required. Please run 'cdk bootstrap' with a recent version of the CDK CLI."
+    }
+   ]
+  }
+ }
+}

packages/@aws-cdk/aws-stepfunctions/test/integ.state-machine-credentials.js.snapshot/aws-stepfunctions-state-machine-credentials-integ.assets.json

@@ -0,0 +1,19 @@
+{
+  "version": "21.0.0",
+  "files": {
+    "d775f19c6469457d54fcd62837c1d84ec75c1b8aea7b635bb10dc74dcc0e474d": {
+      "source": {
+        "path": "aws-stepfunctions-state-machine-credentials-integ.template.json",
+        "packaging": "file"
+      },
+      "destinations": {
+        "current_account-current_region": {
+          "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
+          "objectKey": "d775f19c6469457d54fcd62837c1d84ec75c1b8aea7b635bb10dc74dcc0e474d.json",
+          "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
+        }
+      }
+    }
+  },
+  "dockerImages": {}
+}