From 270eb93324937909481c99c7ef943510e129237d Mon Sep 17 00:00:00 2001
From: SATO Taiki <quietsato@gmail.com>
Date: Tue, 12 Nov 2024 19:31:17 +0900
Subject: [PATCH] use HTTPHeadersCarrier to extract multi-value headers

Signed-off-by: SATO Taiki <quietsato@gmail.com>
---
 pkg/serverless/trace/propagation/carriers.go  | 10 -----
 .../trace/propagation/carriers_test.go        | 45 -------------------
 pkg/serverless/trace/propagation/extractor.go |  2 +-
 .../trace/propagation/extractor_test.go       |  2 +-
 4 files changed, 2 insertions(+), 57 deletions(-)

diff --git a/pkg/serverless/trace/propagation/carriers.go b/pkg/serverless/trace/propagation/carriers.go
index 5d3ce20ee92b7..9cfa1255baabb 100644
--- a/pkg/serverless/trace/propagation/carriers.go
+++ b/pkg/serverless/trace/propagation/carriers.go
@@ -253,16 +253,6 @@ func headersCarrier(hdrs map[string]string) (tracer.TextMapReader, error) {
 	return tracer.TextMapCarrier(hdrs), nil
 }
 
-// multiValueHeadersCarrier returns the tracer.TextMapReader used to extract trace
-// context from a MultiValueHeaders field of form map[string][]string.
-func multiValueHeadersCarrier(hdrs map[string][]string) (tracer.TextMapReader, error) {
-	headers := make(map[string]string)
-	for k, v := range hdrs {
-		headers[k] = v[0]
-	}
-	return headersCarrier(headers)
-}
-
 // 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
index d46b09fdf2d2a..c58b294b74e39 100644
--- a/pkg/serverless/trace/propagation/carriers_test.go
+++ b/pkg/serverless/trace/propagation/carriers_test.go
@@ -826,51 +826,6 @@ func TestHeadersCarrier(t *testing.T) {
 	}
 }
 
-func TestMultiValueHeadersCarrier(t *testing.T) {
-	testcases := []struct {
-		name   string
-		event  map[string]string
-		expMap map[string]string
-		expErr error
-	}{
-		{
-			name:   "nil-map",
-			event:  headersMapNone,
-			expMap: headersMapEmpty,
-			expErr: nil,
-		},
-		{
-			name:   "empty-map",
-			event:  headersMapEmpty,
-			expMap: headersMapEmpty,
-			expErr: nil,
-		},
-		{
-			name:   "headers-all",
-			event:  headersMapAll,
-			expMap: headersMapAll,
-			expErr: nil,
-		},
-	}
-
-	for _, tc := range testcases {
-		t.Run(tc.name, func(t *testing.T) {
-			event := map[string][]string{}
-			for k, v := range tc.event {
-				event[k] = []string{v, "dummy"}
-			}
-
-			tm, err := multiValueHeadersCarrier(event)
-			t.Logf("rawPayloadCarrier 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())
-			}
-			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
index 7de0df08e312a..d0537af1a3e5b 100644
--- a/pkg/serverless/trace/propagation/extractor.go
+++ b/pkg/serverless/trace/propagation/extractor.go
@@ -115,7 +115,7 @@ func (e Extractor) extract(event interface{}) (*TraceContext, error) {
 		if len(ev.Headers) > 0 {
 			carrier, err = headersCarrier(ev.Headers)
 		} else {
-			carrier, err = multiValueHeadersCarrier(ev.MultiValueHeaders)
+			carrier = tracer.HTTPHeadersCarrier(ev.MultiValueHeaders)
 		}
 	case events.LambdaFunctionURLRequest:
 		carrier, err = headersCarrier(ev.Headers)
diff --git a/pkg/serverless/trace/propagation/extractor_test.go b/pkg/serverless/trace/propagation/extractor_test.go
index c598b16befd18..cdcb4bd39b3ec 100644
--- a/pkg/serverless/trace/propagation/extractor_test.go
+++ b/pkg/serverless/trace/propagation/extractor_test.go
@@ -179,7 +179,7 @@ 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"}
+		mvh[k] = []string{v}
 	}
 	return mvh
 }