Skip to content

Commit 90483bb

Browse files
ygreeygree
authored
AWS Payload Tag Extraction (#7811)
AWS Payload Tag Extraction - Support for AWS SDK Java 2 only - Disabled by default - Once enabled, only extracts tags for AWS ApiGateway, EventBridge, Sqs, Sns, S3, Kinesis, unless explicitly enabled for other services - Extracts tags for all JSON data it encounters - Provides some general and service-specific redaction rules - Support for user-defined redaction rules
1 parent 3e45e23 commit 90483bb

File tree

20 files changed

+2518
-2
lines changed

20 files changed

+2518
-2
lines changed

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ addTestSuiteExtendingForDir('dsmForkedTest', 'dsmTest', 'dsmTest')
1919
addTestSuiteForDir('latestDsmTest', 'dsmTest')
2020
addTestSuiteExtendingForDir('latestDsmForkedTest', 'latestDsmTest', 'dsmTest')
2121

22+
addTestSuite('payloadTaggingTest')
23+
addTestSuiteExtendingForDir('payloadTaggingForkedTest', 'payloadTaggingTest', 'payloadTaggingTest')
24+
addTestSuiteForDir('latestPayloadTaggingTest', 'payloadTaggingTest')
25+
addTestSuiteExtendingForDir('latestPayloadTaggingForkedTest', 'latestPayloadTaggingTest', 'payloadTaggingTest')
26+
2227
def fixedSdkVersion = '2.20.33' // 2.20.34 is missing and breaks IDEA import
2328

2429
dependencies {
@@ -40,6 +45,8 @@ dependencies {
4045
testImplementation group: 'org.eclipse.jetty.http2', name: 'http2-server', version: '9.4.56.v20240826'
4146

4247

48+
testImplementation group: 'org.testcontainers', name: 'localstack', version: libs.versions.testcontainers.get()
49+
4350
// First version where dsm traced operations have required StreamARN parameter for kinesis
4451
// and publishBatch is available for SNS
4552
dsmTestImplementation group: 'software.amazon.awssdk', name: 'apache-client', version: '2.18.40'
@@ -49,10 +56,25 @@ dependencies {
4956
latestDsmTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: '+'
5057
latestDsmTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '+'
5158

59+
payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'apigateway', version: '2.19.0'
60+
payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: '2.7.4'
61+
payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sqs', version: '2.18.40'
62+
payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '2.18.40'
63+
payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 's3', version: '2.18.40'
64+
payloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: '2.18.40'
65+
latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'apigateway', version: '2.25.40'
66+
latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: '2.25.40'
67+
latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sqs', version: '2.25.40'
68+
latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '2.25.40'
69+
latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 's3', version: '2.18.40'
70+
latestPayloadTaggingTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: '2.18.40'
71+
5272
latestDepTestImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
5373
latestDepTestImplementation project(':dd-java-agent:instrumentation:netty-4.1')
5474

5575
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'apache-client', version: fixedSdkVersion
76+
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'apigateway', version: fixedSdkVersion
77+
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: fixedSdkVersion
5678
latestDepTestImplementation group: 'software.amazon.awssdk', name: 's3', version: fixedSdkVersion
5779
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'rds', version: fixedSdkVersion
5880
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'ec2', version: fixedSdkVersion
@@ -61,3 +83,7 @@ dependencies {
6183
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'dynamodb', version: fixedSdkVersion
6284
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'kinesis', version: fixedSdkVersion
6385
}
86+
87+
tasks.withType(Test).configureEach {
88+
usesService(testcontainersLimit)
89+
}

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
88

99
import datadog.trace.api.Config;
10+
import datadog.trace.api.ConfigDefaults;
1011
import datadog.trace.api.DDTags;
1112
import datadog.trace.api.cache.DDCache;
1213
import datadog.trace.api.cache.DDCaches;
@@ -23,16 +24,22 @@
2324
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
2425
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
2526
import datadog.trace.core.datastreams.TagsProcessor;
27+
import datadog.trace.payloadtags.PayloadTagsData;
2628
import java.net.URI;
2729
import java.time.Instant;
30+
import java.util.ArrayDeque;
31+
import java.util.ArrayList;
32+
import java.util.Collection;
2833
import java.util.Collections;
2934
import java.util.HashSet;
3035
import java.util.LinkedHashMap;
3136
import java.util.List;
37+
import java.util.Map;
3238
import java.util.Optional;
3339
import java.util.Set;
3440
import javax.annotation.Nonnull;
3541
import software.amazon.awssdk.awscore.AwsResponse;
42+
import software.amazon.awssdk.core.SdkBytes;
3643
import software.amazon.awssdk.core.SdkField;
3744
import software.amazon.awssdk.core.SdkPojo;
3845
import software.amazon.awssdk.core.SdkRequest;
@@ -115,6 +122,12 @@ public AgentSpan onSdkRequest(
115122
final String awsOperationName = attributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
116123
onOperation(span, awsServiceName, awsOperationName);
117124

125+
Config config = Config.get();
126+
if (config.isCloudRequestPayloadTaggingEnabled()
127+
&& config.isCloudPayloadTaggingEnabledFor(awsServiceName)) {
128+
awsPojoToTags(span, ConfigDefaults.DEFAULT_TRACE_CLOUD_PAYLOAD_REQUEST_TAG, request);
129+
}
130+
118131
// S3
119132
request.getValueForField("Bucket", String.class).ifPresent(name -> setBucketName(span, name));
120133
if ("s3".equalsIgnoreCase(awsServiceName) && span.traceConfig().isDataStreamsEnabled()) {
@@ -295,6 +308,14 @@ public AgentSpan onSdkResponse(
295308
final SdkResponse response,
296309
final SdkHttpResponse httpResponse,
297310
final ExecutionAttributes attributes) {
311+
312+
Config config = Config.get();
313+
String serviceName = attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
314+
if (config.isCloudResponsePayloadTaggingEnabled()
315+
&& config.isCloudPayloadTaggingEnabledFor(serviceName)) {
316+
awsPojoToTags(span, ConfigDefaults.DEFAULT_TRACE_CLOUD_PAYLOAD_RESPONSE_TAG, response);
317+
}
318+
298319
if (response instanceof AwsResponse) {
299320
span.setTag(
300321
InstrumentationTags.AWS_REQUEST_ID,
@@ -437,4 +458,48 @@ protected String getRequestHeader(SdkHttpRequest request, String headerName) {
437458
protected String getResponseHeader(SdkHttpResponse response, String headerName) {
438459
return response.firstMatchingHeader(headerName).orElse(null);
439460
}
461+
462+
private void awsPojoToTags(AgentSpan span, String tagsPrefix, Object pojo) {
463+
Collection<PayloadTagsData.PathAndValue> payloadTagsData = new ArrayList<>();
464+
ArrayDeque<Object> path = new ArrayDeque<>();
465+
collectPayloadTagsData(payloadTagsData, path, pojo);
466+
span.setTag(
467+
tagsPrefix,
468+
new PayloadTagsData(payloadTagsData.toArray(new PayloadTagsData.PathAndValue[0])));
469+
}
470+
471+
private void collectPayloadTagsData(
472+
Collection<PayloadTagsData.PathAndValue> payloadTagsData,
473+
ArrayDeque<Object> path,
474+
Object object) {
475+
if (object instanceof SdkPojo) {
476+
SdkPojo pojo = (SdkPojo) object;
477+
for (SdkField<?> field : pojo.sdkFields()) {
478+
Object val = field.getValueOrDefault(pojo);
479+
path.push(field.locationName());
480+
collectPayloadTagsData(payloadTagsData, path, val);
481+
path.pop();
482+
}
483+
} else if (object instanceof Collection) {
484+
int index = 0;
485+
for (Object value : (Collection<?>) object) {
486+
path.push(index);
487+
collectPayloadTagsData(payloadTagsData, path, value);
488+
path.pop();
489+
index++;
490+
}
491+
} else if (object instanceof Map) {
492+
Map<?, ?> map = (Map<?, ?>) object;
493+
for (Map.Entry<?, ?> entry : map.entrySet()) {
494+
path.push(entry.getKey().toString());
495+
collectPayloadTagsData(payloadTagsData, path, entry.getValue());
496+
path.pop();
497+
}
498+
} else if (object instanceof SdkBytes) {
499+
SdkBytes bytes = (SdkBytes) object;
500+
payloadTagsData.add(new PayloadTagsData.PathAndValue(path.toArray(), bytes.asInputStream()));
501+
} else {
502+
payloadTagsData.add(new PayloadTagsData.PathAndValue(path.toArray(), object));
503+
}
504+
}
440505
}

0 commit comments

Comments
 (0)