From 512037474b273f59a0c1878aef4136399d010d8a Mon Sep 17 00:00:00 2001
From: Yiming Luo <yiming.luo@datadoghq.com>
Date: Wed, 18 Sep 2024 13:34:57 -0400
Subject: [PATCH] Handle xfn->Lambda context injection when Payload is used
 without $

---
 .../stepfunctions/__tests__/helpers.test.ts   | 50 +++++++++++++++++
 src/commands/stepfunctions/helpers.ts         | 53 +++++++++++++++++--
 2 files changed, 98 insertions(+), 5 deletions(-)

diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
index 800218db7..2342cf9b4 100644
--- a/src/commands/stepfunctions/__tests__/helpers.test.ts
+++ b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -13,6 +13,7 @@ import {
   StepType,
   injectContextForStepFunctions,
   shouldUpdateStepForStepFunctionContextInjection,
+  PayloadObject,
 } from '../helpers'
 
 import {describeStateMachineFixture} from './fixtures/aws-resources'
@@ -69,6 +70,55 @@ describe('stepfunctions command helpers tests', () => {
       expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
     })
 
+    test('already injected Execution and State into Payload', () => {
+      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: {
+            'Execution.$': '$$.Execution',
+            'State.$': '$$.State',
+            'StateMachine.$': '$$.StateMachine',
+          },
+        },
+        End: true,
+      }
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
+    })
+
+    test('custom State field in Payload', () => {
+      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: {
+            State: {Name: 'Lambda Invoke'},
+          },
+        },
+        End: true,
+      }
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
+    })
+
+    test('no Execution, State, or StateMachine field in Payload', () => {
+      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: {},
+        },
+        End: true,
+      }
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeTruthy()
+      const payload = step.Parameters?.['Payload'] as PayloadObject
+      expect(payload['Execution.$']).toEqual('$$.Execution')
+      expect(payload['State.$']).toEqual('$$.State')
+      expect(payload['StateMachine.$']).toEqual('$$.StateMachine')
+    })
+
     test('default payload field of $', () => {
       const step: StepType = {
         Type: 'Task',
diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
index ceda2a390..28ea10be0 100644
--- a/src/commands/stepfunctions/helpers.ts
+++ b/src/commands/stepfunctions/helpers.ts
@@ -193,9 +193,18 @@ export type StepType = {
   End?: boolean
 }
 
+export type PayloadObject = {
+  'Execution.$'?: any
+  Execution?: any
+  'State.$'?: any
+  State?: any
+  'StateMachine.$'?: any
+  StateMachine?: any
+}
+
 export type ParametersType = {
   'Payload.$'?: string
-  Payload?: string
+  Payload?: string | PayloadObject
   FunctionName?: string
   StateMachineArn?: string
   TableName?: string
@@ -237,13 +246,47 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
     return true
   }
 
-  // payload is not a JSON object
-  if (step.Parameters.hasOwnProperty('Payload') && typeof step.Parameters['Payload'] !== 'object') {
-    context.stdout.write(`[Warn] Step ${stepName}'s Payload field is not a JSON object. Step Functions Context Object \
+  if (step.Parameters.hasOwnProperty('Payload')) {
+    if (typeof step.Parameters['Payload'] !== 'object') {
+      // payload is not a JSON object
+      context.stdout
+        .write(`[Warn] Step ${stepName}'s Payload field is not a JSON object. Step Functions Context Object \
 injection skipped. Your Step Functions trace will not be merged with downstream Lambda traces. To manually \
 merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
 
-    return false
+      return false
+    } else {
+      const payload = step.Parameters.Payload
+      if (
+        payload['Execution.$'] === '$$.Execution' &&
+        payload['State.$'] === '$$.State' &&
+        payload['StateMachine.$'] === '$$.StateMachine'
+      ) {
+        context.stdout.write(`Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
+
+        return false
+      } else if (
+        payload.hasOwnProperty('Execution.$') ||
+        payload.hasOwnProperty('Execution') ||
+        payload.hasOwnProperty('State.$') ||
+        payload.hasOwnProperty('State') ||
+        payload.hasOwnProperty('StateMachine.$') ||
+        payload.hasOwnProperty('StateMachine')
+      ) {
+        context.stdout
+          .write(`[Warn] Step ${stepName} may be using custom Execution, State or StateMachine field. Step Functions Context Object \
+injection skipped. Your Step Functions trace will not be merged with downstream Lambda traces. To manually \
+merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
+
+        return false
+      } else {
+        payload['Execution.$'] = '$$.Execution'
+        payload['State.$'] = '$$.State'
+        payload['StateMachine.$'] = '$$.StateMachine'
+
+        return true
+      }
+    }
   }
 
   // default payload