Create S3 instrumentation + add span pointers

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/spanpointers/SpanPointersHelper.java

@@ -0,0 +1,74 @@
+package datadog.trace.bootstrap.instrumentation.spanpointers;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
+import datadog.trace.bootstrap.instrumentation.api.SpanLink;
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+
+public final class SpanPointersHelper {
+  public static final String S3_PTR_KIND = "aws.s3.object";
+  public static final String LINK_KIND = "span-pointer";
+
+  // The pointer direction will always be down. The agent handles cases where the direction is up.
+  public static final String DOWN_DIRECTION = "d";
+
+  /**
+   * Generates a unique hash from an array of strings by joining them with | before hashing. Used to
+   * uniquely identify AWS requests for span pointers.
+   *
+   * @param components Array of strings to hash
+   * @return A 32-character hash uniquely identifying the components
+   * @throws NoSuchAlgorithmException this should never happen; but should be handled just in case.
+   */
+  private static String generatePointerHash(String[] components) throws NoSuchAlgorithmException {
+    MessageDigest messageDigest = MessageDigest.getInstance("SHA-256");
+
+    // Update the digest incrementally for each component.
+    boolean first = true;
+    for (String component : components) {
+      if (!first) {
+        messageDigest.update((byte) '|');
+      } else {
+        first = false;
+      }
+      messageDigest.update(component.getBytes(StandardCharsets.UTF_8));
+    }
+
+    byte[] fullHash = messageDigest.digest();
+    StringBuilder hex = new StringBuilder(32);
+    for (int i = 0; i < 16; i++) {
+      hex.append(String.format("%02x", fullHash[i]));
+    }
+
+    return hex.toString();
+  }
+
+  /**
+   * Adds a span pointer to the given span, using the SHA-256 hash of the components.
+   *
+   * @param span The span to add the pointer to
+   * @param kind Identifies which hashing rules to follow
+   * @param components Array of strings to hash, following span pointer rules
+   * @throws NoSuchAlgorithmException if unable to calculate hash
+   * @see <a href="https://github.com/DataDog/dd-span-pointer-rules/tree/main">Span pointer
+   *     rules</a>
+   */
+  public static void addSpanPointer(AgentSpan span, String kind, String[] components)
+      throws NoSuchAlgorithmException {
+    SpanAttributes attributes =
+        (SpanAttributes)
+            SpanAttributes.builder()
+                .put("ptr.kind", kind)
+                .put("ptr.dir", DOWN_DIRECTION)
+                .put("ptr.hash", generatePointerHash(components))
+                .put("link.kind", LINK_KIND)
+                .build();
+
+    AgentTracer.NoopContext zeroContext = AgentTracer.NoopContext.INSTANCE;
+    span.addLink(SpanLink.from(zeroContext, AgentSpanLink.DEFAULT_FLAGS, "", attributes));
+  }
+}

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/spanpointers/SpanPointersHelperTest.groovy

@@ -0,0 +1,75 @@
+package datadog.trace.bootstrap.instrumentation.spanpointers
+
+import datadog.trace.api.DDSpanId
+import datadog.trace.api.DDTraceId
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
+import datadog.trace.bootstrap.instrumentation.api.SpanLink
+import spock.lang.Specification
+
+class SpanPointersHelperTest extends Specification {
+  def "addSpanPointer adds correct link to span with basic values"() {
+    given:
+    AgentSpan span = Mock(AgentSpan)
+    String kind = SpanPointersHelper.S3_PTR_KIND
+    String[] components = ["some-bucket", "some-key.data", "ab12ef34"]
+    String expectedHash = "e721375466d4116ab551213fdea08413"
+
+    when:
+    SpanPointersHelper.addSpanPointer(span, kind, components)
+
+    then:
+    1 * span.addLink({ SpanLink link ->
+      assert link.traceId() == DDTraceId.ZERO
+      assert link.spanId() == DDSpanId.ZERO
+      assert link.attributes.asMap().get("ptr.kind") == kind
+      assert link.attributes.asMap().get("ptr.dir") == SpanPointersHelper.DOWN_DIRECTION
+      assert link.attributes.asMap().get("ptr.hash") == expectedHash
+      assert link.attributes.asMap().get("link.kind") == SpanPointersHelper.LINK_KIND
+      true
+    })
+  }
+
+  def "addSpanPointer adds correct link to span with non-ascii key"() {
+    given:
+    AgentSpan span = Mock(AgentSpan)
+    String kind = SpanPointersHelper.S3_PTR_KIND
+    String[] components = ["some-bucket", "some-key.你好", "ab12ef34"]
+    String expectedHash = "d1333a04b9928ab462b5c6cadfa401f4"
+
+    when:
+    SpanPointersHelper.addSpanPointer(span, kind, components)
+
+    then:
+    1 * span.addLink({ SpanLink link ->
+      assert link.traceId() == DDTraceId.ZERO
+      assert link.spanId() == DDSpanId.ZERO
+      assert link.attributes.asMap().get("ptr.kind") == kind
+      assert link.attributes.asMap().get("ptr.dir") == SpanPointersHelper.DOWN_DIRECTION
+      assert link.attributes.asMap().get("ptr.hash") == expectedHash
+      assert link.attributes.asMap().get("link.kind") == SpanPointersHelper.LINK_KIND
+      true
+    })
+  }
+
+  def "addSpanPointer adds correct link to span with multipart-upload"() {
+    given:
+    AgentSpan span = Mock(AgentSpan)
+    String kind = SpanPointersHelper.S3_PTR_KIND
+    String[] components = ["some-bucket", "some-key.data", "ab12ef34-5"]
+    String expectedHash = "2b90dffc37ebc7bc610152c3dc72af9f"
+
+    when:
+    SpanPointersHelper.addSpanPointer(span, kind, components)
+
+    then:
+    1 * span.addLink({ SpanLink link ->
+      assert link.traceId() == DDTraceId.ZERO
+      assert link.spanId() == DDSpanId.ZERO
+      assert link.attributes.asMap().get("ptr.kind") == kind
+      assert link.attributes.asMap().get("ptr.dir") == SpanPointersHelper.DOWN_DIRECTION
+      assert link.attributes.asMap().get("ptr.hash") == expectedHash
+      assert link.attributes.asMap().get("link.kind") == SpanPointersHelper.LINK_KIND
+      true
+    })
+  }
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/build.gradle

@@ -0,0 +1,29 @@
+muzzle {
+  pass {
+    group = "software.amazon.awssdk"
+    module = "s3"
+    versions = "[2.10.36,3)"
+    assertInverse = true
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test')
+
+dependencies {
+  compileOnly group: 'software.amazon.awssdk', name: 's3', version: '2.29.26'
+
+  // Include httpclient instrumentation for testing because it is a dependency for aws-sdk.
+  testRuntimeOnly project(':dd-java-agent:instrumentation:apache-httpclient-4')
+  testRuntimeOnly project(':dd-java-agent:instrumentation:aws-java-sdk-2.2')
+  testImplementation 'software.amazon.awssdk:s3:2.29.26'
+  testImplementation 'org.testcontainers:localstack:1.20.1'
+
+  latestDepTestImplementation group: 'software.amazon.awssdk', name: 's3', version: '+'
+}
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3ClientInstrumentation.java

@@ -0,0 +1,48 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import java.util.List;
+import net.bytebuddy.asm.Advice;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+
+@AutoService(InstrumenterModule.class)
+public final class S3ClientInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public S3ClientInstrumentation() {
+    super("s3", "aws-s3");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("resolveExecutionInterceptors")),
+        S3ClientInstrumentation.class.getName() + "$AwsS3BuilderAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".S3Interceptor", packageName + ".TextMapInjectAdapter"};
+  }
+
+  public static class AwsS3BuilderAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void addHandler(@Advice.Return final List<ExecutionInterceptor> interceptors) {
+      for (ExecutionInterceptor interceptor : interceptors) {
+        if (interceptor instanceof S3Interceptor) {
+          return; // list already has our interceptor, return to builder
+        }
+      }
+      interceptors.add(new S3Interceptor());
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java

@@ -0,0 +1,83 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import static datadog.trace.bootstrap.instrumentation.spanpointers.SpanPointersHelper.S3_PTR_KIND;
+import static datadog.trace.bootstrap.instrumentation.spanpointers.SpanPointersHelper.addSpanPointer;
+
+import datadog.trace.bootstrap.InstanceStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest;
+import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
+import software.amazon.awssdk.services.s3.model.CopyObjectRequest;
+import software.amazon.awssdk.services.s3.model.CopyObjectResponse;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+
+public class S3Interceptor implements ExecutionInterceptor {
+  private static final Logger log = LoggerFactory.getLogger(S3Interceptor.class);
+
+  public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
+      InstanceStore.of(ExecutionAttribute.class)
+          .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
+
+  @Override
+  public void afterExecution(
+      Context.AfterExecution context, ExecutionAttributes executionAttributes) {
+    String flag = System.getenv("DD_AWS_SDK_ADD_SPAN_POINTERS");
+    if (flag != null && flag.equalsIgnoreCase("false")) {
+      return;
+    }
+
+    AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
+    if (span == null) {
+      log.debug("Unable to find S3 request span. Not creating span pointer.");
+      return;
+    }
+
+    String bucket, key, eTag;
+    Object request = context.request();
+    Object response = context.response();
+
+    // Get bucket, key, and eTag for hash calculation.
+    // https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/s3/S3Client.html
+    if (request instanceof PutObjectRequest) {
+      PutObjectRequest putObjectRequest = (PutObjectRequest) request;
+      bucket = putObjectRequest.bucket();
+      key = putObjectRequest.key();
+      eTag = ((PutObjectResponse) response).eTag();
+    } else if (request instanceof CopyObjectRequest) {
+      CopyObjectRequest copyObjectRequest = (CopyObjectRequest) request;
+      bucket = copyObjectRequest.destinationBucket();
+      key = copyObjectRequest.destinationKey();
+      eTag = ((CopyObjectResponse) response).copyObjectResult().eTag();
+    } else if (request instanceof CompleteMultipartUploadRequest) {
+      CompleteMultipartUploadRequest completeMultipartUploadRequest =
+          (CompleteMultipartUploadRequest) request;
+      bucket = completeMultipartUploadRequest.bucket();
+      key = completeMultipartUploadRequest.key();
+      eTag = ((CompleteMultipartUploadResponse) response).eTag();
+    } else {
+      return;
+    }
+
+    // Hash calculation rules:
+    // https://github.com/DataDog/dd-span-pointer-rules/blob/main/AWS/S3/Object/README.md
+    if (eTag != null
+        && !eTag.isEmpty()
+        && eTag.charAt(0) == '"'
+        && eTag.charAt(eTag.length() - 1) == '"') {
+      eTag = eTag.substring(1, eTag.length() - 1);
+    }
+    String[] components = new String[] {bucket, key, eTag};
+    try {
+      addSpanPointer(span, S3_PTR_KIND, components);
+    } catch (Exception e) {
+      log.debug("Failed to add span pointer: {}", e.getMessage());
+    }
+  }
+}

dd-java-agent/instrumentation/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/TextMapInjectAdapter.java

@@ -0,0 +1,13 @@
+package datadog.trace.instrumentation.aws.v2.s3;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+
+public class TextMapInjectAdapter implements AgentPropagation.Setter<StringBuilder> {
+
+  public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter();
+
+  @Override
+  public void set(final StringBuilder builder, final String key, final String value) {
+    builder.append('"').append(key).append("\":\"").append(value).append("\",");
+  }
+}