diff --git a/contrib/aws/datadog-lambda-go/internal/testdata/non-proxy-with-multivalue-headers.json b/contrib/aws/datadog-lambda-go/internal/testdata/non-proxy-with-multivalue-headers.json new file mode 100644 index 0000000000..51fb7b8204 --- /dev/null +++ b/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"] + } +} diff --git a/contrib/aws/datadog-lambda-go/internal/trace/context.go b/contrib/aws/datadog-lambda-go/internal/trace/context.go index dc6a067659..88c88ca792 100644 --- a/contrib/aws/datadog-lambda-go/internal/trace/context.go +++ b/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 } diff --git a/contrib/aws/datadog-lambda-go/internal/trace/context_test.go b/contrib/aws/datadog-lambda-go/internal/trace/context_test.go index fb1ff892ee..58d45ae8f4 100644 --- a/contrib/aws/datadog-lambda-go/internal/trace/context_test.go +++ b/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