Extract trace context from ALB multiValueHeaders (#647)

* extract first value from `event.multiValueHeaders` when `event.headers` is missing

* unit test

* lint

* test flattening multiValueHeaders with real ALB event

Author: nhulston

event_samples/application-load-balancer-multivalue-headers.json

@@ -0,0 +1,65 @@
+{
+    "requestContext": {
+        "elb": {
+            "targetGroupArn": "arn:aws:elasticloadbalancing:us-east-1:1234567890:targetgroup/nhulston-alb-test/dcabb42f66a496e0"
+        }
+    },
+    "httpMethod": "GET",
+    "path": "/",
+    "multiValueQueryStringParameters": {},
+    "multiValueHeaders": {
+        "accept": [
+            "*/*"
+        ],
+        "accept-encoding": [
+            "gzip, deflate"
+        ],
+        "accept-language": [
+            "*"
+        ],
+        "connection": [
+            "keep-alive"
+        ],
+        "host": [
+            "nhulston-test-0987654321.us-east-1.elb.amazonaws.com"
+        ],
+        "sec-fetch-mode": [
+            "cors"
+        ],
+        "traceparent": [
+            "00-68126c4300000000125a7f065cf9a530-1c6dcc8ab8a6e99d-01"
+        ],
+        "tracestate": [
+            "dd=t.dm:-0;t.tid:68126c4300000000;s:1;p:1c6dcc8ab8a6e99d"
+        ],
+        "user-agent": [
+            "node"
+        ],
+        "x-amzn-trace-id": [
+            "Root=1-68126c45-01b175997ab51c4c47a2d643"
+        ],
+        "x-datadog-parent-id": [
+            "1234567890"
+        ],
+        "x-datadog-sampling-priority": [
+            "1"
+        ],
+        "x-datadog-tags": [
+            "_dd.p.tid=68126c4300000000,_dd.p.dm=-0"
+        ],
+        "x-datadog-trace-id": [
+            "0987654321"
+        ],
+        "x-forwarded-for": [
+            "18.204.55.6"
+        ],
+        "x-forwarded-port": [
+            "80"
+        ],
+        "x-forwarded-proto": [
+            "http"
+        ]
+    },
+    "body": "",
+    "isBase64Encoded": false
+}

src/trace/context/extractor.ts

@@ -75,7 +75,8 @@ export class TraceContextExtractor {
   private getTraceEventExtractor(event: any): EventTraceExtractor | undefined {
     if (!event || typeof event !== "object") return;
 
-    if (event.headers !== null && typeof event.headers === "object") {
+    const headers = event.headers ?? event.multiValueHeaders;
+    if (headers !== null && typeof headers === "object") {
       return new HTTPEventTraceExtractor(this.tracerWrapper, this.config.decodeAuthorizerContext);
     }
 

src/trace/context/extractors/http.spec.ts

@@ -1,5 +1,6 @@
 import { TracerWrapper } from "../../tracer-wrapper";
 import { HTTPEventSubType, HTTPEventTraceExtractor } from "./http";
+const albMultivalueHeadersEvent = require("../../../../event_samples/application-load-balancer-multivalue-headers.json");
 
 let mockSpanContext: any = null;
 
@@ -97,6 +98,66 @@ describe("HTTPEventTraceExtractor", () => {
       expect(traceContext?.source).toBe("event");
     });
 
+    it("extracts trace context from payload with multiValueHeaders", () => {
+      mockSpanContext = {
+        toTraceId: () => "123",
+        toSpanId: () => "456",
+        _sampling: { priority: "1" },
+      };
+      const tracerWrapper = new TracerWrapper();
+      const payload = {
+        multiValueHeaders: {
+          "X-Datadog-Trace-Id": ["123", "789"],
+          "X-Datadog-Parent-Id": ["456"],
+          "X-Datadog-Sampling-Priority": ["1"],
+        },
+      };
+      const extractor = new HTTPEventTraceExtractor(tracerWrapper);
+      const traceContext = extractor.extract(payload);
+
+      expect(traceContext).not.toBeNull();
+      expect(spyTracerWrapper).toHaveBeenCalledWith({
+        "x-datadog-trace-id": "123",
+        "x-datadog-parent-id": "456",
+        "x-datadog-sampling-priority": "1",
+      });
+
+      expect(traceContext?.toTraceId()).toBe("123");
+      expect(traceContext?.toSpanId()).toBe("456");
+      expect(traceContext?.sampleMode()).toBe("1");
+    });
+
+    it("flattens a real ALB multiValueHeaders payload into a lowercase, single-value map", () => {
+      const tracerWrapper = new TracerWrapper();
+      const extractor = new HTTPEventTraceExtractor(tracerWrapper);
+
+      spyTracerWrapper.mockClear();
+      extractor.extract(albMultivalueHeadersEvent);
+      expect(spyTracerWrapper).toHaveBeenCalled();
+
+      const captured = spyTracerWrapper.mock.calls[0][0] as Record<string, string>;
+
+      expect(captured).toEqual({
+        accept: "*/*",
+        "accept-encoding": "gzip, deflate",
+        "accept-language": "*",
+        connection: "keep-alive",
+        host: "nhulston-test-0987654321.us-east-1.elb.amazonaws.com",
+        "sec-fetch-mode": "cors",
+        "user-agent": "node",
+        traceparent: "00-68126c4300000000125a7f065cf9a530-1c6dcc8ab8a6e99d-01",
+        tracestate: "dd=t.dm:-0;t.tid:68126c4300000000;s:1;p:1c6dcc8ab8a6e99d",
+        "x-amzn-trace-id": "Root=1-68126c45-01b175997ab51c4c47a2d643",
+        "x-datadog-tags": "_dd.p.tid=68126c4300000000,_dd.p.dm=-0",
+        "x-datadog-sampling-priority": "1",
+        "x-datadog-trace-id": "0987654321",
+        "x-datadog-parent-id": "1234567890",
+        "x-forwarded-for": "18.204.55.6",
+        "x-forwarded-port": "80",
+        "x-forwarded-proto": "http",
+      });
+    });
+
     it("extracts trace context from payload with authorizer", () => {
       mockSpanContext = {
         toTraceId: () => "2389589954026090296",

src/trace/context/extractors/http.ts

@@ -41,11 +41,17 @@ export class HTTPEventTraceExtractor implements EventTraceExtractor {
       }
     }
 
-    const headers = event.headers;
+    const headers = event.headers ?? event.multiValueHeaders;
     const lowerCaseHeaders: { [key: string]: string } = {};
 
-    for (const key of Object.keys(headers)) {
-      lowerCaseHeaders[key.toLowerCase()] = headers[key];
+    for (const [key, val] of Object.entries(headers)) {
+      if (Array.isArray(val)) {
+        // MultiValueHeaders: take the first value
+        lowerCaseHeaders[key.toLowerCase()] = val[0] ?? "";
+      } else if (typeof val === "string") {
+        // Single‐value header
+        lowerCaseHeaders[key.toLowerCase()] = val;
+      }
     }
 
     const traceContext = this.tracerWrapper.extract(lowerCaseHeaders);