From 2a2f4758ed2f51ec2e7f90c9710a3a906fd439d9 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Tue, 18 Nov 2025 15:27:35 -0500 Subject: [PATCH 1/4] close any existing active span when extracting the tracecontext --- src/trace/trace-context-service.spec.ts | 51 +++++++++++++++++++++++++ src/trace/trace-context-service.ts | 3 +- src/trace/tracer-wrapper.spec.ts | 31 +++++++++++++++ src/trace/tracer-wrapper.ts | 18 +++++++++ 4 files changed, 102 insertions(+), 1 deletion(-) diff --git a/src/trace/trace-context-service.spec.ts b/src/trace/trace-context-service.spec.ts index 2e9c51bb..ef877d81 100644 --- a/src/trace/trace-context-service.spec.ts +++ b/src/trace/trace-context-service.spec.ts @@ -14,6 +14,7 @@ describe("TraceContextService", () => { mockXRayShouldThrow = false; const traceWrapper = { traceContext: () => spanContextWrapper, + closeScope: jest.fn(), }; traceContextService = new TraceContextService(traceWrapper as any, {} as any); }); @@ -126,4 +127,54 @@ describe("TraceContextService", () => { expect(result?.toTraceId()).toBe("newTraceId"); expect(traceContextService.traceSource).toBe("event"); }); + + it("should not leak dd-trace context from previous invocation when extract is called", async () => { + // Simulate dd-trace having a stale active span from a previous invocation (warm start scenario) + const staleDdTraceContext = { + toTraceId: () => "staleTraceId_999", + toSpanId: () => "staleSpanId_888", + sampleMode: () => 1, + source: TraceSource.DdTrace, + spanContext: {}, + }; + + // Mock tracerWrapper that returns stale context initially, then null after closeScope is called + let traceContextValue: any = staleDdTraceContext; + const mockCloseScopeFn = jest.fn(() => { + // After closeScope is called, traceContext should return null + traceContextValue = null; + }); + + const mockTracerWrapper = { + traceContext: jest.fn(() => traceContextValue), + closeScope: mockCloseScopeFn, + }; + + const service = new TraceContextService(mockTracerWrapper as any, {} as any); + + // Mock the extractor to return a NEW context for the current invocation + const newEventContext = { + toTraceId: () => "newTraceId_123", + toSpanId: () => "newSpanId_456", + sampleMode: () => 2, + source: TraceSource.Event, + spanContext: {}, + }; + const mockExtract = jest.fn().mockResolvedValue(newEventContext); + service["traceExtractor"] = { extract: mockExtract } as any; + + // Call extract for the new invocation + await service.extract({}, {} as any); + + // Verify that closeScope was called to clear the stale context + expect(mockCloseScopeFn).toHaveBeenCalled(); + + // After the fix: currentTraceHeaders should return the NEW context from the event + // not the stale dd-trace context from the previous invocation + const headers = service.currentTraceHeaders; + + expect(headers["x-datadog-trace-id"]).toBe("newTraceId_123"); + expect(headers["x-datadog-parent-id"]).toBe("newSpanId_456"); + expect(headers["x-datadog-sampling-priority"]).toBe("2"); + }); }); diff --git a/src/trace/trace-context-service.ts b/src/trace/trace-context-service.ts index 247e7ff7..83ec8c0b 100644 --- a/src/trace/trace-context-service.ts +++ b/src/trace/trace-context-service.ts @@ -50,8 +50,9 @@ export class TraceContextService { } async extract(event: any, context: Context): Promise { - // Reset trace context from previous invocation to prevent caching + // Reset trace context and close dd-trace scope to prevent stale context from previous invocation due to unfinished spans this.rootTraceContext = null; + this.tracerWrapper.closeScope(); this.rootTraceContext = await this.traceExtractor?.extract(event, context); diff --git a/src/trace/tracer-wrapper.spec.ts b/src/trace/tracer-wrapper.spec.ts index 7281d3ef..01c5f4f3 100644 --- a/src/trace/tracer-wrapper.spec.ts +++ b/src/trace/tracer-wrapper.spec.ts @@ -127,4 +127,35 @@ describe("TracerWrapper", () => { expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(eventType, arn, contextJson, false); }); + + it("should finish active span when closing scope to prevent context leakage", () => { + const mockFinishFn = jest.fn(); + mockSpan = { + context: () => ({ + toSpanId: () => "1234", + toTraceId: () => "45678", + _sampling: { + priority: "2", + }, + }), + finish: mockFinishFn, + }; + + const wrapper = new TracerWrapper(); + wrapper.closeScope(); + + expect(mockFinishFn).toHaveBeenCalled(); + }); + + it("should not error when closing scope with no active span", () => { + mockSpan = null; + const wrapper = new TracerWrapper(); + expect(() => wrapper.closeScope()).not.toThrow(); + }); + + it("should not error when closing scope with tracer unavailable", () => { + mockNoTracer = true; + const wrapper = new TracerWrapper(); + expect(() => wrapper.closeScope()).not.toThrow(); + }); }); diff --git a/src/trace/tracer-wrapper.ts b/src/trace/tracer-wrapper.ts index bdc5ebec..82937a8b 100644 --- a/src/trace/tracer-wrapper.ts +++ b/src/trace/tracer-wrapper.ts @@ -94,6 +94,24 @@ export class TracerWrapper { return new SpanContextWrapper(span.context(), TraceSource.DdTrace); } + public closeScope(): void { + if (!this.isTracerAvailable) { + return; + } + try { + const activeSpan = this.currentSpan; + if (activeSpan && typeof activeSpan.finish === "function") { + logDebug("Finishing stale dd-trace span to prevent context leakage between invocations"); + // Finish any stale span from previous invocation due to unfinished spans + activeSpan.finish(); + } + } catch (err) { + if (err instanceof Object || err instanceof Error) { + logDebug("Failed to close dd-trace scope", err); + } + } + } + public injectSpan(span: SpanContext): any { const dest = {}; this.tracer.inject(span, "text_map", dest); From f5cd92879adf00b5d0f4288af3a4858426724fc1 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 19 Nov 2025 14:28:12 -0500 Subject: [PATCH 2/4] refactoring and fixing --- src/trace/trace-context-service.ts | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/trace/trace-context-service.ts b/src/trace/trace-context-service.ts index 83ec8c0b..97bc7a75 100644 --- a/src/trace/trace-context-service.ts +++ b/src/trace/trace-context-service.ts @@ -55,8 +55,8 @@ export class TraceContextService { this.tracerWrapper.closeScope(); this.rootTraceContext = await this.traceExtractor?.extract(event, context); - - return this.currentTraceContext; + // Return the extracted context, not the current context which may not be related to the event or context + return this.rootTraceContext; } get currentTraceHeaders(): Partial { @@ -71,7 +71,6 @@ export class TraceContextService { } get currentTraceContext(): SpanContextWrapper | null { - if (this.rootTraceContext === null) return null; const traceContext = this.rootTraceContext; const currentDatadogContext = this.tracerWrapper.traceContext(); From 2aef2e78679ced2df89713882fe65d5db279e346 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 19 Nov 2025 14:31:54 -0500 Subject: [PATCH 3/4] lint --- src/trace/trace-context-service.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/trace/trace-context-service.ts b/src/trace/trace-context-service.ts index 97bc7a75..d39f62ad 100644 --- a/src/trace/trace-context-service.ts +++ b/src/trace/trace-context-service.ts @@ -71,7 +71,6 @@ export class TraceContextService { } get currentTraceContext(): SpanContextWrapper | null { - const traceContext = this.rootTraceContext; const currentDatadogContext = this.tracerWrapper.traceContext(); if (currentDatadogContext) { From cb3a3d66eff0d5a354cd8d7943f7311e9c04b10f Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 19 Nov 2025 14:48:33 -0500 Subject: [PATCH 4/4] clean up the logic --- src/trace/tracer-wrapper.ts | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/trace/tracer-wrapper.ts b/src/trace/tracer-wrapper.ts index 82937a8b..1ee8b844 100644 --- a/src/trace/tracer-wrapper.ts +++ b/src/trace/tracer-wrapper.ts @@ -1,4 +1,4 @@ -import { logDebug } from "../utils"; +import { logDebug, logWarning } from "../utils"; import { SpanContextWrapper } from "./span-context-wrapper"; import { TraceSource } from "./trace-context-service"; @@ -95,14 +95,10 @@ export class TracerWrapper { } public closeScope(): void { - if (!this.isTracerAvailable) { - return; - } try { const activeSpan = this.currentSpan; if (activeSpan && typeof activeSpan.finish === "function") { - logDebug("Finishing stale dd-trace span to prevent context leakage between invocations"); - // Finish any stale span from previous invocation due to unfinished spans + logDebug("Detected stale span from previous invocation, finishing it to prevent trace context leakage"); activeSpan.finish(); } } catch (err) {