From cb068c079954c48ecdb587b24ef1866847ceeba9 Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Mon, 30 Oct 2023 13:32:35 -0700
Subject: [PATCH] Fix the raw-message-delivery case.

From https://github.com/DataDog/datadog-agent/pull/19888
---
 pkg/serverless/trace/propagation/carriers.go  | 21 ++++++---
 .../trace/propagation/carriers_test.go        | 46 +++++++++++++++++--
 2 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/pkg/serverless/trace/propagation/carriers.go b/pkg/serverless/trace/propagation/carriers.go
index 7c2ba772f7fdb..6691912891a48 100644
--- a/pkg/serverless/trace/propagation/carriers.go
+++ b/pkg/serverless/trace/propagation/carriers.go
@@ -93,14 +93,23 @@ func sqsMessageCarrier(event events.SQSMessage) (tracer.TextMapReader, error) {
 // context from the events.SQSMessageAttribute field on an events.SQSMessage
 // type.
 func sqsMessageAttrCarrier(attr events.SQSMessageAttribute) (tracer.TextMapReader, error) {
-	var carrier tracer.TextMapCarrier
-	if attr.DataType != "String" {
+	var bytes []byte
+	switch attr.DataType {
+	case "String":
+		if attr.StringValue == nil {
+			return nil, errors.New("String value not found in _datadog payload")
+		}
+		bytes = []byte(*attr.StringValue)
+	case "Binary":
+		// SNS => SQS => Lambda with SQS's subscription to SNS has enabled RAW
+		// MESSAGE DELIVERY option
+		bytes = attr.BinaryValue // No need to decode base64 because already decoded
+	default:
 		return nil, errors.New("Unsupported DataType in _datadog payload")
 	}
-	if attr.StringValue == nil {
-		return nil, errors.New("String value not found in _datadog payload")
-	}
-	if err := json.Unmarshal([]byte(*attr.StringValue), &carrier); err != nil {
+
+	var carrier tracer.TextMapCarrier
+	if err := json.Unmarshal(bytes, &carrier); err != nil {
 		return nil, fmt.Errorf("Error unmarshaling payload value: %w", err)
 	}
 	return carrier, nil
diff --git a/pkg/serverless/trace/propagation/carriers_test.go b/pkg/serverless/trace/propagation/carriers_test.go
index 0b3e54a4ddd75..261590799ad00 100644
--- a/pkg/serverless/trace/propagation/carriers_test.go
+++ b/pkg/serverless/trace/propagation/carriers_test.go
@@ -36,7 +36,7 @@ func TestSQSMessageAttrCarrier(t *testing.T) {
 		expNoErr bool
 	}{
 		{
-			name: "datadog-map",
+			name: "string-datadog-map",
 			attr: events.SQSMessageAttribute{
 				DataType:    "String",
 				StringValue: aws.String(headersAll),
@@ -45,7 +45,7 @@ func TestSQSMessageAttrCarrier(t *testing.T) {
 			expNoErr: true,
 		},
 		{
-			name: "empty-map",
+			name: "string-empty-map",
 			attr: events.SQSMessageAttribute{
 				DataType:    "String",
 				StringValue: aws.String("{}"),
@@ -54,7 +54,7 @@ func TestSQSMessageAttrCarrier(t *testing.T) {
 			expNoErr: true,
 		},
 		{
-			name: "empty-string",
+			name: "string-empty-string",
 			attr: events.SQSMessageAttribute{
 				DataType:    "String",
 				StringValue: aws.String(""),
@@ -63,7 +63,7 @@ func TestSQSMessageAttrCarrier(t *testing.T) {
 			expNoErr: false,
 		},
 		{
-			name: "nil-string",
+			name: "string-nil-string",
 			attr: events.SQSMessageAttribute{
 				DataType:    "String",
 				StringValue: nil,
@@ -71,10 +71,46 @@ func TestSQSMessageAttrCarrier(t *testing.T) {
 			expMap:   nil,
 			expNoErr: false,
 		},
+		{
+			name: "binary-datadog-map",
+			attr: events.SQSMessageAttribute{
+				DataType:    "Binary",
+				BinaryValue: []byte(headersAll),
+			},
+			expMap:   headersMapAll,
+			expNoErr: true,
+		},
+		{
+			name: "binary-empty-map",
+			attr: events.SQSMessageAttribute{
+				DataType:    "Binary",
+				BinaryValue: []byte("{}"),
+			},
+			expMap:   map[string]string{},
+			expNoErr: true,
+		},
+		{
+			name: "binary-empty-string",
+			attr: events.SQSMessageAttribute{
+				DataType:    "Binary",
+				BinaryValue: []byte(""),
+			},
+			expMap:   nil,
+			expNoErr: false,
+		},
+		{
+			name: "binary-nil-string",
+			attr: events.SQSMessageAttribute{
+				DataType:    "Binary",
+				BinaryValue: nil,
+			},
+			expMap:   nil,
+			expNoErr: false,
+		},
 		{
 			name: "wrong-data-type",
 			attr: events.SQSMessageAttribute{
-				DataType: "Binary",
+				DataType: "Purple",
 			},
 			expMap:   nil,
 			expNoErr: false,