From 24503ecba46babf6b29722c844ab0c9354458cd0 Mon Sep 17 00:00:00 2001
From: Luca Pizzini <lpizzini7@gmail.com>
Date: Tue, 20 Jun 2023 14:41:14 +0200
Subject: [PATCH] fix(stepfunctions): nested arrays are not serialized
 correctly

---
 .../lib/private/json-path.ts                  | 27 +++++++-----
 .../aws-stepfunctions/test/fields.test.ts     | 42 +++++++++++++++++++
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions/lib/private/json-path.ts b/packages/aws-cdk-lib/aws-stepfunctions/lib/private/json-path.ts
index 484ccc1552c1e..8a7be470b8b08 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions/lib/private/json-path.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions/lib/private/json-path.ts
@@ -183,18 +183,25 @@ function recurseArray(key: string, arr: any[], handlers: FieldHandlers, visited:
   }
 
   return {
-    [key]: arr.map(value => {
-      if ((typeof value === 'string' && jsonPathString(value) !== undefined)
+    [key]: resolveArray(arr, handlers, visited),
+  };
+}
+
+function resolveArray(arr: any[], handlers: FieldHandlers, visited: object[] = []): any[] {
+  return arr.map(value => {
+    if ((typeof value === 'string' && jsonPathString(value) !== undefined)
         || (typeof value === 'number' && jsonPathNumber(value) !== undefined)
         || (isStringArray(value) && jsonPathStringList(value) !== undefined)) {
-        throw new Error('Cannot use JsonPath fields in an array, they must be used in objects');
-      }
-      if (typeof value === 'object' && value !== null) {
-        return recurseObject(value, handlers, visited);
-      }
-      return value;
-    }),
-  };
+      throw new Error('Cannot use JsonPath fields in an array, they must be used in objects');
+    }
+    if (Array.isArray(value)) {
+      return resolveArray(value, handlers, visited);
+    }
+    if (typeof value === 'object' && value !== null) {
+      return recurseObject(value, handlers, visited);
+    }
+    return value;
+  });
 }
 
 function isStringArray(x: any): x is string[] {
diff --git a/packages/aws-cdk-lib/aws-stepfunctions/test/fields.test.ts b/packages/aws-cdk-lib/aws-stepfunctions/test/fields.test.ts
index 024e01e0025df..062a226266f3b 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions/test/fields.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions/test/fields.test.ts
@@ -470,6 +470,48 @@ describe('intrinsics constructors', () => {
       'Field.$': 'States.JsonToString($.Obj)',
     });
   });
+
+  test('correctly serialize a nested array', () => {
+    expect(
+      FieldUtils.renderObject({
+        nestedArray: [
+          [
+            [123, 123],
+            [456, 456],
+          ],
+        ],
+      }),
+    ).toStrictEqual({
+      nestedArray: [
+        [
+          [123, 123],
+          [456, 456],
+        ],
+      ],
+    });
+  });
+
+  test('deep replace correctly handles fields in nested arrays', () => {
+    expect(
+      FieldUtils.renderObject({
+        deep: [
+          [
+            {
+              deepField: JsonPath.numberAt('$.numField'),
+            },
+          ],
+        ],
+      }),
+    ).toStrictEqual({
+      deep: [
+        [
+          {
+            'deepField.$': '$.numField',
+          },
+        ],
+      ],
+    });
+  });
 });
 
 test('find task token even if nested in intrinsic functions', () => {