support extraction from MultiValueHeaders in ALB target group requests

DataDog · Nov 5, 2024 · 6eda679 · 6eda679
1 parent c5769f7
commit 6eda679
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 5 deletions.
diff --git a/pkg/serverless/trace/propagation/extractor.go b/pkg/serverless/trace/propagation/extractor.go
@@ -112,7 +112,11 @@ func (e Extractor) extract(event interface{}) (*TraceContext, error) {
 	case events.APIGatewayCustomAuthorizerRequestTypeRequest:
 		carrier, err = headersCarrier(ev.Headers)
 	case events.ALBTargetGroupRequest:
-		carrier, err = headersCarrier(ev.Headers)
+		if len(ev.Headers) > 0 {
+			carrier, err = headersCarrier(ev.Headers)
+		} else {
+			carrier, err = multiValueHeadersCarrier(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, "dummy"}
+	}
+	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/extractor.go b/pkg/serverless/trigger/extractor.go
@@ -188,11 +188,20 @@ 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"]
+
+	headers := event.Headers
+	if headers == nil && event.MultiValueHeaders != nil {
+		headers = make(map[string]string)
+		for k, v := range event.MultiValueHeaders {
+			headers[k] = v[0]
 		}
-		if ua := event.Headers["User-Agent"]; ua != "" {
+	}
+
+	if headers != nil {
+		if headers["Referer"] != "" {
+			httpTags["http.referer"] = headers["Referer"]
+		}
+		if ua := headers["User-Agent"]; ua != "" {
 			httpTags["http.useragent"] = ua
 		}
 	}

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{