Skip to content

Commit a146e03

Browse files
ryanstratryanstrat
committed
Add support for Eventbridge -> SNS -> Lambda pattern for both header and SFN context
1 parent 952887b commit a146e03

File tree

6 files changed

+301
-0
lines changed

6 files changed

+301
-0
lines changed

‎src/trace/context/extractor.spec.ts

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { TraceSource } from "../trace-context-service";
1111
import {
1212
AppSyncEventTraceExtractor,
1313
EventBridgeEventTraceExtractor,
14+
EventBridgeSNSEventTraceExtractor,
1415
EventBridgeSQSEventTraceExtractor,
1516
HTTPEventTraceExtractor,
1617
KinesisEventTraceExtractor,
@@ -658,6 +659,58 @@ describe("TraceContextExtractor", () => {
658659
expect(traceContext?.source).toBe("event");
659660
});
660661

662+
// EventBridge message delivered to SNS event
663+
it("extracts trace context from EventBridge to SNS event", async () => {
664+
mockSpanContext = {
665+
toTraceId: () => "1234567890123456789",
666+
toSpanId: () => "9876543210987654321",
667+
_sampling: {
668+
priority: "1",
669+
},
670+
};
671+
672+
const event: SNSEvent = {
673+
Records: [
674+
{
675+
EventSource: "aws:sns",
676+
EventVersion: "1.0",
677+
EventSubscriptionArn: "arn:aws:sns:us-east-1:123456123456:my-topic:12345678-1234-1234-1234-123456789012",
678+
Sns: {
679+
Type: "Notification",
680+
MessageId: "12345678-1234-1234-1234-123456789012",
681+
TopicArn: "arn:aws:sns:us-east-1:123456123456:my-topic",
682+
Message: '{"version":"0","id":"12345678-1234-1234-1234-123456789012","detail-type":"my.Detail","source":"my.Source","account":"123456123456","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"1234567890123456789","x-datadog-parent-id":"9876543210987654321","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0"}}}',
683+
Timestamp: "2023-08-03T22:49:03.123Z",
684+
SignatureVersion: "1",
685+
Signature: "EXAMPLE",
686+
SigningCertUrl: "https://sns.us-east-1.amazonaws.com/SimpleNotificationService-123456789012.pem",
687+
Subject: undefined,
688+
UnsubscribeUrl: "https://sns.us-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:us-east-1:123456123456:my-topic:12345678-1234-1234-1234-123456789012",
689+
MessageAttributes: {}
690+
}
691+
}
692+
]
693+
};
694+
695+
const tracerWrapper = new TracerWrapper();
696+
const extractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
697+
698+
const traceContext = await extractor.extract(event, {} as Context);
699+
expect(traceContext).not.toBeNull();
700+
701+
expect(spyTracerWrapper).toHaveBeenCalledWith({
702+
"x-datadog-parent-id": "9876543210987654321",
703+
"x-datadog-sampling-priority": "1",
704+
"x-datadog-tags": "_dd.p.dm=-0",
705+
"x-datadog-trace-id": "1234567890123456789",
706+
});
707+
708+
expect(traceContext?.toTraceId()).toBe("1234567890123456789");
709+
expect(traceContext?.toSpanId()).toBe("9876543210987654321");
710+
expect(traceContext?.sampleMode()).toBe("1");
711+
expect(traceContext?.source).toBe("event");
712+
});
713+
661714
// StepFunction context event
662715
it("extracts trace context from StepFunction event", async () => {
663716
const event = {
@@ -821,6 +874,20 @@ describe("TraceContextExtractor", () => {
821874
],
822875
},
823876
],
877+
[
878+
"EventBridgeSNSTraceExtractor",
879+
"EventBridge to SNS event",
880+
EventBridgeSNSEventTraceExtractor,
881+
{
882+
Records: [
883+
{
884+
Sns: {
885+
Message: '{"detail-type":"some-detail-type"}',
886+
},
887+
},
888+
],
889+
},
890+
],
824891
[
825892
"AppSyncEventTraceExtractor",
826893
"AppSync event",

‎src/trace/context/extractor.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
AppSyncEventTraceExtractor,
77
CustomTraceExtractor,
88
EventBridgeEventTraceExtractor,
9+
EventBridgeSNSEventTraceExtractor,
910
EventBridgeSQSEventTraceExtractor,
1011
HTTPEventTraceExtractor,
1112
KinesisEventTraceExtractor,
@@ -80,6 +81,7 @@ export class TraceContextExtractor {
8081
return new HTTPEventTraceExtractor(this.tracerWrapper, this.config.decodeAuthorizerContext);
8182
}
8283

84+
if (EventValidator.isEventBridgeSNSEvent(event)) return new EventBridgeSNSEventTraceExtractor(this.tracerWrapper);
8385
if (EventValidator.isSNSEvent(event)) return new SNSEventTraceExtractor(this.tracerWrapper);
8486
if (EventValidator.isSNSSQSEvent(event)) return new SNSSQSEventTraceExtractor(this.tracerWrapper);
8587
if (EventValidator.isEventBridgeSQSEvent(event)) return new EventBridgeSQSEventTraceExtractor(this.tracerWrapper);

‎src/trace/context/extractors/event-bridge-sns.spec.ts

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
import { TracerWrapper } from "../../tracer-wrapper";
2+
import { EventBridgeSNSEventTraceExtractor } from "./event-bridge-sns";
3+
import { StepFunctionContextService } from "../../step-function-service";
4+
import { SNSEvent } from "aws-lambda";
5+
6+
let mockSpanContext: any = null;
7+
8+
// Mocking extract is needed, due to dd-trace being a No-op
9+
// if the detected environment is testing. This is expected, since
10+
// we don't want to test dd-trace extraction, but our components.
11+
const ddTrace = require("dd-trace");
12+
jest.mock("dd-trace", () => {
13+
return {
14+
...ddTrace,
15+
_tracer: { _service: {} },
16+
extract: (_carrier: any, _headers: any) => mockSpanContext,
17+
};
18+
});
19+
const spyTracerWrapper = jest.spyOn(TracerWrapper.prototype, "extract");
20+
21+
describe("EventBridgeSNSEventTraceExtractor", () => {
22+
describe("extract", () => {
23+
beforeEach(() => {
24+
mockSpanContext = null;
25+
});
26+
27+
afterEach(() => {
28+
jest.resetModules();
29+
});
30+
31+
it("extracts trace context with valid payload", () => {
32+
mockSpanContext = {
33+
toTraceId: () => "1234567890123456789",
34+
toSpanId: () => "9876543210987654321",
35+
_sampling: {
36+
priority: "1",
37+
},
38+
};
39+
const tracerWrapper = new TracerWrapper();
40+
41+
const payload = {
42+
Records: [
43+
{
44+
EventSource: "aws:sns",
45+
EventVersion: "1.0",
46+
EventSubscriptionArn: "arn:aws:sns:us-east-1:123456123456:my-topic:12345678-1234-1234-1234-123456789012",
47+
Sns: {
48+
Type: "Notification",
49+
MessageId: "12345678-1234-1234-1234-123456789012",
50+
TopicArn: "arn:aws:sns:us-east-1:123456123456:my-topic",
51+
Message: '{"version":"0","id":"12345678-1234-1234-1234-123456789012","detail-type":"my.Detail","source":"my.Source","account":"123456123456","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"1234567890123456789","x-datadog-parent-id":"9876543210987654321","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0"}}}',
52+
Timestamp: "2023-08-03T22:49:03.123Z",
53+
SignatureVersion: "1",
54+
Signature: "EXAMPLE",
55+
SigningCertUrl: "https://sns.us-east-1.amazonaws.com/SimpleNotificationService-123456789012.pem",
56+
Subject: undefined,
57+
UnsubscribeUrl: "https://sns.us-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:us-east-1:123456123456:my-topic:12345678-1234-1234-1234-123456789012",
58+
MessageAttributes: {}
59+
}
60+
}
61+
]
62+
};
63+
64+
const extractor = new EventBridgeSNSEventTraceExtractor(tracerWrapper);
65+
66+
const traceContext = extractor.extract(payload);
67+
expect(traceContext).not.toBeNull();
68+
69+
expect(spyTracerWrapper).toHaveBeenCalledWith({
70+
"x-datadog-parent-id": "9876543210987654321",
71+
"x-datadog-sampling-priority": "1",
72+
"x-datadog-tags": "_dd.p.dm=-0",
73+
"x-datadog-trace-id": "1234567890123456789",
74+
});
75+
76+
expect(traceContext?.toTraceId()).toBe("1234567890123456789");
77+
expect(traceContext?.toSpanId()).toBe("9876543210987654321");
78+
expect(traceContext?.sampleMode()).toBe("1");
79+
expect(traceContext?.source).toBe("event");
80+
});
81+
82+
it.each([
83+
["Records", {}],
84+
["Records first entry", { Records: [] }],
85+
["Records first entry Sns", { Records: [{}] }],
86+
["Records first entry Sns Message", { Records: [{ Sns: {} }] }],
87+
["valid JSON in Message", { Records: [{ Sns: { Message: "{" } }] }], // JSON.parse should fail
88+
["detail in Message", { Records: [{ Sns: { Message: "{}" } }] }],
89+
["_datadog in detail", { Records: [{ Sns: { Message: '{"detail":{"text":"Hello, world!"}}' } }] }],
90+
])("returns null and skips extracting when payload is missing '%s'", (_, payload) => {
91+
const tracerWrapper = new TracerWrapper();
92+
const extractor = new EventBridgeSNSEventTraceExtractor(tracerWrapper);
93+
94+
const traceContext = extractor.extract(payload as any);
95+
expect(traceContext).toBeNull();
96+
});
97+
98+
it("returns null when extracted span context by tracer is null", () => {
99+
const tracerWrapper = new TracerWrapper();
100+
101+
const payload = {
102+
Records: [
103+
{
104+
EventSource: "aws:sns",
105+
EventVersion: "1.0",
106+
EventSubscriptionArn: "arn:aws:sns:us-east-1:123456123456:my-topic:12345678-1234-1234-1234-123456789012",
107+
Sns: {
108+
Type: "Notification",
109+
MessageId: "12345678-1234-1234-1234-123456789012",
110+
TopicArn: "arn:aws:sns:us-east-1:123456123456:my-topic",
111+
// Message is missing _datadog
112+
Message: '{"version":"0","id":"12345678-1234-1234-1234-123456789012","detail-type":"my.Detail","source":"my.Source","account":"123456123456","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!"}}',
113+
Timestamp: "2023-08-03T22:49:03.123Z",
114+
SignatureVersion: "1",
115+
Signature: "EXAMPLE",
116+
SigningCertUrl: "https://sns.us-east-1.amazonaws.com/SimpleNotificationService-123456789012.pem",
117+
Subject: undefined,
118+
UnsubscribeUrl: "https://sns.us-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:us-east-1:123456123456:my-topic:12345678-1234-1234-1234-123456789012",
119+
MessageAttributes: {}
120+
}
121+
}
122+
]
123+
};
124+
125+
const extractor = new EventBridgeSNSEventTraceExtractor(tracerWrapper);
126+
127+
const traceContext = extractor.extract(payload as SNSEvent);
128+
expect(traceContext).toBeNull();
129+
});
130+
131+
it("extracts trace context from Step Function EventBridge-SNS event", () => {
132+
// Reset StepFunctionContextService instance
133+
StepFunctionContextService["_instance"] = undefined as any;
134+
135+
const tracerWrapper = new TracerWrapper();
136+
137+
const payload = {
138+
"Records": [
139+
{
140+
"EventSource": "aws:sns",
141+
"EventVersion": "1.0",
142+
"EventSubscriptionArn": "arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-evb-sns-demo-dev-process-event-topic:8257bfde-3426-4901-9ace-6fbb180875b1",
143+
"Sns": {
144+
"Type": "Notification",
145+
"MessageId": "5d4b7d39-8b8d-5427-b7d9-1dc9e3d270aa",
146+
"TopicArn": "arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-evb-sns-demo-dev-process-event-topic",
147+
"Message": "{\"version\":\"0\",\"id\":\"c967411c-9066-225e-1d37-527fed26f847\",\"detail-type\":\"ProcessEvent\",\"source\":\"demo.stepfunction\",\"account\":\"123456123456\",\"time\":\"2025-07-15T14:30:55Z\",\"region\":\"sa-east-1\",\"resources\":[\"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-evb-sns-demo-dev-state-machine\",\"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sns-demo-dev-state-machine:test-execution-1752589852\"],\"detail\":{\"message\":\"Event from Step Functions\",\"timestamp\":\"2025-07-15T14:30:55.311Z\",\"executionName\":\"test-execution-1752589852\",\"stateMachineName\":\"rstrat-sfn-evb-sns-demo-dev-state-machine\",\"input\":{\"testData\":\"Hello from SNS integration\"},\"_datadog\":{\"Execution\":{\"Id\":\"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sns-demo-dev-state-machine:test-execution-1752589852\",\"StartTime\":\"2025-07-15T14:30:55.271Z\",\"Name\":\"test-execution-1752589852\",\"RoleArn\":\"arn:aws:iam::123456123456:role/rstrat-sfn-evb-sns-demo-d-StepFunctionsExecutionRol-5uB2zncEHF8x\",\"RedriveCount\":0},\"StateMachine\":{\"Id\":\"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-evb-sns-demo-dev-state-machine\",\"Name\":\"rstrat-sfn-evb-sns-demo-dev-state-machine\"},\"State\":{\"Name\":\"PublishToEventBridge\",\"EnteredTime\":\"2025-07-15T14:30:55.311Z\",\"RetryCount\":0},\"RootExecutionId\":\"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sns-demo-dev-state-machine:test-execution-1752589852\",\"serverless-version\":\"v1\"}}}",
148+
"Timestamp": "2025-07-15T14:30:55.534Z",
149+
"SignatureVersion": "1",
150+
"Signature": "N+8rhLGKH/oj9wWBkvOrKHpH5icaWNxa68I0gmvxutTm+vBRZiT18GnncRR8yYPdrA6nNndb7PoNhzFdofNw3PhWx8GvaYNXQ41W4qBvM5fqg7XwrnS/+YBmmwI0Mq8uRq/+Uen/J0W8tBDxSmVkmZV8LiZceV113U/YQ8cZOorVUrbKQrCEG9c7+dydJLrFfJqT4sQ+yAxePC1ilme/OBpa2c8T9amnvuXvKRBDk2/uigvTLC9POC45p42q5jbweNzr3igXRHL20WDlR5fdqoKpZiELNeiUZKbA+caK73smQhtMUwY1vNmY8QghSSVGdOK6MgmalRsDAMF6GZRLEQ==",
151+
"SigningCertUrl": "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem",
152+
"Subject": undefined,
153+
"UnsubscribeUrl": "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-evb-sns-demo-dev-process-event-topic:8257bfde-3426-4901-9ace-6fbb180875b1",
154+
"MessageAttributes": {}
155+
}
156+
}
157+
]
158+
};
159+
160+
const extractor = new EventBridgeSNSEventTraceExtractor(tracerWrapper);
161+
162+
const traceContext = extractor.extract(payload);
163+
expect(traceContext).not.toBeNull();
164+
165+
// The trace IDs are deterministically generated from the Step Function execution context
166+
expect(traceContext?.toTraceId()).toBe("2458939194637197377");
167+
expect(traceContext?.toSpanId()).toBe("6978708559187765983");
168+
expect(traceContext?.sampleMode()).toBe("1");
169+
expect(traceContext?.source).toBe("event");
170+
});
171+
});
172+
});

‎src/trace/context/extractors/event-bridge-sns.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
import { EventBridgeEvent, SNSEvent } from "aws-lambda";
2+
import { EventTraceExtractor } from "../extractor";
3+
import { logDebug } from "../../../utils";
4+
import { TracerWrapper } from "../../tracer-wrapper";
5+
import { SpanContextWrapper } from "../../span-context-wrapper";
6+
import { StepFunctionContextService } from "../../step-function-service";
7+
8+
export class EventBridgeSNSEventTraceExtractor implements EventTraceExtractor {
9+
constructor(private tracerWrapper: TracerWrapper) {}
10+
11+
extract(event: SNSEvent): SpanContextWrapper | null {
12+
const message = event?.Records?.[0]?.Sns?.Message;
13+
if (message === undefined) return null;
14+
15+
try {
16+
const parsedMessage = JSON.parse(message) as EventBridgeEvent<any, any>;
17+
const headers = parsedMessage?.detail?._datadog;
18+
if (headers === undefined) return null;
19+
20+
// First try to extract as regular trace headers
21+
const traceContext = this.tracerWrapper.extract(headers);
22+
if (traceContext !== null) {
23+
logDebug("Extracted trace context from EventBridge-SNS event", { traceContext, event });
24+
return traceContext;
25+
}
26+
27+
// If that fails, check if this is a Step Function context
28+
// The StepFunctionContextService can handle the Step Function format
29+
const stepFunctionInstance = StepFunctionContextService.instance(headers);
30+
const stepFunctionContext = stepFunctionInstance.context;
31+
32+
if (stepFunctionContext !== undefined) {
33+
const spanContext = stepFunctionInstance.spanContext;
34+
if (spanContext !== null) {
35+
logDebug("Extracted Step Function trace context from EventBridge-SNS event", { spanContext, event });
36+
return spanContext;
37+
}
38+
}
39+
} catch (error) {
40+
if (error instanceof Error) {
41+
logDebug("Unable to extract trace context from EventBridge-SNS event", error);
42+
}
43+
}
44+
45+
return null;
46+
}
47+
}

‎src/trace/context/extractors/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
export { AppSyncEventTraceExtractor } from "./app-sync";
22
export { EventBridgeEventTraceExtractor } from "./event-bridge";
33
export { EventBridgeSQSEventTraceExtractor } from "./event-bridge-sqs";
4+
export { EventBridgeSNSEventTraceExtractor } from "./event-bridge-sns";
45
export { KinesisEventTraceExtractor } from "./kinesis";
56
export { HTTPEventTraceExtractor } from "./http";
67
export { SQSEventTraceExtractor } from "./sqs";

‎src/utils/event-validator.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,18 @@ export class EventValidator {
5959
return false;
6060
}
6161

62+
static isEventBridgeSNSEvent(event: any): event is SNSEvent {
63+
if (Array.isArray(event.Records) && event.Records.length > 0 && event.Records[0].Sns !== undefined) {
64+
try {
65+
const message = JSON.parse(event.Records[0].Sns.Message) as EventBridgeEvent<any, any>;
66+
return message["detail-type"] !== undefined;
67+
} catch (_) {
68+
return false;
69+
}
70+
}
71+
return false;
72+
}
73+
6274
static isAppSyncResolverEvent(event: any): event is AppSyncResolverEvent<any> {
6375
return event.info !== undefined && event.info.selectionSetGraphQL !== undefined;
6476
}

0 commit comments

Comments
 (0)