From 12cbc165e362cb6600c03f5eb6172037f805ab24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Fri, 5 Jul 2024 13:39:45 +0200 Subject: [PATCH] remove binary datadog attribute if present in JMS SQS instrumentation --- .../v1/sqs/SqsJmsMessageInstrumentation.java | 23 +++++++++++++++++-- .../src/test/groovy/SqsClientTest.groovy | 7 +++++- .../v2/sqs/SqsJmsMessageInstrumentation.java | 23 +++++++++++++++++-- .../src/test/groovy/SqsClientTest.groovy | 17 ++++++++------ 4 files changed, 58 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java index 03eae7b7d26..d645cb40bf9 100644 --- a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsJmsMessageInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sqs-1.0/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 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 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 index 5042347483f..c3e810ab360 100644 --- 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/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsJmsMessageInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsJmsMessageInstrumentation.java index 287994a4356..070094b35ba 100644 --- a/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsJmsMessageInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sqs-2.0/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 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 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 index 3e505cc12ac..859aee2881f 100644 --- 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()