Extract AWS payload tags

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

@@ -19,6 +19,11 @@ addTestSuiteExtendingForDir('dsmForkedTest', 'dsmTest', 'dsmTest')
 addTestSuiteForDir('latestDsmTest', 'dsmTest')
 addTestSuiteExtendingForDir('latestDsmForkedTest', 'latestDsmTest', 'dsmTest')
 
+addTestSuite('payloadTaggingTest')
+addTestSuiteExtendingForDir('payloadTaggingForkedTest', 'payloadTaggingTest', 'payloadTaggingTest')
+addTestSuiteForDir('latestPayloadTaggingTest', 'payloadTaggingTest')
+addTestSuiteExtendingForDir('latestPayloadTaggingForkedTest', 'latestPayloadTaggingTest', 'payloadTaggingTest')
+
 def fixedSdkVersion = '2.20.33' // 2.20.34 is missing and breaks IDEA import
 
 dependencies {
@@ -40,6 +45,8 @@ dependencies {
   testImplementation group: 'org.eclipse.jetty.http2', name: 'http2-server', version: '9.4.56.v20240826'
 
 
+  testImplementation group: 'org.testcontainers', name: 'localstack', version: libs.versions.testcontainers.get()
+
   // First version where dsm traced operations have required StreamARN parameter for kinesis
   // and publishBatch is available for SNS
   dsmTestImplementation group: 'software.amazon.awssdk', name: 'apache-client', version: '2.18.40'
@@ -49,10 +56,25 @@ dependencies {
   latestDsmTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: '+'
   latestDsmTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '+'
 
+  payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'apigateway', version: '2.19.0'
+  payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: '2.7.4'
+  payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sqs', version: '2.18.40'
+  payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '2.18.40'
+  payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 's3', version: '2.18.40'
+  payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: '2.18.40'
+  latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'apigateway', version: '2.25.40'
+  latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: '2.25.40'
+  latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sqs', version: '2.25.40'
+  latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '2.25.40'
+  latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 's3', version: '2.18.40'
+  latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: '2.18.40'
+
   latestDepTestImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
   latestDepTestImplementation project(':dd-java-agent:instrumentation:netty-4.1')
 
   latestDepTestImplementation group: 'software.amazon.awssdk', name: 'apache-client', version: fixedSdkVersion
+  latestDepTestImplementation group: 'software.amazon.awssdk', name: 'apigateway', version: fixedSdkVersion
+  latestDepTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: fixedSdkVersion
   latestDepTestImplementation group: 'software.amazon.awssdk', name: 's3', version: fixedSdkVersion
   latestDepTestImplementation group: 'software.amazon.awssdk', name: 'rds', version: fixedSdkVersion
   latestDepTestImplementation group: 'software.amazon.awssdk', name: 'ec2', version: fixedSdkVersion
@@ -61,3 +83,7 @@ dependencies {
   latestDepTestImplementation group: 'software.amazon.awssdk', name: 'dynamodb', version: fixedSdkVersion
   latestDepTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: fixedSdkVersion
 }
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+}

dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java

@@ -7,6 +7,7 @@
 import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
 
 import datadog.trace.api.Config;
+import datadog.trace.api.ConfigDefaults;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.cache.DDCache;
 import datadog.trace.api.cache.DDCaches;
@@ -23,16 +24,22 @@
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
 import datadog.trace.core.datastreams.TagsProcessor;
+import datadog.trace.payloadtags.PayloadTagsData;
 import java.net.URI;
 import java.time.Instant;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nonnull;
 import software.amazon.awssdk.awscore.AwsResponse;
+import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.SdkField;
 import software.amazon.awssdk.core.SdkPojo;
 import software.amazon.awssdk.core.SdkRequest;
@@ -115,6 +122,12 @@ public AgentSpan onSdkRequest(
     final String awsOperationName = attributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
     onOperation(span, awsServiceName, awsOperationName);
 
+    Config config = Config.get();
+    if (config.isCloudRequestPayloadTaggingEnabled()
+        && config.isCloudPayloadTaggingEnabledFor(awsServiceName)) {
+      awsPojoToTags(span, ConfigDefaults.DEFAULT_TRACE_CLOUD_PAYLOAD_REQUEST_TAG, request);
+    }
+
     // S3
     request.getValueForField("Bucket", String.class).ifPresent(name -> setBucketName(span, name));
     if ("s3".equalsIgnoreCase(awsServiceName) && span.traceConfig().isDataStreamsEnabled()) {
@@ -295,6 +308,14 @@ public AgentSpan onSdkResponse(
       final SdkResponse response,
       final SdkHttpResponse httpResponse,
       final ExecutionAttributes attributes) {
+
+    Config config = Config.get();
+    String serviceName = attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
+    if (config.isCloudResponsePayloadTaggingEnabled()
+        && config.isCloudPayloadTaggingEnabledFor(serviceName)) {
+      awsPojoToTags(span, ConfigDefaults.DEFAULT_TRACE_CLOUD_PAYLOAD_RESPONSE_TAG, response);
+    }
+
     if (response instanceof AwsResponse) {
       span.setTag(
           InstrumentationTags.AWS_REQUEST_ID,
@@ -437,4 +458,48 @@ protected String getRequestHeader(SdkHttpRequest request, String headerName) {
   protected String getResponseHeader(SdkHttpResponse response, String headerName) {
     return response.firstMatchingHeader(headerName).orElse(null);
   }
+
+  private void awsPojoToTags(AgentSpan span, String tagsPrefix, Object pojo) {
+    Collection<PayloadTagsData.PathAndValue> payloadTagsData = new ArrayList<>();
+    ArrayDeque<Object> path = new ArrayDeque<>();
+    collectPayloadTagsData(payloadTagsData, path, pojo);
+    span.setTag(
+        tagsPrefix,
+        new PayloadTagsData(payloadTagsData.toArray(new PayloadTagsData.PathAndValue[0])));
+  }
+
+  private void collectPayloadTagsData(
+      Collection<PayloadTagsData.PathAndValue> payloadTagsData,
+      ArrayDeque<Object> path,
+      Object object) {
+    if (object instanceof SdkPojo) {
+      SdkPojo pojo = (SdkPojo) object;
+      for (SdkField<?> field : pojo.sdkFields()) {
+        Object val = field.getValueOrDefault(pojo);
+        path.push(field.locationName());
+        collectPayloadTagsData(payloadTagsData, path, val);
+        path.pop();
+      }
+    } else if (object instanceof Collection) {
+      int index = 0;
+      for (Object value : (Collection<?>) object) {
+        path.push(index);
+        collectPayloadTagsData(payloadTagsData, path, value);
+        path.pop();
+        index++;
+      }
+    } else if (object instanceof Map) {
+      Map<?, ?> map = (Map<?, ?>) object;
+      for (Map.Entry<?, ?> entry : map.entrySet()) {
+        path.push(entry.getKey().toString());
+        collectPayloadTagsData(payloadTagsData, path, entry.getValue());
+        path.pop();
+      }
+    } else if (object instanceof SdkBytes) {
+      SdkBytes bytes = (SdkBytes) object;
+      payloadTagsData.add(new PayloadTagsData.PathAndValue(path.toArray(), bytes.asInputStream()));
+    } else {
+      payloadTagsData.add(new PayloadTagsData.PathAndValue(path.toArray(), object));
+    }
+  }
 }