fix(stepfunctions): task token integration cannot be used with API Gateway (#18595)

To pass the Task Token in headers to an API Gateway, the token must
be wrapped in an array (because that's the value type of headers).

Because JSONPath evaluation needs to happen to resolve the token,
we need to use the `States.Array()` function in a `JsonPathToken`
to properly resolve this. However, doing that makes the existing
validation code fail the validation checking that you are passing
the task token somewhere.

Add convenience methods for the intrinsics, and update the checker
to also find paths referenced inside intrinsic functions.

Fixes #14184, fixes #14181.


----

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

Author: rix0rrr

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

@@ -208,19 +208,20 @@ and invokes it asynchronously.
 
 ```ts
 declare const fn: lambda.Function;
+
 const submitJob = new tasks.LambdaInvoke(this, 'Invoke Handler', {
   lambdaFunction: fn,
   payload: sfn.TaskInput.fromJsonPathAt('$.input'),
   invocationType: tasks.LambdaInvocationType.EVENT,
 });
 ```
 
-You can also use [intrinsic functions](https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html) with `JsonPath.stringAt()`.
+You can also use [intrinsic functions](https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html) available on `JsonPath`, for example `JsonPath.format()`.
 Here is an example of starting an Athena query that is dynamically created using the task input:
 
 ```ts
 const startQueryExecutionJob = new tasks.AthenaStartQueryExecution(this, 'Athena Start Query', {
-  queryString: sfn.JsonPath.stringAt("States.Format('select contacts where year={};', $.year)"),
+  queryString: sfn.JsonPath.format('select contacts where year={};', sfn.JsonPath.stringAt('$.year')),
   queryExecutionContext: {
     databaseName: 'interactions',
   },
@@ -305,6 +306,25 @@ const invokeTask = new tasks.CallApiGatewayRestApiEndpoint(this, 'Call REST API'
 });
 ```
 
+Be aware that the header values must be arrays. When passing the Task Token
+in the headers field `WAIT_FOR_TASK_TOKEN` integration, use
+`JsonPath.array()` to wrap the token in an array:
+
+```ts
+import * as apigateway from '@aws-cdk/aws-apigateway';
+declare const api: apigateway.RestApi;
+
+new tasks.CallApiGatewayRestApiEndpoint(this, 'Endpoint', {
+  api,
+  stageName: 'Stage',
+  method: tasks.HttpMethod.PUT,
+  integrationPattern: sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN,
+  headers: sfn.TaskInput.fromObject({
+    TaskToken: sfn.JsonPath.array(sfn.JsonPath.taskToken),
+  }),
+});
+```
+
 ### Call HTTP API Endpoint
 
 The `CallApiGatewayHttpApiEndpoint` calls the HTTP API endpoint.
@@ -798,7 +818,7 @@ The service integration APIs correspond to Amazon EMR on EKS APIs, but differ in
 
 ### Create Virtual Cluster
 
-The [CreateVirtualCluster](https://docs.aws.amazon.com/emr-on-eks/latest/APIReference/API_CreateVirtualCluster.html) API creates a single virtual cluster that's mapped to a single Kubernetes namespace. 
+The [CreateVirtualCluster](https://docs.aws.amazon.com/emr-on-eks/latest/APIReference/API_CreateVirtualCluster.html) API creates a single virtual cluster that's mapped to a single Kubernetes namespace.
 
 The EKS cluster containing the Kubernetes namespace where the virtual cluster will be mapped can be passed in from the task input.
 

packages/@aws-cdk/aws-stepfunctions-tasks/lib/apigateway/call-rest-api.ts

@@ -24,6 +24,25 @@ export interface CallApiGatewayRestApiEndpointProps extends CallApiGatewayEndpoi
 /**
  * Call REST API endpoint as a Task
  *
+ * Be aware that the header values must be arrays. When passing the Task Token
+ * in the headers field `WAIT_FOR_TASK_TOKEN` integration, use
+ * `JsonPath.array()` to wrap the token in an array:
+ *
+ * ```ts
+ * import * as apigateway from '@aws-cdk/aws-apigateway';
+ * declare const api: apigateway.RestApi;
+ *
+ * new tasks.CallApiGatewayRestApiEndpoint(this, 'Endpoint', {
+ *   api,
+ *   stageName: 'Stage',
+ *   method: tasks.HttpMethod.PUT,
+ *   integrationPattern: sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN,
+ *   headers: sfn.TaskInput.fromObject({
+ *     TaskToken: sfn.JsonPath.array(sfn.JsonPath.taskToken),
+ *   }),
+ * });
+ * ```
+ *
  * @see https://docs.aws.amazon.com/step-functions/latest/dg/connect-api-gateway.html
  */
 export class CallApiGatewayRestApiEndpoint extends CallApiGatewayEndpointBase {

packages/@aws-cdk/aws-stepfunctions-tasks/test/apigateway/call-rest-api.test.ts

@@ -69,7 +69,7 @@ describe('CallApiGatewayRestApiEndpoint', () => {
       method: HttpMethod.GET,
       stageName: 'dev',
       integrationPattern: sfn.IntegrationPattern.WAIT_FOR_TASK_TOKEN,
-      headers: sfn.TaskInput.fromObject({ TaskToken: sfn.JsonPath.taskToken }),
+      headers: sfn.TaskInput.fromObject({ TaskToken: sfn.JsonPath.array(sfn.JsonPath.taskToken) }),
     });
 
     // THEN
@@ -97,7 +97,7 @@ describe('CallApiGatewayRestApiEndpoint', () => {
         },
         AuthType: 'NO_AUTH',
         Headers: {
-          'TaskToken.$': '$$.Task.Token',
+          'TaskToken.$': 'States.Array($$.Task.Token)',
         },
         Method: 'GET',
         Stage: 'dev',

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

@@ -95,6 +95,51 @@ properly (for example, permissions to invoke any Lambda functions you add to
 your workflow). A role will be created by default, but you can supply an
 existing one as well.
 
+## Accessing State (the JsonPath class)
+
+Every State Machine execution has [State Machine
+Data](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-state-machine-data.html):
+a JSON document containing keys and values that is fed into the state machine,
+gets modified as the state machine progresses, and finally is produced as output.
+
+You can pass fragments of this State Machine Data into Tasks of the state machine.
+To do so, use the static methods on the `JsonPath` class. For example, to pass
+the value that's in the data key of `OrderId` to a Lambda function as you invoke
+it, use `JsonPath.stringAt('$.OrderId')`, like so:
+
+```ts
+import * as lambda from '@aws-cdk/aws-lambda';
+
+declare const orderFn: lambda.Function;
+
+const submitJob = new tasks.LambdaInvoke(this, 'InvokeOrderProcessor', {
+  lambdaFunction: orderFn,
+  payload: sfn.TaskInput.fromObject({
+    OrderId: sfn.JsonPath.stringAt('$.OrderId'),
+  }),
+});
+```
+
+The following methods are available:
+
+| Method | Purpose |
+|--------|---------|
+| `JsonPath.stringAt('$.Field')` | reference a field, return the type as a `string`. |
+| `JsonPath.listAt('$.Field')` | reference a field, return the type as a list of strings. |
+| `JsonPath.numberAt('$.Field')` | reference a field, return the type as a number. Use this for functions that expect a number argument. |
+| `JsonPath.objectAt('$.Field')` | reference a field, return the type as an `IResolvable`. Use this for functions that expect an object argument. |
+| `JsonPath.entirePayload` | reference the entire data object (equivalent to a path of `$`). |
+| `JsonPath.taskToken` | reference the [Task Token](https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#connect-wait-token), used for integration patterns that need to run for a long time. |
+
+You can also call [intrinsic functions](https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html) using the methods on `JsonPath`:
+
+| Method | Purpose |
+|--------|---------|
+| `JsonPath.array(JsonPath.stringAt('$.Field'), ...)` | make an array from other elements. |
+| `JsonPath.format('The value is {}.', JsonPath.stringAt('$.Value'))` | insert elements into a format string. |
+| `JsonPath.stringToJson(JsonPath.stringAt('$.ObjStr'))` | parse a JSON string to an object |
+| `JsonPath.jsonToString(JsonPath.objectAt('$.Obj'))` | stringify an object to a JSON string |
+
 ## Amazon States Language
 
 This library comes with a set of classes that model the [Amazon States
@@ -603,8 +648,8 @@ new cloudwatch.Alarm(this, 'ThrottledAlarm', {
 
 ## Error names
 
-Step Functions identifies errors in the Amazon States Language using case-sensitive strings, known as error names. 
-The Amazon States Language defines a set of built-in strings that name well-known errors, all beginning with the `States.` prefix. 
+Step Functions identifies errors in the Amazon States Language using case-sensitive strings, known as error names.
+The Amazon States Language defines a set of built-in strings that name well-known errors, all beginning with the `States.` prefix.
 
 * `States.ALL` - A wildcard that matches any known error name.
 * `States.Runtime` - An execution failed due to some exception that could not be processed. Often these are caused by errors at runtime, such as attempting to apply InputPath or OutputPath on a null JSON payload. A `States.Runtime` error is not retriable, and will always cause the execution to fail. A retry or catch on `States.ALL` will NOT catch States.Runtime errors.

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

@@ -1,5 +1,5 @@
-import { Token } from '@aws-cdk/core';
-import { findReferencedPaths, jsonPathString, JsonPathToken, renderObject } from './json-path';
+import { Token, IResolvable } from '@aws-cdk/core';
+import { findReferencedPaths, jsonPathString, JsonPathToken, renderObject, renderInExpression, jsonPathFromAny } from './private/json-path';
 
 /**
  * Extract a field from the State Machine data or context
@@ -38,6 +38,14 @@ export class JsonPath {
     return Token.asNumber(new JsonPathToken(path));
   }
 
+  /**
+   * Reference a complete (complex) object in a JSON path location
+   */
+  public static objectAt(path: string): IResolvable {
+    validateJsonPath(path);
+    return new JsonPathToken(path);
+  }
+
   /**
    * Use the entire data structure
    *
@@ -78,6 +86,82 @@ export class JsonPath {
     return new JsonPathToken('$$').toString();
   }
 
+  /**
+   * Make an intrinsic States.Array expression
+   *
+   * Combine any number of string literals or JsonPath expressions into an array.
+   *
+   * Use this function if the value of an array element directly has to come
+   * from a JSON Path expression (either the State object or the Context object).
+   *
+   * If the array contains object literals whose values come from a JSON path
+   * expression, you do not need to use this function.
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html
+   */
+  public static array(...values: string[]): string {
+    return new JsonPathToken(`States.Array(${values.map(renderInExpression).join(', ')})`).toString();
+  }
+
+  /**
+   * Make an intrinsic States.Format expression
+   *
+   * This can be used to embed JSON Path variables inside a format string.
+   *
+   * For example:
+   *
+   * ```ts
+   * sfn.JsonPath.format('Hello, my name is {}.', sfn.JsonPath.stringAt('$.name'))
+   * ```
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html
+   */
+  public static format(formatString: string, ...values: string[]): string {
+    const allArgs = [formatString, ...values];
+    return new JsonPathToken(`States.Format(${allArgs.map(renderInExpression).join(', ')})`).toString();
+  }
+
+  /**
+   * Make an intrinsic States.StringToJson expression
+   *
+   * During the execution of the Step Functions state machine, parse the given
+   * argument as JSON into its object form.
+   *
+   * For example:
+   *
+   * ```ts
+   * sfn.JsonPath.stringToJson(sfn.JsonPath.stringAt('$.someJsonBody'))
+   * ```
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html
+   */
+  public static stringToJson(jsonString: string): IResolvable {
+    return new JsonPathToken(`States.StringToJson(${renderInExpression(jsonString)})`);
+  }
+
+  /**
+   * Make an intrinsic States.JsonToString expression
+   *
+   * During the execution of the Step Functions state machine, encode the
+   * given object into a JSON string.
+   *
+   * For example:
+   *
+   * ```ts
+   * sfn.JsonPath.jsonToString(sfn.JsonPath.objectAt('$.someObject'))
+   * ```
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-intrinsic-functions.html
+   */
+  public static jsonToString(value: any): string {
+    const path = jsonPathFromAny(value);
+    if (!path) {
+      throw new Error('Argument to JsonPath.jsonToString() must be a JsonPath object');
+    }
+
+    return new JsonPathToken(`States.JsonToString(${path})`).toString();
+  }
+
   private constructor() {}
 }