From 10f168e115ead22fdbff791e84e20d77682ef54d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=F0=9F=A7=91=F0=9F=8F=BB=E2=80=8D=F0=9F=92=BB=20Romain=20M?=
 =?UTF-8?q?arcadier?= <rmuller@amazon.com>
Date: Mon, 13 Mar 2023 14:23:47 +0100
Subject: [PATCH] fix(sfn): can't override toStateJson() from other languages

If any part of a state's JSON representation is `null`, that value will
be replaced by `undefined` when jsii sends data to the other language,
resulting in a change of semantics.

Multi-language APIs cannot differentiate between `null` and `undefined`
as non-JS languages typically fail to distinguish between them... In
order to address that, a `JsonNull` value was added which serializes to
`null` (via Javascript's standard `toJSON` method), which must be used
in such cases where `null` may need to cross the language boundary.

The `JsonPath.DISCARD` value is now a string-token representation of the
`JsonNull` instance.

Fixes #14639
---
 .../@aws-cdk/aws-stepfunctions/lib/fields.ts  |  6 ++--
 .../aws-stepfunctions/lib/states/state.ts     |  3 +-
 .../aws-stepfunctions/test/state.test.ts      | 30 +++++++++++++++++++
 packages/@aws-cdk/core/lib/token.ts           | 21 +++++++++++++
 4 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 packages/@aws-cdk/aws-stepfunctions/test/state.test.ts

diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts b/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts
index 0d152d15c9656..2d1405ffd1b72 100644
--- a/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts
@@ -1,4 +1,4 @@
-import { Token, IResolvable } from '@aws-cdk/core';
+import { Token, IResolvable, JsonNull } from '@aws-cdk/core';
 import { findReferencedPaths, jsonPathString, JsonPathToken, renderObject, renderInExpression, jsonPathFromAny } from './private/json-path';
 
 /**
@@ -9,9 +9,9 @@ import { findReferencedPaths, jsonPathString, JsonPathToken, renderObject, rende
  */
 export class JsonPath {
   /**
-   * Special string value to discard state input, output or result
+   * Special string value to discard state input, output or result.
    */
-  public static readonly DISCARD = 'DISCARD';
+  public static readonly DISCARD = Token.asString(JsonNull.INSTANCE, { displayHint: 'DISCARD (JSON `null`)' });
 
   /**
    * Instead of using a literal string, get the value from a JSON path
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/states/state.ts b/packages/@aws-cdk/aws-stepfunctions/lib/states/state.ts
index 5fd6a7f0b2ce8..66869e282b163 100644
--- a/packages/@aws-cdk/aws-stepfunctions/lib/states/state.ts
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/states/state.ts
@@ -1,3 +1,4 @@
+import { Token } from '@aws-cdk/core';
 import { IConstruct, Construct, Node } from 'constructs';
 import { Condition } from '../condition';
 import { FieldUtils, JsonPath } from '../fields';
@@ -579,7 +580,7 @@ export function renderJsonPath(jsonPath?: string): undefined | null | string {
   if (jsonPath === undefined) { return undefined; }
   if (jsonPath === JsonPath.DISCARD) { return null; }
 
-  if (!jsonPath.startsWith('$')) {
+  if (!Token.isUnresolved(jsonPath) && !jsonPath.startsWith('$')) {
     throw new Error(`Expected JSON path to start with '$', got: ${jsonPath}`);
   }
   return jsonPath;
diff --git a/packages/@aws-cdk/aws-stepfunctions/test/state.test.ts b/packages/@aws-cdk/aws-stepfunctions/test/state.test.ts
new file mode 100644
index 0000000000000..db13dfb573412
--- /dev/null
+++ b/packages/@aws-cdk/aws-stepfunctions/test/state.test.ts
@@ -0,0 +1,30 @@
+import * as cdk from '@aws-cdk/core';
+import { FakeTask } from './integ.state-machine-credentials';
+import { renderGraph } from './private/render-util';
+import { JsonPath } from '../lib';
+
+test('JsonPath.DISCARD can be used to discard a state\'s output', () => {
+  const stack = new cdk.Stack();
+
+  const task = new FakeTask(stack, 'my-state', {
+    inputPath: JsonPath.DISCARD,
+    outputPath: JsonPath.DISCARD,
+    resultPath: JsonPath.DISCARD,
+  });
+
+  expect(renderGraph(task)).toEqual({
+    StartAt: 'my-state',
+    States: {
+      'my-state': {
+        End: true,
+        Type: 'Task',
+        Resource: expect.any(String),
+        Parameters: expect.any(Object),
+        // The important bits:
+        InputPath: null,
+        OutputPath: null,
+        ResultPath: null,
+      },
+    },
+  });
+});
diff --git a/packages/@aws-cdk/core/lib/token.ts b/packages/@aws-cdk/core/lib/token.ts
index af4bdcf11eacb..0ff003b2c28c5 100644
--- a/packages/@aws-cdk/core/lib/token.ts
+++ b/packages/@aws-cdk/core/lib/token.ts
@@ -232,6 +232,27 @@ export class Tokenization {
   }
 }
 
+/**
+ * An object which serializes to the JSON `null` literal, and which can safely
+ * be passed across languages where `undefined` and `null` are not different.
+ */
+export class JsonNull {
+  /** The canonical instance of `JsonNull`. */
+  public static readonly INSTANCE = new JsonNull();
+
+  private constructor() { }
+
+  /** Obtains the JSON representation of this object (`null`) */
+  public toJSON(): any {
+    return null;
+  }
+
+  /** Obtains the string representation of this object (`'null'`) */
+  public toString(): string {
+    return 'null';
+  }
+}
+
 /**
  * Options for the 'reverse()' operation
  */