Add aws sns instrumentation

dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle

@@ -0,0 +1,32 @@
+
+muzzle {
+  pass {
+    group = "com.amazonaws"
+    module = "aws-java-sdk-sns"
+    versions = "[1.0.0,)"
+//    extraDependency 'com.amazonaws:amazon-sqs-java-messaging-lib:1.0.0'
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+//addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test')
+
+dependencies {
+  compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '1.12.676'
+//  compileOnly group: 'com.amazonaws', name: 'amazon-sqs-java-messaging-lib', version: '1.0.0'
+//
+//  // Include httpclient instrumentation for testing because it is a dependency for aws-sdk.
+//  testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
+//  testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-1.11.0')
+//  testImplementation project(':dd-java-agent:instrumentation:jms')
+
+  // SQS<->JMS testing:
+//  testImplementation group: 'org.elasticmq', name: 'elasticmq-rest-sqs_2.13', version: '1.4.7'
+//  testImplementation group: 'com.amazonaws', name: 'amazon-sqs-java-messaging-lib', version: '1.0.8'
+
+  latestDepTestImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '+'
+//  latestDepTestImplementation group: 'com.amazonaws', name: 'amazon-sqs-java-messaging-lib', version: '1.1.2'
+}

dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java

@@ -0,0 +1,9 @@
+package datadog.trace.instrumentation.aws.v1.sns;
+
+import datadog.trace.agent.tooling.InstrumenterModule;
+
+public abstract class AbstractSnsInstrumentation extends InstrumenterModule.Tracing {
+  public AbstractSnsInstrumentation() {
+    super("sns", "aws-sdk");
+  }
+}

dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java

@@ -0,0 +1,24 @@
+package datadog.trace.instrumentation.aws.v1.sns;
+
+import static datadog.trace.bootstrap.instrumentation.api.PathwayContext.DATADOG_KEY;
+
+import com.amazonaws.services.sns.model.MessageAttributeValue;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import java.util.Map;
+
+public class MessageAttributeInjector
+    implements AgentPropagation.Setter<Map<String, MessageAttributeValue>> {
+
+  public static final MessageAttributeInjector SETTER = new MessageAttributeInjector();
+
+  @Override
+  public void set(
+      final Map<String, MessageAttributeValue> carrier, final String key, final String value) {
+    if (carrier.size() < 10 && !carrier.containsKey(DATADOG_KEY)) {
+      String jsonPathway = String.format("{\"%s\": \"%s\"}", key, value);
+      carrier.put(
+          DATADOG_KEY,
+          new MessageAttributeValue().withDataType("String").withStringValue(jsonPathway));
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java

@@ -0,0 +1,70 @@
+package datadog.trace.instrumentation.aws.v1.sns;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.amazonaws.handlers.RequestHandler2;
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+
+/** AWS SDK v1 instrumentation */
+@AutoService(InstrumenterModule.class)
+public final class SnsClientInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType {
+  private static final String INSTRUMENTATION_NAME = "aws-sdk";
+
+  public SnsClientInstrumentation() {
+    super(INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "com.amazonaws.handlers.HandlerChainFactory";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("newRequestHandler2Chain")),
+        SnsClientInstrumentation.class.getName() + "$HandlerChainAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector"
+    };
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap(
+        "com.amazonaws.AmazonWebServiceRequest",
+        "datadog.trace.bootstrap.instrumentation.api.AgentSpan");
+  }
+
+  public static class HandlerChainAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void addHandler(@Advice.Return final List<RequestHandler2> handlers) {
+      if (Config.get().isDataStreamsEnabled()) {
+        for (RequestHandler2 interceptor : handlers) {
+          if (interceptor instanceof SnsInterceptor) {
+            return; // list already has our interceptor, return to builder
+          }
+        }
+        handlers.add(
+            new SnsInterceptor(
+                InstrumentationContext.get(
+                    "com.amazonaws.AmazonWebServiceRequest",
+                    "datadog.trace.bootstrap.instrumentation.api.AgentSpan")));
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java

@@ -0,0 +1,75 @@
+package datadog.trace.instrumentation.aws.v1.sns;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.bootstrap.instrumentation.api.PathwayContext.DATADOG_KEY;
+import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName;
+import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
+import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
+import static datadog.trace.core.datastreams.TagsProcessor.TOPIC_TAG;
+import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
+import static datadog.trace.instrumentation.aws.v1.sns.MessageAttributeInjector.SETTER;
+
+import com.amazonaws.AmazonWebServiceRequest;
+import com.amazonaws.handlers.RequestHandler2;
+import com.amazonaws.services.sns.model.PublishBatchRequest;
+import com.amazonaws.services.sns.model.PublishBatchRequestEntry;
+import com.amazonaws.services.sns.model.PublishRequest;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.LinkedHashMap;
+
+public class SnsInterceptor extends RequestHandler2 {
+
+  private final ContextStore<AmazonWebServiceRequest, AgentSpan> contextStore;
+
+  public SnsInterceptor(ContextStore<AmazonWebServiceRequest, AgentSpan> contextStore) {
+    this.contextStore = contextStore;
+  }
+
+  @Override
+  public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request) {
+    if (request instanceof PublishRequest) {
+      PublishRequest pRequest = (PublishRequest) request;
+
+      String topicArn = pRequest.getTopicArn();
+      if (topicArn == null) return request;
+
+      LinkedHashMap<String, String> sortedTags = getTags(topicArn);
+
+      final AgentSpan span = newSpan(request);
+      // note: modifying message attributes has to be done before marshalling, otherwise the changes
+      // are not reflected in the actual request (and the MD5 check on send will fail).
+      propagate().injectPathwayContext(span, pRequest.getMessageAttributes(), SETTER, sortedTags);
+    } else if (request instanceof PublishBatchRequest) {
+      PublishBatchRequest pmbRequest = (PublishBatchRequest) request;
+
+      String topicArn = pmbRequest.getTopicArn();
+      if (topicArn == null) return request;
+
+      LinkedHashMap<String, String> sortedTags = getTags(topicArn);
+
+      final AgentSpan span = newSpan(request);
+      for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) {
+        propagate().injectPathwayContext(span, entry.getMessageAttributes(), SETTER, sortedTags);
+      }
+    }
+    return request;
+  }
+
+  private AgentSpan newSpan(AmazonWebServiceRequest request) {
+    final AgentSpan span = startSpan("aws.sns.send");
+    // pass the span to TracingRequestHandler in the sdk instrumentation where it'll be enriched &
+    // activated
+    contextStore.put(request, span);
+    return span;
+  }
+
+  private static LinkedHashMap<String, String> getTags(String topicArn) {
+    LinkedHashMap<String, String> sortedTags = new LinkedHashMap<>();
+    sortedTags.put(DIRECTION_TAG, DIRECTION_OUT);
+    sortedTags.put(TOPIC_TAG, urlFileName(topicArn));
+    sortedTags.put(TYPE_TAG, "sns");
+    return sortedTags;
+  }
+}

dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy

@@ -0,0 +1,46 @@
+
+import com.amazonaws.SDKGlobalConfiguration
+import com.amazonaws.auth.AWSStaticCredentialsProvider
+import com.amazonaws.auth.AnonymousAWSCredentials
+import datadog.trace.agent.test.naming.VersionedNamingTestBase
+import datadog.trace.api.config.GeneralConfig
+import spock.lang.Shared
+
+
+abstract class SnsClientTest extends VersionedNamingTestBase {
+
+  def setup() {
+    System.setProperty(SDKGlobalConfiguration.ACCESS_KEY_SYSTEM_PROPERTY, "my-access-key")
+    System.setProperty(SDKGlobalConfiguration.SECRET_KEY_SYSTEM_PROPERTY, "my-secret-key")
+  }
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    // Set a service name that gets sorted early with SORT_BY_NAMES
+    injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service")
+    injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, isDataStreamsEnabled().toString())
+  }
+
+  @Shared
+  def credentialsProvider = new AWSStaticCredentialsProvider(new AnonymousAWSCredentials())
+
+  @Override
+  String operation() {
+    null
+  }
+
+  @Override
+  String service() {
+    null
+  }
+
+  boolean hasTimeInQueueSpan() {
+    false
+  }
+
+  abstract String expectedOperation(String awsService, String awsOperation)
+
+  abstract String expectedService(String awsService, String awsOperation)
+
+}

settings.gradle

@@ -169,6 +169,7 @@ include ':dd-java-agent:instrumentation:aws-java-sdk-1.11.0'
 include ':dd-java-agent:instrumentation:aws-java-sdk-2.2'
 include ':dd-java-agent:instrumentation:aws-java-sqs-1.0'
 include ':dd-java-agent:instrumentation:aws-java-sqs-2.0'
+include ':dd-java-agent:instrumentation:aws-java-sns-1.0'
 include ':dd-java-agent:instrumentation:aws-lambda-handler'
 include ':dd-java-agent:instrumentation:axis-2'
 include ':dd-java-agent:instrumentation:axway-api'