Add implementation and test for SFN -> SNS -> Lambda

Author: ryanstrat

src/trace/context/extractors/sns.spec.ts

@@ -1,6 +1,7 @@
 import { SNSEvent } from "aws-lambda";
 import { TracerWrapper } from "../../tracer-wrapper";
 import { SNSEventTraceExtractor } from "./sns";
+import { StepFunctionContextService } from "../../step-function-service";
 
 let mockSpanContext: any = null;
 
@@ -256,5 +257,51 @@ describe("SNSEventTraceExtractor", () => {
       expect(traceContext?.sampleMode()).toBe("1");
       expect(traceContext?.source).toBe("event");
     });
+
+    it("extracts trace context from Step Function SNS event", () => {
+      // Reset StepFunctionContextService instance
+      StepFunctionContextService["_instance"] = undefined as any;
+
+      const tracerWrapper = new TracerWrapper();
+
+      const payload: SNSEvent = {
+        Records: [
+          {
+            EventSource: "aws:sns",
+            EventVersion: "1.0",
+            EventSubscriptionArn: "arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-sns-demo-dev-process-event-topic:f18241f8-a4f7-4586-80db-97bd1939a557",
+            Sns: {
+              Type: "Notification",
+              MessageId: "46d2665c-7ee2-50ba-a4cd-06acf35f5d5f",
+              TopicArn: "arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-sns-demo-dev-process-event-topic",
+              Message: "{\"source\":\"demo.stepfunction\",\"detailType\":\"ProcessEvent\",\"message\":\"Test event from Step Functions\",\"customData\":{\"userId\":\"12345\",\"action\":\"test\"}}",
+              Timestamp: "2025-07-15T17:10:21.503Z",
+              SignatureVersion: "1",
+              Signature: "fHeJta0GWCs/lHhI6wesXiT+66i1SZ+XH58pyd8mKHKD8bepXsnWvfQdDsOkO2AVv2CqPBF58sAWQae6yob2aMawe/vo8eeahJCaguK8a/3HLj7kP+nXGjgSGvzQm4CdYEyAUco453/mfE/BSf0SkdctxW0rjMs27T2l964Lt2Y/vJeiXVibs/AqEIu3ImekbM8+EIfNMOLBdRBVE47650vawazMkcpPtg5o/8LCA/jNUNj9VCTJrvzep8/vVJEcuHbZ3pcmajA9UJmP3000G0+to0cXwZ5YaakOxQTv81I+cfC99yQJoogLklbgiu+4bqEeNWbwW9KdQz1U+79NgA==",
+              SigningCertUrl: "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem",
+              Subject: "Event from Step Functions",
+              UnsubscribeUrl: "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-sns-demo-dev-process-event-topic:f18241f8-a4f7-4586-80db-97bd1939a557",
+              MessageAttributes: {
+                _datadog: {
+                  Type: "String",
+                  Value: "{\"Execution\":{\"Id\":\"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-demo-dev-state-machine:79049e80-5cc6-49da-9dc0-f19ba2921772\",\"StartTime\":\"2025-07-15T17:10:21.328Z\",\"Name\":\"79049e80-5cc6-49da-9dc0-f19ba2921772\",\"RoleArn\":\"arn:aws:iam::123456123456:role/rstrat-sfn-sns-demo-dev-StepFunctionsExecutionRole-LrsdDm6wMmBh\",\"RedriveCount\":0},\"StateMachine\":{\"Id\":\"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sns-demo-dev-state-machine\",\"Name\":\"rstrat-sfn-sns-demo-dev-state-machine\"},\"State\":{\"Name\":\"PublishToSNS\",\"EnteredTime\":\"2025-07-15T17:10:21.354Z\",\"RetryCount\":0},\"RootExecutionId\":\"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-demo-dev-state-machine:79049e80-5cc6-49da-9dc0-f19ba2921772\",\"serverless-version\":\"v1\"}"
+                }
+              }
+            }
+          }
+        ]
+      };
+
+      const extractor = new SNSEventTraceExtractor(tracerWrapper);
+
+      const traceContext = extractor.extract(payload);
+      expect(traceContext).not.toBeNull();
+
+      // The StepFunctionContextService generates deterministic trace IDs
+      expect(traceContext?.toTraceId()).toBe("3995810302240690842");
+      expect(traceContext?.toSpanId()).toBe("8347071195300897803");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
   });
 });

src/trace/context/extractors/sns.ts

@@ -4,6 +4,7 @@ import { logDebug } from "../../../utils";
 import { EventTraceExtractor } from "../extractor";
 import { SpanContextWrapper } from "../../span-context-wrapper";
 import { XrayService, AMZN_TRACE_ID_ENV_VAR } from "../../xray-service";
+import { StepFunctionContextService } from "../../step-function-service";
 
 export class SNSEventTraceExtractor implements EventTraceExtractor {
   constructor(private tracerWrapper: TracerWrapper) {}
@@ -22,13 +23,26 @@ export class SNSEventTraceExtractor implements EventTraceExtractor {
           headers = JSON.parse(decodedValue);
         }
 
+        // First try to extract as regular trace headers
         const traceContext = this.tracerWrapper.extract(headers);
         if (traceContext) {
           logDebug("Extracted trace context from SNS event");
           return traceContext;
-        } else {
-          logDebug("Failed to extract trace context from SNS event");
         }
+
+        // If that fails, check if this is a Step Function context
+        const stepFunctionInstance = StepFunctionContextService.instance(headers);
+        const stepFunctionContext = stepFunctionInstance.context;
+
+        if (stepFunctionContext !== undefined) {
+          const spanContext = stepFunctionInstance.spanContext;
+          if (spanContext !== null) {
+            logDebug("Extracted Step Function trace context from SNS event", { spanContext, event });
+            return spanContext;
+          }
+        }
+
+        logDebug("Failed to extract trace context from SNS event");
       }
       // Then try to extract trace context from _X_AMZN_TRACE_ID header (Upstream Java apps can
       // pass down Datadog trace id (parent id wrong) in the env in SNS case)