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..d39f62ad 100644 --- a/src/trace/trace-context-service.ts +++ b/src/trace/trace-context-service.ts @@ -50,12 +50,13 @@ 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); - - 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 { @@ -70,8 +71,6 @@ export class TraceContextService { } get currentTraceContext(): SpanContextWrapper | null { - if (this.rootTraceContext === null) return null; - const traceContext = this.rootTraceContext; const currentDatadogContext = this.tracerWrapper.traceContext(); if (currentDatadogContext) { 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..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"; @@ -94,6 +94,20 @@ export class TracerWrapper { return new SpanContextWrapper(span.context(), TraceSource.DdTrace); } + public closeScope(): void { + try { + const activeSpan = this.currentSpan; + if (activeSpan && typeof activeSpan.finish === "function") { + logDebug("Detected stale span from previous invocation, finishing it to prevent trace context leakage"); + 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);