diff --git a/internal/testdata/non-proxy-with-multivalue-headers.json b/internal/testdata/non-proxy-with-multivalue-headers.json new file mode 100644 index 00000000..51fb7b82 --- /dev/null +++ b/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/internal/trace/context.go b/internal/trace/context.go index ac8320ea..0bbc9e62 100644 --- a/internal/trace/context.go +++ b/internal/trace/context.go @@ -27,7 +27,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. @@ -167,8 +168,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{} @@ -181,10 +183,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/internal/trace/context_test.go b/internal/trace/context_test.go index 7c988397..9b979be1 100644 --- a/internal/trace/context_test.go +++ b/internal/trace/context_test.go @@ -130,6 +130,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