fix(appsec): send the normalized status code on end-invocation

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,
+      });
+    });
   });
 });

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<string, string> | undefined,
   });
 }

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();
       }

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");