From 521e1541d25d5181445734c48c10537ddb198777 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 ++-
 packages/@aws-cdk/core/lib/token.ts           | 19 +++++++++++++++++++
 3 files changed, 24 insertions(+), 4 deletions(-)

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/core/lib/token.ts b/packages/@aws-cdk/core/lib/token.ts
index af4bdcf11eacb..99c4499f708bb 100644
--- a/packages/@aws-cdk/core/lib/token.ts
+++ b/packages/@aws-cdk/core/lib/token.ts
@@ -232,6 +232,25 @@ 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() { }
+
+  public toJSON(): any {
+    return null;
+  }
+
+  public toString(): string {
+    return 'null';
+  }
+}
+
 /**
  * Options for the 'reverse()' operation
  */