[serverless] Support trace context propagation for ALB target groups …

…with MultiValueHeaders (DataDog#31542) Co-authored-by: quietsato <[email protected]>
kkpan11 · Dec 6, 2024 · 4b5c8b9 · 4b5c8b9
1 parent af201ab
commit 4b5c8b9
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 8 deletions.
diff --git a/pkg/serverless/trace/propagation/carriers.go b/pkg/serverless/trace/propagation/carriers.go
@@ -253,6 +253,16 @@ func headersCarrier(hdrs map[string]string) (tracer.TextMapReader, error) {
 	return tracer.TextMapCarrier(hdrs), nil
 }
 
+// headersOrMultiheadersCarrier returns the tracer.TextMapReader used to extract
+// trace context from a Headers field of form map[string]string or MultiValueHeaders
+// field of form map[string][]string.
+func headersOrMultiheadersCarrier(hdrs map[string]string, multiHdrs map[string][]string) (tracer.TextMapReader, error) {
+	if len(hdrs) > 0 {
+		return headersCarrier(hdrs)
+	}
+	return tracer.HTTPHeadersCarrier(multiHdrs), nil
+}
+
 // extractTraceContextFromStepFunctionContext extracts the execution ARN, state name, and state entered time and uses them to generate Trace ID and Parent ID
 // The logic is based on the trace context conversion in Logs To Traces, dd-trace-py, dd-trace-js, etc.
 func extractTraceContextFromStepFunctionContext(event events.StepFunctionPayload) (*TraceContext, error) {

diff --git a/pkg/serverless/trace/propagation/carriers_test.go b/pkg/serverless/trace/propagation/carriers_test.go
@@ -816,7 +816,7 @@ func TestHeadersCarrier(t *testing.T) {
 	for _, tc := range testcases {
 		t.Run(tc.name, func(t *testing.T) {
 			tm, err := headersCarrier(tc.event)
-			t.Logf("rawPayloadCarrier returned TextMapReader=%#v error=%#v", tm, err)
+			t.Logf("headersCarrier returned TextMapReader=%#v error=%#v", tm, err)
 			assert.Equal(t, tc.expErr != nil, err != nil)
 			if tc.expErr != nil && err != nil {
 				assert.Equal(t, tc.expErr.Error(), err.Error())
@@ -826,6 +826,55 @@ func TestHeadersCarrier(t *testing.T) {
 	}
 }
 
+func TestHeadersOrMultiheadersCarrier(t *testing.T) {
+	testcases := []struct {
+		name      string
+		hdrs      map[string]string
+		multiHdrs map[string][]string
+		expMap    map[string]string
+	}{
+		{
+			name:      "nil-map",
+			hdrs:      headersMapNone,
+			multiHdrs: toMultiValueHeaders(headersMapNone),
+			expMap:    headersMapEmpty,
+		},
+		{
+			name:      "empty-map",
+			hdrs:      headersMapEmpty,
+			multiHdrs: toMultiValueHeaders(headersMapEmpty),
+			expMap:    headersMapEmpty,
+		},
+		{
+			name:      "headers-and-multiheaders",
+			hdrs:      headersMapDD,
+			multiHdrs: toMultiValueHeaders(headersMapW3C),
+			expMap:    headersMapDD,
+		},
+		{
+			name:      "just-headers",
+			hdrs:      headersMapDD,
+			multiHdrs: toMultiValueHeaders(headersMapEmpty),
+			expMap:    headersMapDD,
+		},
+		{
+			name:      "just-multiheaders",
+			hdrs:      headersMapEmpty,
+			multiHdrs: toMultiValueHeaders(headersMapW3C),
+			expMap:    headersMapW3C,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			tm, err := headersOrMultiheadersCarrier(tc.hdrs, tc.multiHdrs)
+			t.Logf("headersOrMultiheadersCarrier returned TextMapReader=%#v error=%#v", tm, err)
+			assert.Nil(t, err)
+			assert.Equal(t, tc.expMap, getMapFromCarrier(tm))
+		})
+	}
+}
+
 func Test_stringToDdSpanId(t *testing.T) {
 	type args struct {
 		execArn          string

diff --git a/pkg/serverless/trace/propagation/extractor.go b/pkg/serverless/trace/propagation/extractor.go
@@ -112,7 +112,7 @@ func (e Extractor) extract(event interface{}) (*TraceContext, error) {
 	case events.APIGatewayCustomAuthorizerRequestTypeRequest:
 		carrier, err = headersCarrier(ev.Headers)
 	case events.ALBTargetGroupRequest:
-		carrier, err = headersCarrier(ev.Headers)
+		carrier, err = headersOrMultiheadersCarrier(ev.Headers, ev.MultiValueHeaders)
 	case events.LambdaFunctionURLRequest:
 		carrier, err = headersCarrier(ev.Headers)
 	case events.StepFunctionPayload:

diff --git a/pkg/serverless/trace/propagation/extractor_test.go b/pkg/serverless/trace/propagation/extractor_test.go
@@ -176,6 +176,14 @@ var (
 	}
 )
 
+func toMultiValueHeaders(headers map[string]string) map[string][]string {
+	mvh := make(map[string][]string)
+	for k, v := range headers {
+		mvh[k] = []string{v}
+	}
+	return mvh
+}
+
 func TestNilPropagator(t *testing.T) {
 	var extractor Extractor
 	tc, err := extractor.Extract([]byte(`{"headers":` + headersAll + `}`))
@@ -533,6 +541,16 @@ func TestExtractorExtract(t *testing.T) {
 			expCtx:   ddTraceContext,
 			expNoErr: true,
 		},
+		{
+			name: "ALBTargetGroupRequestMultiValueHeaders",
+			events: []interface{}{
+				events.ALBTargetGroupRequest{
+					MultiValueHeaders: toMultiValueHeaders(headersMapAll),
+				},
+			},
+			expCtx:   ddTraceContext,
+			expNoErr: true,
+		},
 
 		// events.LambdaFunctionURLRequest:
 		{

diff --git a/pkg/serverless/trigger/events/events.go b/pkg/serverless/trigger/events/events.go
@@ -122,10 +122,11 @@ type APIGatewayCustomAuthorizerRequestTypeRequestContext struct {
 // ALBTargetGroupRequest mirrors events.ALBTargetGroupRequest type, removing
 // unused fields.
 type ALBTargetGroupRequest struct {
-	HTTPMethod     string
-	Path           string
-	Headers        map[string]string
-	RequestContext ALBTargetGroupRequestContext
+	HTTPMethod        string
+	Path              string
+	Headers           map[string]string
+	MultiValueHeaders map[string][]string
+	RequestContext    ALBTargetGroupRequestContext
 }
 
 // ALBTargetGroupRequestContext mirrors events.ALBTargetGroupRequestContext

diff --git a/pkg/serverless/trigger/extractor.go b/pkg/serverless/trigger/extractor.go
@@ -188,14 +188,23 @@ func GetTagsFromALBTargetGroupRequest(event events.ALBTargetGroupRequest) map[st
 	httpTags := make(map[string]string)
 	httpTags["http.url_details.path"] = event.Path
 	httpTags["http.method"] = event.HTTPMethod
+
 	if event.Headers != nil {
-		if event.Headers["Referer"] != "" {
-			httpTags["http.referer"] = event.Headers["Referer"]
+		if r := event.Headers["Referer"]; r != "" {
+			httpTags["http.referer"] = r
 		}
 		if ua := event.Headers["User-Agent"]; ua != "" {
 			httpTags["http.useragent"] = ua
 		}
+	} else if event.MultiValueHeaders != nil {
+		if r := event.MultiValueHeaders["Referer"]; len(r) > 0 && r[0] != "" {
+			httpTags["http.referer"] = r[0]
+		}
+		if ua := event.MultiValueHeaders["User-Agent"]; len(ua) > 0 && ua[0] != "" {
+			httpTags["http.useragent"] = ua[0]
+		}
 	}
+
 	return httpTags
 }
 

diff --git a/pkg/serverless/trigger/extractor_test.go b/pkg/serverless/trigger/extractor_test.go
@@ -353,6 +353,25 @@ func TestGetTagsFromALBTargetGroupRequest(t *testing.T) {
 	}, httpTags)
 }
 
+func TestGetTagsFromALBTargetGroupRequestMultiValueHeaders(t *testing.T) {
+	event := events.ALBTargetGroupRequest{
+		MultiValueHeaders: map[string][]string{
+			"key":     {"val"},
+			"Referer": {"referer"},
+		},
+		Path:       "path",
+		HTTPMethod: "http-method",
+	}
+
+	httpTags := GetTagsFromALBTargetGroupRequest(event)
+
+	assert.Equal(t, map[string]string{
+		"http.url_details.path": "path",
+		"http.method":           "http-method",
+		"http.referer":          "referer",
+	}, httpTags)
+}
+
 func TestGetTagsFromFunctionURLRequest(t *testing.T) {
 	event := events.LambdaFunctionURLRequest{
 		Headers: map[string]string{