From 6eda6791cf50470472e9a3616360d2c33be7c2a6 Mon Sep 17 00:00:00 2001
From: SATO Taiki <quietsato@gmail.com>
Date: Tue, 5 Nov 2024 20:58:02 +0900
Subject: [PATCH] support extraction from MultiValueHeaders in ALB target group
 requests

---
 pkg/serverless/trace/propagation/extractor.go |  6 +++++-
 .../trace/propagation/extractor_test.go       | 18 ++++++++++++++++++
 pkg/serverless/trigger/extractor.go           | 17 +++++++++++++----
 pkg/serverless/trigger/extractor_test.go      | 19 +++++++++++++++++++
 4 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/pkg/serverless/trace/propagation/extractor.go b/pkg/serverless/trace/propagation/extractor.go
index d6c756ae1693f..7de0df08e312a 100644
--- 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
index 8cdd07c6d9027..c598b16befd18 100644
--- 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
index 1652c176245e3..3cfd56683a7ac 100644
--- 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
index 34cb6f4b3c816..91234c92ed997 100644
--- 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{