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

src/appsec/event-data-extractor.ts

@@ -1,4 +1,5 @@
 import * as eventType from "../utils/event-type-guards";
+import { normalizeHeaders } from "./headers";
 
 export interface ExtractedHTTPData {
   headers: Record<string, string>;
@@ -127,34 +128,6 @@ function extractFromLambdaUrl(event: any): ExtractedHTTPData {
   return result;
 }
 
-function normalizeHeaders(
-  headers?: Record<string, string>,
-  multiValueHeaders?: Record<string, string[]>,
-): Record<string, string> {
-  if (!headers && !multiValueHeaders) return {};
-
-  const result: Record<string, string> = {};
-
-  if (multiValueHeaders) {
-    for (const [key, values] of Object.entries(multiValueHeaders)) {
-      if (values && values.length > 0) {
-        result[key.toLowerCase()] = values.join(", ");
-      }
-    }
-  }
-
-  if (headers) {
-    for (const [key, value] of Object.entries(headers)) {
-      const lowerKey = key.toLowerCase();
-      if (!(lowerKey in result) && value !== undefined) {
-        result[lowerKey] = value;
-      }
-    }
-  }
-
-  return result;
-}
-
 function separateCookies(headers: Record<string, string>): {
   cookies: Record<string, string> | undefined;
   headersNoCookies: Record<string, string>;

src/appsec/headers.ts

@@ -0,0 +1,35 @@
+/**
+ * Normalizes header names to lowercase and collapses multi-value headers into a single
+ * comma-separated value, which is the shape the WAF expects. When a name appears in both maps the
+ * multi-value one wins.
+ *
+ * API Gateway v1 and application load balancer payloads carry `headers` and `multiValueHeaders`,
+ * both on the request event and on the handler result.
+ */
+export function normalizeHeaders(
+  headers?: Record<string, unknown>,
+  multiValueHeaders?: Record<string, unknown[]>,
+): Record<string, string> {
+  if (!headers && !multiValueHeaders) return {};
+
+  const result: Record<string, string> = {};
+
+  if (multiValueHeaders) {
+    for (const [key, values] of Object.entries(multiValueHeaders)) {
+      if (Array.isArray(values) && values.length > 0) {
+        result[key.toLowerCase()] = values.join(", ");
+      }
+    }
+  }
+
+  if (headers) {
+    for (const [key, value] of Object.entries(headers)) {
+      const lowerKey = key.toLowerCase();
+      if (!(lowerKey in result) && value !== undefined && value !== null) {
+        result[lowerKey] = String(value);
+      }
+    }
+  }
+
+  return result;
+}

src/appsec/index.spec.ts

@@ -78,10 +78,10 @@ describe("AppSec orchestrator", () => {
       expect(mockPublish).not.toHaveBeenCalled();
     });
 
-    it("should extract status code and headers from the result and publish them", () => {
+    it("should publish the normalized status code and the response headers", () => {
       const span = { setTag: jest.fn() };
 
-      processAppsecResponse(span, { statusCode: 200, headers: { "content-type": "application/json" } });
+      processAppsecResponse(span, { statusCode: 200, headers: { "content-type": "application/json" } }, "200");
 
       expect(mockPublish).toHaveBeenCalledWith({
         span,
@@ -101,5 +101,109 @@ describe("AppSec orchestrator", () => {
         responseHeaders: undefined,
       });
     });
+
+    it("should ignore the status code carried by 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 lowercase the response header names", () => {
+      const span = { setTag: jest.fn() };
+
+      processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": "test_value" } }, "200");
+
+      expect(mockPublish).toHaveBeenCalledWith({
+        span,
+        statusCode: "200",
+        responseHeaders: { "x-option": "test_value" },
+      });
+    });
+
+    it("should merge multi value response headers", () => {
+      const span = { setTag: jest.fn() };
+
+      processAppsecResponse(
+        span,
+        {
+          statusCode: 200,
+          headers: { "Content-Type": "application/json" },
+          multiValueHeaders: { "X-Option": ["a", "b"] },
+        },
+        "200",
+      );
+
+      expect(mockPublish).toHaveBeenCalledWith({
+        span,
+        statusCode: "200",
+        responseHeaders: { "content-type": "application/json", "x-option": "a, b" },
+      });
+    });
+
+    it("should ignore multi value response headers that are not arrays", () => {
+      const span = { setTag: jest.fn() };
+
+      processAppsecResponse(span, { statusCode: 200, multiValueHeaders: { "Set-Cookie": "a=b" } }, "200");
+
+      expect(mockPublish).toHaveBeenCalledWith({
+        span,
+        statusCode: "200",
+        responseHeaders: {},
+      });
+    });
+
+    it("should stringify non string response header values", () => {
+      const span = { setTag: jest.fn() };
+
+      processAppsecResponse(span, { statusCode: 200, headers: { "Content-Length": 42, "X-Flag": true } }, "200");
+
+      expect(mockPublish).toHaveBeenCalledWith({
+        span,
+        statusCode: "200",
+        responseHeaders: { "content-length": "42", "x-flag": "true" },
+      });
+    });
+
+    it("should skip response headers with a null value", () => {
+      const span = { setTag: jest.fn() };
+
+      processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": null } }, "200");
+
+      expect(mockPublish).toHaveBeenCalledWith({
+        span,
+        statusCode: "200",
+        responseHeaders: {},
+      });
+    });
+
+    it("should publish no status code when none is normalized, even if the result carries one", () => {
+      const span = { setTag: jest.fn() };
+
+      processAppsecResponse(span, { statusCode: 204 }, undefined);
+
+      expect(mockPublish).toHaveBeenCalledWith({
+        span,
+        statusCode: undefined,
+        responseHeaders: undefined,
+      });
+    });
   });
 });

src/appsec/index.ts

@@ -2,6 +2,7 @@
 const dc = require("dc-polyfill");
 
 import { extractHTTPDataFromEvent } from "./event-data-extractor";
+import { normalizeHeaders } from "./headers";
 
 const startInvocationChannel = dc.channel("datadog:lambda:start-invocation");
 const endInvocationChannel = dc.channel("datadog:lambda:end-invocation");
@@ -30,12 +31,29 @@ 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.
+ */
+export function processAppsecResponse(span: any, result: any, statusCode?: string): void {
   if (!span || !endInvocationChannel.hasSubscribers) return;
 
   endInvocationChannel.publish({
     span,
-    statusCode: result?.statusCode?.toString(),
-    responseHeaders: result?.headers as Record<string, string> | undefined,
+    statusCode,
+    responseHeaders: normalizeResponseHeaders(result),
   });
 }
+
+/**
+ * Response headers reach the tracer in the same shape as the request ones
+ */
+function normalizeResponseHeaders(result: any): Record<string, string> | undefined {
+  const headers = result?.headers as Record<string, unknown> | undefined;
+  const multiValueHeaders = result?.multiValueHeaders as Record<string, unknown[]> | undefined;
+
+  if (!headers && !multiValueHeaders) return undefined;
+
+  return normalizeHeaders(headers, multiValueHeaders);
+}

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