feat: add multivalue headers changes from datadog-lambda-go

contrib/aws/datadog-lambda-go/internal/testdata/non-proxy-with-multivalue-headers.json

@@ -0,0 +1,15 @@
+{
+  "my-custom-event": {
+    "hello": 100
+  },
+  "fake-id": "12345678910",
+  "headers": {
+    "x-datadog-trace-id": "1231452342",
+    "x-datadog-parent-id": "45678910",
+    "x-datadog-sampling-priority": "2"
+  },
+  "multivalueheaders": {
+    "x-datadog-origin": ["origin1", "origin2"],
+    "x-datadog-trace-id": ["duplicate"]
+  }
+}

contrib/aws/datadog-lambda-go/internal/trace/context.go

@@ -24,7 +24,8 @@ import (
 
 type (
 	eventWithHeaders struct {
-		Headers map[string]string `json:"headers"`
+		Headers           map[string]string   `json:"headers"`
+		MultiValueHeaders map[string][]string `json:"multiValueHeaders"`
 	}
 
 	// TraceContext is map of headers containing a Datadog trace context.
@@ -164,8 +165,9 @@ func getTraceContext(ctx context.Context, headers map[string]string) (TraceConte
 	return tc, true
 }
 
-// getHeadersFromEventHeaders extracts the Datadog trace context from an incoming Lambda event payload
-// and creates a dummy X-Ray subsegment containing this information.
+// getHeadersFromEventHeaders extracts the Datadog trace context from an incoming
+// Lambda event payload's headers and multivalueHeaders, with headers taking precedence
+// then creates a dummy X-Ray subsegment containing this information.
 // This is used as the DefaultTraceExtractor.
 func getHeadersFromEventHeaders(ctx context.Context, ev json.RawMessage) map[string]string {
 	eh := eventWithHeaders{}
@@ -178,10 +180,22 @@ func getHeadersFromEventHeaders(ctx context.Context, ev json.RawMessage) map[str
 	}
 
 	lowercaseHeaders := map[string]string{}
+
+	// extract values from event headers into lowercaseheaders
 	for k, v := range eh.Headers {
 		lowercaseHeaders[strings.ToLower(k)] = v
 	}
 
+	// now extract from multivalue headers
+	for k, v := range eh.MultiValueHeaders {
+		if len(v) > 0 {
+			// If this key was not already extracted from event headers, extract first value from multivalue headers
+			if _, ok := lowercaseHeaders[strings.ToLower(k)]; !ok {
+				lowercaseHeaders[strings.ToLower(k)] = v[0]
+			}
+		}
+	}
+
 	return lowercaseHeaders
 }
 

contrib/aws/datadog-lambda-go/internal/trace/context_test.go

@@ -127,6 +127,23 @@ func TestGetDatadogTraceContextForMissingData(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestGetDatadogTraceContextWithMultivalueHeaders(t *testing.T) {
+	// test that multivalue headers are properly extracted from given context
+	// single value headers should take precedence in the case of duplicates
+
+	ctx := mockLambdaXRayTraceContext(context.Background(), mockXRayTraceID, mockXRayEntityID, true)
+	ev := loadRawJSON(t, "../testdata/non-proxy-with-multivalue-headers.json")
+	expected := TraceContext{
+		"x-datadog-trace-id":          "1231452342",
+		"x-datadog-parent-id":         "45678910",
+		"x-datadog-sampling-priority": "2",
+	}
+
+	actual, ok := getTraceContext(ctx, getHeadersFromEventHeaders(ctx, *ev))
+	assert.True(t, ok)
+	assert.Equal(t, expected, actual)
+}
+
 func TestGetDatadogTraceContextFromContextObject(t *testing.T) {
 	testcases := []struct {
 		traceID          string