diff --git a/src/appsec/index.spec.ts b/src/appsec/index.spec.ts index 51b72316..d44795c7 100644 --- a/src/appsec/index.spec.ts +++ b/src/appsec/index.spec.ts @@ -101,5 +101,41 @@ describe("AppSec orchestrator", () => { responseHeaders: undefined, }); }); + + it("should prefer the normalized status code over the raw one from the result", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { statusCode: 200 }, "502"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "502", + responseHeaders: undefined, + }); + }); + + it("should publish the normalized status code when the result carries none", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { headers: { "content-type": "application/json" } }, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, + }); + }); + + it("should fall back to the raw status code when no normalized one is given", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { statusCode: 204 }, undefined); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "204", + responseHeaders: undefined, + }); + }); }); }); diff --git a/src/appsec/index.ts b/src/appsec/index.ts index ee3d8687..2aa4bbd1 100644 --- a/src/appsec/index.ts +++ b/src/appsec/index.ts @@ -30,12 +30,18 @@ export function processAppsecRequest(event: any, span: any): void { }); } -export function processAppsecResponse(span: any, result: any): void { +/** + * @param span + * @param result + * @param statusCode Status code already normalized by the trigger layer. Falls back to the raw + * `result.statusCode` when the caller has none, which happens for non-HTTP triggers. + */ +export function processAppsecResponse(span: any, result: any, statusCode?: string): void { if (!span || !endInvocationChannel.hasSubscribers) return; endInvocationChannel.publish({ span, - statusCode: result?.statusCode?.toString(), + statusCode: statusCode ?? result?.statusCode?.toString(), responseHeaders: result?.headers as Record | undefined, }); } diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index 5e1399dd..9cede1f8 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -1,3 +1,4 @@ +import { readFileSync } from "fs"; import { TraceListener } from "./listener"; import { ddtraceVersion, parentSpanFinishTimeHeader } from "./constants"; import { datadogLambdaVersion } from "../constants"; @@ -698,7 +699,102 @@ describe("TraceListener", () => { listener.onEndingInvocation(event, result, false); expect(mockProcessAppsecResponse).toHaveBeenCalledTimes(1); - expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result); + // Non-HTTP trigger: there is no normalized status code to hand over. + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, undefined); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("passes the normalized status code instead of the raw one for API Gateway v2", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/api-gateway-v2.json", "utf8")); + // API Gateway v2 defaults to 200 when the handler omits the status code, so the raw + // result value (undefined) is not what AppSec should see. + const result = { body: "ok" }; + await listener.onStartInvocation(event, context as any); + listener.onEndingInvocation(event, result, false); + + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, "200"); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("passes the normalized 502 when a buffered function returned no result", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/application-load-balancer.json", "utf8")); + await listener.onStartInvocation(event, context as any); + listener.onEndingInvocation(event, undefined, false); + + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "502"); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("passes the normalized 200 when a streaming function returned no result", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/api-gateway-v2.json", "utf8")); + await listener.onStartInvocation(event, context as any); + listener.onEndingInvocation(event, undefined, true); + + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "200"); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("tags http.status_code on the span before calling processAppsecResponse", async () => { + const callOrder: string[] = []; + mockProcessAppsecResponse.mockImplementation(() => callOrder.push("appsec")); + + const mockSetTag = jest.fn((key: string) => { + if (key === "http.status_code") callOrder.push("tag"); + }); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/api-gateway-v2.json", "utf8")); + await listener.onStartInvocation(event, context as any); + listener.onEndingInvocation(event, { statusCode: 201 }, false); + + expect(callOrder).toEqual(["tag", "appsec"]); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("still calls processAppsecResponse on a 5xx response that short-circuits onEndingInvocation", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/api-gateway-v2.json", "utf8")); + await listener.onStartInvocation(event, context as any); + const responseIs5xxError = listener.onEndingInvocation(event, { statusCode: 500 }, false); + + expect(responseIs5xxError).toBe(true); + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, { statusCode: 500 }, "500"); } finally { currentSpanSpy.mockRestore(); } diff --git a/src/trace/listener.ts b/src/trace/listener.ts index 1483a6ad..2b2010aa 100644 --- a/src/trace/listener.ts +++ b/src/trace/listener.ts @@ -231,25 +231,29 @@ export class TraceListener { // Always clear the tree to prevent memory leaks, even if we skip span creation clearTraceTree(); } - if (this.config.appsecEnabled) { - processAppsecResponse(this.tracerWrapper.currentSpan, result); - } + // The status code has to be resolved and tagged before AppSec runs: the WAF needs the + // normalized value (raw result.statusCode is wrong for ALB, API Gateway v2 and response + // streaming), and API Security samples off the span, so http.status_code must already be + // there when the sampling decision is taken. + let statusCode: string | undefined; if (this.triggerTags) { - const statusCode = extractHTTPStatusCodeTag(this.triggerTags, result, isResponseStreamFunction); + statusCode = extractHTTPStatusCodeTag(this.triggerTags, result, isResponseStreamFunction); // Store the status tag in the listener to send to Xray on invocation completion this.triggerTags["http.status_code"] = statusCode!; if (this.tracerWrapper.currentSpan) { this.tracerWrapper.currentSpan.setTag("http.status_code", statusCode); } - if (this.inferredSpan) { - this.inferredSpan.setTag("http.status_code", statusCode); - - if (statusCode?.length === 3 && statusCode?.startsWith("5")) { - this.wrappedCurrentSpan.setTag("error", 1); - return true; - } - } + this.inferredSpan?.setTag("http.status_code", statusCode); + } + if (this.config.appsecEnabled) { + processAppsecResponse(this.tracerWrapper.currentSpan, result, statusCode); + } + // Kept behind AppSec so 5xx responses still reach the WAF, and still nested on inferredSpan + // so the early return only happens when there is an inferred span, as before. + if (this.inferredSpan && statusCode?.length === 3 && statusCode?.startsWith("5")) { + this.wrappedCurrentSpan.setTag("error", 1); + return true; } if (this.durableFunctionContext) { logDebug("Applying durable function context to the aws.lambda span");