Merge pull request #206 from DataDog/tyler/aws

Prevent headers from being added for AWS client calls

Author: tylerbenson

dd-java-agent-ittests/dd-java-agent-ittests.gradle

@@ -37,14 +37,9 @@ dependencies {
   testCompile group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '8.2.0.v20160908'
   testCompile group: 'org.apache.tomcat.embed', name: 'tomcat-embed-core', version: '8.0.41'
   testCompile group: 'org.apache.tomcat.embed', name: 'tomcat-embed-jasper', version: '8.0.41'
-  // note: aws transitively pulls in apache-httpclient, which we also test against:
-  // testCompile group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.5.3'
-  testCompile(group: 'com.amazonaws', name: 'aws-java-sdk', version: '1.11.119')
+  testCompile group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.3'
   testCompile group: 'com.squareup.okhttp3', name: 'okhttp', version: '3.6.0'
-  testCompile(group: 'com.datastax.cassandra', name: 'cassandra-driver-core', version: '3.2.0')
-  testCompile(group: 'org.cassandraunit', name: 'cassandra-unit', version: '3.1.3.2') {
-    exclude(module: 'netty-handler')
-  }
+
   testCompile group: 'org.apache.logging.log4j', name: 'log4j-api', version: '2.8.2'
   testCompile group: 'org.apache.logging.log4j', name: 'log4j-core', version: '2.8.2'
 

dd-java-agent-ittests/src/test/java/datadog/trace/agent/integration/AWSInstrumentationTest.java

@@ -127,8 +127,15 @@ private CloseableHttpResponse createNetworkSpan(
             .startActive(true);
 
     final Span networkSpan = networkScope.span();
-    tracer.inject(
-        networkSpan.context(), Format.Builtin.HTTP_HEADERS, new HttpHeadersInjectAdapter(request));
+
+    final boolean awsClientCall = request.getHeaders("amz-sdk-invocation-id").length > 0;
+    // AWS calls are often signed, so we can't add headers without breaking the signature.
+    if (!awsClientCall) {
+      tracer.inject(
+          networkSpan.context(),
+          Format.Builtin.HTTP_HEADERS,
+          new HttpHeadersInjectAdapter(request));
+    }
 
     try {
       // request tags

dd-java-agent/instrumentation/apache-httpclient-4.3/src/main/java/datadog/trace/instrumentation/apachehttpclient/DDTracingClientExec.java

@@ -15,6 +15,17 @@ versionScan {
 
 apply from: "${rootDir}/gradle/java.gradle"
 
+if (JavaVersion.current() != JavaVersion.VERSION_1_8) {
+  sourceSets {
+    test {
+      groovy {
+        // These classes use Ratpack which requires Java 8. (Currently also incompatible with Java 9.)
+        exclude '**/AWSClientTest.groovy'
+      }
+    }
+  }
+}
+
 dependencies {
   compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-core', version: '1.11.119'
   compile('io.opentracing.contrib:opentracing-aws-sdk:0.0.3') {
@@ -27,4 +38,9 @@ dependencies {
   compile deps.bytebuddy
   compile deps.opentracing
   compile deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+  // Include httpclient instrumentation for testing because it is a dependency for aws-sdk.
+  testCompile project(':dd-java-agent:instrumentation:apache-httpclient-4.3')
+  testCompile group: 'com.amazonaws', name: 'aws-java-sdk', version: '1.11.119'
 }

dd-java-agent/instrumentation/aws-sdk/aws-sdk.gradle

@@ -17,6 +17,7 @@
 import io.opentracing.contrib.aws.TracingRequestHandler;
 import io.opentracing.util.GlobalTracer;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
@@ -48,21 +49,22 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
   }
 
   public static class AWSClientAdvice {
-    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void addHandler(@Advice.This final AwsClientBuilder<?, ?> builder) {
       List<RequestHandler2> handlers = builder.getRequestHandlers();
       boolean hasDDHandler = false;
       if (null == handlers) {
-        handlers = new ArrayList<RequestHandler2>(1);
+        handlers = Collections.emptyList();
       } else {
-        for (RequestHandler2 handler : handlers) {
+        for (final RequestHandler2 handler : handlers) {
           if (handler instanceof TracingRequestHandler) {
             hasDDHandler = true;
             break;
           }
         }
       }
       if (!hasDDHandler) {
+        handlers = new ArrayList<>(handlers); // copy since returned list is unmodifiable
         handlers.add(new TracingRequestHandler(GlobalTracer.get()));
         builder.setRequestHandlers(handlers.toArray(new RequestHandler2[0]));
       }

dd-java-agent/instrumentation/aws-sdk/src/main/java/datadog/trace/instrumentation/aws/AWSClientInstrumentation.java

@@ -0,0 +1,142 @@
+import com.amazonaws.AmazonWebServiceClient
+import com.amazonaws.auth.AWSStaticCredentialsProvider
+import com.amazonaws.auth.AnonymousAWSCredentials
+import com.amazonaws.client.builder.AwsClientBuilder
+import com.amazonaws.handlers.RequestHandler2
+import com.amazonaws.regions.Regions
+import com.amazonaws.services.s3.AmazonS3ClientBuilder
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.DDTags
+import io.opentracing.tag.Tags
+import ratpack.http.Headers
+
+import java.util.concurrent.atomic.AtomicReference
+
+import static ratpack.groovy.test.embed.GroovyEmbeddedApp.ratpack
+
+class AWSClientTest extends AgentTestRunner {
+
+  def "request handler is hooked up"() {
+    setup:
+    def builder = AmazonS3ClientBuilder.standard()
+      .withRegion(Regions.US_EAST_1)
+    if (addHandler) {
+      builder.withRequestHandlers(new RequestHandler2() {})
+    }
+    AmazonWebServiceClient client = builder.build()
+
+    expect:
+    client.requestHandler2s != null
+    client.requestHandler2s.size() == size
+    client.requestHandler2s.get(position).getClass().getSimpleName() == "TracingRequestHandler"
+
+    where:
+    addHandler | size | position
+    true       | 2    | 1
+    false      | 1    | 0
+  }
+
+  def "send request with mocked back end"() {
+    setup:
+    def receivedHeaders = new AtomicReference<Headers>()
+    def server = ratpack {
+      handlers {
+        all {
+          receivedHeaders.set(request.headers)
+          response.status(200).send("pong")
+        }
+      }
+    }
+    AwsClientBuilder.EndpointConfiguration endpoint = new AwsClientBuilder.EndpointConfiguration("http://localhost:$server.address.port", "us-west-2");
+
+    AmazonWebServiceClient client = AmazonS3ClientBuilder
+      .standard()
+      .withPathStyleAccessEnabled(true)
+      .withEndpointConfiguration(endpoint)
+      .withCredentials(new AWSStaticCredentialsProvider(new AnonymousAWSCredentials()))
+      .build()
+
+    def bucket = client.createBucket("testbucket")
+
+    expect:
+    bucket != null
+
+    client.requestHandler2s != null
+    client.requestHandler2s.size() == 1
+    client.requestHandler2s.get(0).getClass().getSimpleName() == "TracingRequestHandler"
+
+    TEST_WRITER.size() == 2
+
+    def trace = TEST_WRITER.get(0)
+    trace.size() == 2
+
+    and: // span 0 - from apache-httpclient instrumentation
+    def span1 = trace[0]
+
+    span1.context().operationName == "PUT"
+    span1.serviceName == "unnamed-java-app"
+    span1.resourceName == "PUT"
+    span1.type == null
+    !span1.context().getErrorFlag()
+    span1.context().parentId == 0
+
+
+    def tags1 = span1.context().tags
+    tags1["component"] == "apache-httpclient"
+    tags1["thread.name"] != null
+    tags1["thread.id"] != null
+    tags1.size() == 3
+
+    and: // span 1 - from aws instrumentation
+    def span2 = trace[1]
+
+    span2.context().operationName == "PUT"
+    span2.serviceName == "unnamed-java-app"
+    span2.resourceName == "PUT"
+    span2.type == null
+    !span2.context().getErrorFlag()
+    span2.context().parentId == span1.spanId
+
+
+    def tags2 = span2.context().tags
+    tags2[Tags.SPAN_KIND.key] == Tags.SPAN_KIND_CLIENT
+    tags2[Tags.HTTP_METHOD.key] == "PUT"
+    tags2[Tags.HTTP_URL.key] == "http://localhost:$server.address.port/testbucket/"
+    tags2[Tags.PEER_HOSTNAME.key] == "localhost"
+    tags2[Tags.PEER_PORT.key] == server.address.port
+    tags2[DDTags.THREAD_NAME] != null
+    tags2[DDTags.THREAD_ID] != null
+    tags2.size() == 8
+
+    and:
+
+    def trace2 = TEST_WRITER.get(1)
+    trace2.size() == 1
+
+    and: // span 0 - from aws instrumentation
+    def span = trace2[0]
+
+    span.context().operationName == "Amazon S3"
+    span.serviceName == "unnamed-java-app"
+    span.resourceName == "Amazon S3"
+    span.type == null
+    !span.context().getErrorFlag()
+    span.context().parentId == 0
+
+    def tags = span.context().tags
+    tags["component"] == "java-aws-sdk"
+    tags2[Tags.SPAN_KIND.key] == Tags.SPAN_KIND_CLIENT
+    tags2[Tags.HTTP_METHOD.key] == "PUT"
+    tags2[Tags.HTTP_URL.key] == "http://localhost:$server.address.port/testbucket/"
+    tags2[Tags.HTTP_STATUS.key] == 200
+    tags["thread.name"] != null
+    tags["thread.id"] != null
+    tags.size() == 7
+
+    receivedHeaders.get().get("x-datadog-trace-id") == "$span.traceId"
+    receivedHeaders.get().get("x-datadog-parent-id") == "$span.spanId"
+
+    cleanup:
+    server.close()
+  }
+}

dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy

@@ -9,7 +9,7 @@
 import javax.jms.Message;
 
 public class MessagePropertyTextMap implements TextMap {
-  static final String DASH = "_$dash$_";
+  static final String DASH = "__dash__";
 
   private final Message message;