remove binary datadog attribute if present in JMS SQS instrumentation

DataDog · Jul 5, 2024 · 12cbc16 · 12cbc16
1 parent 0327c4d
commit 12cbc16
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 12 deletions.
diff --git a/.../src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java
@@ -7,10 +7,12 @@
 
 import com.amazon.sqs.javamessaging.message.SQSMessage;
 import com.amazonaws.services.sqs.model.Message;
+import com.amazonaws.services.sqs.model.MessageAttributeValue;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
+import java.util.HashMap;
 import java.util.Map;
 import javax.jms.JMSException;
 import net.bytebuddy.asm.Advice;
@@ -28,10 +30,27 @@ public String instrumentedType() {
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isConstructor().and(takesArgument(2, named("com.amazonaws.services.sqs.model.Message"))),
-        getClass().getName() + "$CopyTracePropertyAdvice");
+        getClass().getName() + "$JmsSqsMessageConstructorAdvice");
   }
 
-  public static class CopyTracePropertyAdvice {
+  public static class JmsSqsMessageConstructorAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(2) Message sqsMessage) {
+      Map<String, MessageAttributeValue> messageAttributes = sqsMessage.getMessageAttributes();
+      MessageAttributeValue ddAttribute = messageAttributes.get("_datadog");
+      if (ddAttribute != null && "Binary".equals(ddAttribute.getDataType())) {
+        // binary message attributes are not supported by amazon-sqs-java-messaging-lib, and there
+        // is a chance we might introduce one, either when the message was sent from SNS or from a
+        // DD-instrumented Javascript app.
+        // When we reach this point, the value would already have been used by the aws-sqs
+        // instrumentation, so we can safely remove it.
+        Map<String, MessageAttributeValue> messageAttributesCopy = new HashMap<>(messageAttributes);
+        // need to copy to remove because the original is an UnmodifiableMap
+        messageAttributesCopy.remove("_datadog");
+        sqsMessage.withMessageAttributes(messageAttributesCopy);
+      }
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(
         @Advice.Argument(2) Message sqsMessage, @Advice.FieldValue("properties") Map properties)

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
@@ -7,6 +7,7 @@ import com.amazonaws.client.builder.AwsClientBuilder
 import com.amazonaws.services.sqs.AmazonSQSClientBuilder
 import com.amazonaws.services.sqs.model.Message
 import com.amazonaws.services.sqs.model.MessageAttributeValue
+import com.amazonaws.services.sqs.model.SendMessageRequest
 import datadog.trace.agent.test.naming.VersionedNamingTestBase
 import datadog.trace.agent.test.utils.TraceUtils
 import datadog.trace.api.Config
@@ -24,6 +25,8 @@ import spock.lang.IgnoreIf
 import spock.lang.Shared
 
 import javax.jms.Session
+import java.nio.ByteBuffer
+import java.nio.charset.Charset
 
 import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 import static java.nio.charset.StandardCharsets.UTF_8
@@ -281,7 +284,8 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
     when:
     connection.start()
     TraceUtils.runUnderTrace('parent', {
-      client.sendMessage(queue.queueUrl, 'sometext')
+      def ddMsgAttribute = new MessageAttributeValue().withDataType("Binary").withBinaryValue(ByteBuffer.wrap("hello world".getBytes(Charset.defaultCharset())))
+      client.sendMessage(new SendMessageRequest(queue.queueUrl, 'sometext').withMessageAttributes([_datadog: ddMsgAttribute]))
     })
     def message = consumer.receive()
     consumer.receiveNoWait()
@@ -411,6 +415,7 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
     def expectedTraceProperty = 'X-Amzn-Trace-Id'.toLowerCase(Locale.ENGLISH).replace('-', '__dash__')
     assert message.getStringProperty(expectedTraceProperty) =~
     /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/
+    assert !message.propertyExists("_datadog")
 
     cleanup:
     session.close()

diff --git a/.../src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsJmsMessageInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsJmsMessageInstrumentation.java
@@ -10,10 +10,12 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
+import java.util.HashMap;
 import java.util.Map;
 import javax.jms.JMSException;
 import net.bytebuddy.asm.Advice;
 import software.amazon.awssdk.services.sqs.model.Message;
+import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
 
 @AutoService(InstrumenterModule.class)
 public class SqsJmsMessageInstrumentation extends AbstractSqsInstrumentation
@@ -29,10 +31,27 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isConstructor()
             .and(takesArgument(2, named("software.amazon.awssdk.services.sqs.model.Message"))),
-        getClass().getName() + "$CopyTracePropertyAdvice");
+        getClass().getName() + "$JmsSqsMessageConstructorAdvice");
   }
 
-  public static class CopyTracePropertyAdvice {
+  public static class JmsSqsMessageConstructorAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(value = 2, readOnly = false) Message sqsMessage) {
+      Map<String, MessageAttributeValue> messageAttributes = sqsMessage.messageAttributes();
+      MessageAttributeValue ddAttribute = messageAttributes.get("_datadog");
+      if (ddAttribute != null && "Binary".equals(ddAttribute.dataType())) {
+        // binary message attributes are not supported by amazon-sqs-java-messaging-lib, and there
+        // is a chance we might introduce one, either when the message was sent from SNS or from a
+        // DD-instrumented Javascript app.
+        // When we reach this point, the value would already have been used by the aws-sqs
+        // instrumentation, so we can safely remove it.
+        Map<String, MessageAttributeValue> messageAttributesCopy = new HashMap<>(messageAttributes);
+        // need to copy to remove because the original is an UnmodifiableMap
+        messageAttributesCopy.remove("_datadog");
+        sqsMessage = sqsMessage.toBuilder().messageAttributes(messageAttributesCopy).build();
+      }
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(
         @Advice.Argument(2) Message sqsMessage, @Advice.FieldValue("properties") Map properties)

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
@@ -1,14 +1,11 @@
-import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
-import static java.nio.charset.StandardCharsets.UTF_8
-
 import com.amazon.sqs.javamessaging.ProviderConfiguration
 import com.amazon.sqs.javamessaging.SQSConnectionFactory
 import datadog.trace.agent.test.naming.VersionedNamingTestBase
 import datadog.trace.agent.test.utils.TraceUtils
 import datadog.trace.api.Config
-import datadog.trace.api.DDTags
 import datadog.trace.api.DDSpanId
 import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTags
 import datadog.trace.api.config.GeneralConfig
 import datadog.trace.api.naming.SpanNaming
 import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
@@ -22,16 +19,20 @@ import software.amazon.awssdk.core.SdkSystemSetting
 import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.sqs.SqsClient
 import software.amazon.awssdk.services.sqs.model.CreateQueueRequest
-import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest
-import software.amazon.awssdk.services.sqs.model.SendMessageRequest
 import software.amazon.awssdk.services.sqs.model.Message
 import software.amazon.awssdk.services.sqs.model.MessageAttributeValue
+import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest
+import software.amazon.awssdk.services.sqs.model.SendMessageRequest
 import spock.lang.IgnoreIf
 import spock.lang.Shared
 
 import javax.jms.Session
 
+import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
+import static java.nio.charset.StandardCharsets.UTF_8
+
 abstract class SqsClientTest extends VersionedNamingTestBase {
+  private MessageAttributeValue ddMsgAttribute
 
   def setup() {
     System.setProperty(SdkSystemSetting.AWS_ACCESS_KEY_ID.property(), "my-access-key")
@@ -292,7 +293,8 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
     when:
     connection.start()
     TraceUtils.runUnderTrace('parent', {
-      client.sendMessage(SendMessageRequest.builder().queueUrl(queue.queueUrl).messageBody('sometext').build())
+      def ddMsgAttribute = MessageAttributeValue.builder().dataType("Binary").binaryValue(SdkBytes.fromUtf8String("hello world")).build()
+      client.sendMessage(SendMessageRequest.builder().queueUrl(queue.queueUrl).messageBody('sometext').messageAttributes([_datadog: ddMsgAttribute]).build())
     })
     def message = consumer.receive()
     consumer.receiveNoWait()
@@ -421,6 +423,7 @@ abstract class SqsClientTest extends VersionedNamingTestBase {
     def expectedTraceProperty = 'X-Amzn-Trace-Id'.toLowerCase(Locale.ENGLISH).replace('-', '__dash__')
     assert message.getStringProperty(expectedTraceProperty) =~
     /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/
+    assert !message.propertyExists("_datadog")
 
     cleanup:
     session.close()