Fix the raw-message-delivery case.

From #19888
DataDog · Nov 3, 2023 · cb068c0 · cb068c0
1 parent d4d74aa
commit cb068c0
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 11 deletions.
diff --git 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
@@ -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,18 +63,54 @@ func TestSQSMessageAttrCarrier(t *testing.T) {
 			expNoErr: false,
 		},
 		{
-			name: "nil-string",
+			name: "string-nil-string",
 			attr: events.SQSMessageAttribute{
 				DataType:    "String",
 				StringValue: nil,
 			},
 			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,