Fix parsing of binary datadog headers in SQS (#7324)

DataDog · Jul 18, 2024 · 24b87f7 · 24b87f7
1 parent 144efa8
commit 24b87f7
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 3 deletions.
diff --git a/...c/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java b/...c/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java
@@ -39,7 +39,18 @@ public static void forEachProperty(AgentPropagation.KeyClassifier classifier, By
       return;
     }
     try {
-      forEachProperty(classifier, UTF_8.decode(BASE_64.decode(json)).toString());
+      String jsonStr;
+      // peak at the first character to know if we're dealing with json or base64 (since '{' is not
+      // a valid b64 char).
+      if (json.get(0) == '{') {
+        jsonStr = UTF_8.decode(json).toString();
+      } else {
+        // TODO remove this branch once we are sure that this is never happening.
+        // passing a base64 string in a byte buffer is a nonsense, since the base64 is a less
+        // efficient representation.
+        jsonStr = UTF_8.decode(BASE_64.decode(json)).toString();
+      }
+      forEachProperty(classifier, jsonStr);
     } catch (Exception e) {
       log.debug("Problem decoding _datadog context", e);
     }

diff --git a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy
@@ -229,7 +229,7 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
     when:
     def message = new Message()
     message.addMessageAttributesEntry('_datadog', new MessageAttributeValue().withDataType('Binary').withBinaryValue(
-      UTF_8.encode('eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI0OTQ4Mzc3MzE2MzU3MjkxNDIxIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjY3NDY5OTgwMTUwMzc0Mjk1MTIiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0=')
+      headerValue
       ))
     def messages = new TracingList([message], "http://localhost:${address.port}/000000000000/somequeue")
 
@@ -260,6 +260,13 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
         }
       }
     }
+
+    where:
+    headerValue << [
+      UTF_8.encode('{"x-datadog-trace-id":"4948377316357291421","x-datadog-parent-id":"6746998015037429512","x-datadog-sampling-priority":"1"}'),
+      // not sure this test case with base 64 corresponds to an actual use case
+      UTF_8.encode('eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI0OTQ4Mzc3MzE2MzU3MjkxNDIxIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjY3NDY5OTgwMTUwMzc0Mjk1MTIiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0=')
+    ]
   }
 
   @IgnoreIf({ instance.isDataStreamsEnabled() })

diff --git a/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy
@@ -236,7 +236,7 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
 
     when:
     def message = Message.builder().messageAttributes(['_datadog': MessageAttributeValue.builder().dataType('Binary').binaryValue(SdkBytes.fromByteBuffer(
-      UTF_8.encode('eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI0OTQ4Mzc3MzE2MzU3MjkxNDIxIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjY3NDY5OTgwMTUwMzc0Mjk1MTIiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0=')
+      headerValue
       )).build()]).build()
     def messages = new TracingList([message],
     "http://localhost:${address.port}/000000000000/somequeue",
@@ -270,6 +270,13 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
         }
       }
     }
+
+    where:
+    headerValue << [
+      UTF_8.encode('{"x-datadog-trace-id":"4948377316357291421","x-datadog-parent-id":"6746998015037429512","x-datadog-sampling-priority":"1"}'),
+      // not sure this test case with base 64 corresponds to an actual use case
+      UTF_8.encode('eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiI0OTQ4Mzc3MzE2MzU3MjkxNDIxIiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjY3NDY5OTgwMTUwMzc0Mjk1MTIiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIn0=')
+    ]
   }
 
   @IgnoreIf({instance.isDataStreamsEnabled()})