From a9f365f4c7c97ae72b999276d7a3c20f60c96374 Mon Sep 17 00:00:00 2001
From: Yiming Luo <yiming.luo@datadoghq.com>
Date: Thu, 12 Sep 2024 15:15:33 -0400
Subject: [PATCH] Skip xfn->Lambda context propagation if custom Payload is
 used

---
 .../stepfunctions/__tests__/helpers.test.ts   | 26 +++++++++++++++++++
 src/commands/stepfunctions/helpers.ts         |  3 ++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
index 1d0b05024..89bb450a2 100644
--- a/src/commands/stepfunctions/__tests__/helpers.test.ts
+++ b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -68,6 +68,32 @@ describe('stepfunctions command helpers tests', () => {
       expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeTruthy()
     })
 
+    test('custom payload field not using JsonPath expression', () => {
+      const step: StepType = {
+        Type: 'Task',
+        Resource: 'arn:aws:states:::lambda:invoke',
+        Parameters: {
+          FunctionName: 'arn:aws:lambda:sa-east-1:425362991234:function:unit-test-lambda-function',
+          Payload: '{"action": "service/delete_customer"}',
+        },
+        End: true,
+      }
+      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+    })
+
+    test('custom payload field using JsonPath expression', () => {
+      const step: StepType = {
+        Type: 'Task',
+        Resource: 'arn:aws:states:::lambda:invoke',
+        Parameters: {
+          FunctionName: 'arn:aws:lambda:sa-east-1:425362991234:function:unit-test-lambda-function',
+          'Payload.$': '{"customer.$": "$.customer"}',
+        },
+        End: true,
+      }
+      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+    })
+
     test('none-lambda step should not be updated', () => {
       const step: StepType = {
         Type: 'Task',
diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
index 5da8aba5c..01ec84c52 100644
--- a/src/commands/stepfunctions/helpers.ts
+++ b/src/commands/stepfunctions/helpers.ts
@@ -138,7 +138,7 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
     }
 
     // payload field not set
-    if (!step.Parameters.hasOwnProperty('Payload.$')) {
+    if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
       return true
     }
 
@@ -228,6 +228,7 @@ export type StepType = {
 
 export type ParametersType = {
   'Payload.$'?: string
+  Payload?: string
   FunctionName?: string
   StateMachineArn?: string
   TableName?: string