From ac1a522c50b5e06666d15ebe5b45e3c6db357a52 Mon Sep 17 00:00:00 2001 From: "eunhye.kim" Date: Fri, 4 Sep 2026 14:37:39 +0900 Subject: [PATCH 1/6] feat(lambda): strip _datadog trace context injected into EventBridge event detail --- .../lambda/LambdaHandlerInstrumentation.java | 13 +- .../LambdaHandlerInstrumentationTest.java | 44 ++++ .../datadog/trace/api/ConfigDefaults.java | 2 + .../trace/api/config/GeneralConfig.java | 2 + .../lambda/StripInjectedContextBenchmark.java | 48 ++++ .../java/datadog/trace/core/CoreTracer.java | 7 + .../trace/lambda/StripInjectedContext.java | 239 ++++++++++++++++++ .../lambda/StripInjectedContextFuzzTest.java | 101 ++++++++ .../lambda/StripInjectedContextTest.java | 204 +++++++++++++++ .../main/java/datadog/trace/api/Config.java | 13 + .../instrumentation/api/AgentTracer.java | 8 + .../java/datadog/trace/api/ConfigTest.java | 18 ++ metadata/supported-configurations.json | 8 + 13 files changed, 706 insertions(+), 1 deletion(-) create mode 100644 dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java create mode 100644 dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java create mode 100644 dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index 7f62f4337d9..5c9f8c48793 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -20,6 +20,7 @@ import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.Config; import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -27,6 +28,7 @@ import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities; import datadog.trace.config.inversion.ConfigHelper; +import java.io.InputStream; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -76,7 +78,7 @@ public static class ExtensionCommunicationAdvice { @OnMethodEnter static ContextScope enter( @This final Object that, - @Advice.Argument(0) final Object in, + @Advice.Argument(value = 0, readOnly = false) InputStream in, @Advice.Argument(1) final Object out, @Advice.Argument(2) final Context awsContext, @Origin("#m") final String methodName) { @@ -86,6 +88,15 @@ static ContextScope enter( } String lambdaRequestId = awsContext.getAwsRequestId(); AgentSpanContext lambdaContext = AgentTracer.get().notifyLambdaStart(in, lambdaRequestId); + + // Skip the strip pass entirely (extra read/parse/re-serialize) unless explicitly enabled. + if (Config.get().isLambdaStripInjectedContextEnabled()) { + InputStream stripped = AgentTracer.get().stripLambdaInjectedContext(in); + if (stripped != null) { + in = stripped; + } + } + final AgentSpan span; if (null == lambdaContext) { span = startSpan("java-aws-sdk", INVOCATION_SPAN_NAME); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java index 0aff2114e51..d46198725be 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java @@ -15,9 +15,11 @@ import com.amazonaws.services.lambda.runtime.CognitoIdentity; import com.amazonaws.services.lambda.runtime.Context; import com.amazonaws.services.lambda.runtime.LambdaLogger; +import com.amazonaws.services.lambda.runtime.RequestStreamHandler; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.api.DDSpanTypes; import datadog.trace.api.DDTags; +import datadog.trace.api.config.GeneralConfig; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; import datadog.trace.api.gateway.Flow; @@ -34,6 +36,8 @@ import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; @@ -634,6 +638,46 @@ void responseCallbacksReceiveNoDataWhenHandlerThrows() { error(Error.class, "Some error")))); } + @WithConfig(key = GeneralConfig.LAMBDA_STRIP_INJECTED_CONTEXT, value = "true") + @Test + void stripsInjectedDatadogContextFromEventBridgeDetail() throws IOException { + String eventJson = + "{" + + "\"detail-type\": \"order.created\"," + + "\"detail\": {" + + " \"orderId\": 42," + + " \"_datadog\": {\"x-datadog-trace-id\": \"123\"}" + + "}" + + "}"; + + ByteArrayInputStream input = + new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8)); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + class CapturingHandler implements RequestStreamHandler { + String capturedBody; + + @Override + public void handleRequest(InputStream in, OutputStream out, Context context) + throws IOException { + ByteArrayOutputStream buffer = new ByteArrayOutputStream(); + byte[] chunk = new byte[8192]; + int n; + while ((n = in.read(chunk)) != -1) { + buffer.write(chunk, 0, n); + } + capturedBody = new String(buffer.toByteArray(), StandardCharsets.UTF_8); + out.write("{}".getBytes(StandardCharsets.UTF_8)); + } + } + + CapturingHandler handler = new CapturingHandler(); + handler.handleRequest(input, output, newContext()); + + assertFalse(handler.capturedBody.contains("_datadog")); + assertTrue(handler.capturedBody.contains("\"orderId\":42")); + } + private static final class TestContext implements Context { private final String requestId; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index bf5bab30d05..b1971cbd1db 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java @@ -342,6 +342,8 @@ public final class ConfigDefaults { static final boolean DEFAULT_WEBSOCKET_MESSAGES_SEPARATE_TRACES = true; static final boolean DEFAULT_WEBSOCKET_TAG_SESSION_ID = false; + static final boolean DEFAULT_LAMBDA_STRIP_INJECTED_CONTEXT = false; + static final Set DEFAULT_TRACE_CLOUD_PAYLOAD_TAGGING_SERVICES = new HashSet<>( Arrays.asList( diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java index a32f997cebe..a17dd65e580 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java @@ -140,5 +140,7 @@ public final class GeneralConfig { public static final String INSTRUMENTATION_SOURCE = "instrumentation.source"; public static final String APP_LOGS_COLLECTION_ENABLED = "app.logs.collection.enabled"; + public static final String LAMBDA_STRIP_INJECTED_CONTEXT = "lambda.strip.injected.context"; + private GeneralConfig() {} } diff --git a/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java b/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java new file mode 100644 index 00000000000..f0f570319b0 --- /dev/null +++ b/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java @@ -0,0 +1,48 @@ +package datadog.trace.lambda; + +import static java.util.concurrent.TimeUnit.MICROSECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; + +import java.nio.charset.StandardCharsets; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@State(Scope.Benchmark) +@Warmup(iterations = 3, time = 5, timeUnit = SECONDS) +@Measurement(iterations = 3, time = 5, timeUnit = SECONDS) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(MICROSECONDS) +@Fork(value = 1) +public class StripInjectedContextBenchmark { + + @Param({"true", "false"}) + boolean containsDatadogCarrier; + + byte[] payload; + + @Setup(Level.Trial) + public void setUp() { + String detail = + containsDatadogCarrier + ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":{\"x-datadog-trace-id\":\"123\",\"x-datadog-parent-id\":\"456\"}}" + : "{\"orderId\":42,\"customer\":\"acme\"}"; + String envelope = "{\"detail-type\":\"order.created\",\"detail\":" + detail + "}"; + payload = envelope.getBytes(StandardCharsets.UTF_8); + } + + @Benchmark + public void stripInternal(Blackhole blackhole) { + blackhole.consume(StripInjectedContext.stripInternal(payload)); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 9ed3f34aceb..22d6f60e4ef 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -119,8 +119,10 @@ import datadog.trace.core.traceinterceptor.LatencyTraceInterceptor; import datadog.trace.lambda.LambdaAppSecHandler; import datadog.trace.lambda.LambdaHandler; +import datadog.trace.lambda.StripInjectedContext; import datadog.trace.util.AgentTaskScheduler; import java.io.IOException; +import java.io.InputStream; import java.lang.ref.WeakReference; import java.util.ArrayList; import java.util.Collection; @@ -1366,6 +1368,11 @@ public void notifyAppSecEnd(AgentSpan span, Object result) { } } + @Override + public InputStream stripLambdaInjectedContext(InputStream in) { + return StripInjectedContext.replaceInputStream(in); + } + @Override public AgentDataStreamsMonitoring getDataStreamsMonitoring() { return dataStreamsMonitoring; diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java b/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java new file mode 100644 index 00000000000..6b9097dbcc8 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java @@ -0,0 +1,239 @@ +package datadog.trace.lambda; + +import com.squareup.moshi.JsonAdapter; +import com.squareup.moshi.JsonReader; +import com.squareup.moshi.JsonWriter; +import com.squareup.moshi.Moshi; +import com.squareup.moshi.Types; +import datadog.trace.api.Config; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.lang.reflect.Type; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Strips the injected "_datadog" trace carrier from a Lambda event's "detail" field before the user + * handler runs, when {@code DD_LAMBDA_STRIP_INJECTED_CONTEXT} is enabled. Fail-open: any error or + * invalid result returns the original payload unchanged. + */ +public final class StripInjectedContext { + + private static final Logger log = LoggerFactory.getLogger(StripInjectedContext.class); + + private static final String DETAIL_FIELD = "detail"; + private static final String DATADOG_CARRIER_KEY = "_datadog"; + + // Custom Object adapter preserves Long (not Double) for whole numbers and key order via + // LinkedHashMap, so re-serializing the envelope doesn't corrupt untouched fields. + private static final Moshi MOSHI = + new Moshi.Builder().add(Object.class, new NumberPreservingObjectAdapter()).build(); + + private static final Type STRING_OBJECT_MAP_TYPE = + Types.newParameterizedType(Map.class, String.class, Object.class); + + @SuppressWarnings("unchecked") + private static final JsonAdapter> MAP_ADAPTER = + (JsonAdapter>) (JsonAdapter) MOSHI.adapter(STRING_OBJECT_MAP_TYPE); + + private StripInjectedContext() {} + + /** Reads {@code in} fully, strips "_datadog", and returns a new stream over the result. */ + public static ByteArrayInputStream replaceInputStream(InputStream in) { + if (in == null) { + return null; + } + try { + return new ByteArrayInputStream(strip(readAllBytes(in))); + } catch (IOException e) { + log.debug("Unable to read Lambda payload for _datadog stripping", e); + return null; + } + } + + /** Strips "_datadog" from {@code payload}'s "detail" field when the config is enabled. */ + public static byte[] strip(byte[] payload) { + if (!Config.get().isLambdaStripInjectedContextEnabled()) { + return payload; + } + return stripInternal(payload); + } + + /** Unconditional strip logic, package-private so tests can call it directly. */ + static byte[] stripInternal(byte[] payload) { + if (payload == null || payload.length == 0) { + return payload; + } + + // Fast path: skip JSON parsing when there's clearly nothing to strip. + String json = new String(payload, StandardCharsets.UTF_8); + if (!json.contains(DATADOG_CARRIER_KEY)) { + return payload; + } + + try { + Map envelope = MAP_ADAPTER.fromJson(json); + if (envelope == null || !envelope.containsKey(DETAIL_FIELD)) { + return payload; + } + + Object strippedDetail = stripDetail(envelope.get(DETAIL_FIELD)); + if (strippedDetail == null) { + return payload; // no "_datadog" found in detail + } + + envelope.put(DETAIL_FIELD, strippedDetail); + return MAP_ADAPTER.toJson(envelope).getBytes(StandardCharsets.UTF_8); + } catch (IOException | RuntimeException e) { + // Fail open: stripping must never break the handler invocation. + log.debug("Unable to strip injected _datadog context from Lambda payload", e); + return payload; + } + } + + /** Removes "_datadog" from an object or string-encoded detail value. Null if unchanged. */ + @SuppressWarnings("unchecked") + private static Object stripDetail(Object detail) { + // Object detail: "detail": {"foo": "bar", "_datadog": {...}} + if (detail instanceof Map) { + Map detailMap = (Map) detail; + return detailMap.remove(DATADOG_CARRIER_KEY) != null ? detailMap : null; + } + // String-encoded detail: "detail": "{\"foo\":\"bar\",\"_datadog\":{...}}" + if (detail instanceof String) { + return stripFromJsonString((String) detail); + } + return null; + } + + /** Parses, strips "_datadog", and re-encodes a JSON-object-shaped string. Null if unchanged. */ + private static String stripFromJsonString(String value) { + if (!value.contains(DATADOG_CARRIER_KEY)) { + return null; + } + try { + Map inner = MAP_ADAPTER.fromJson(value); + if (inner == null || inner.remove(DATADOG_CARRIER_KEY) == null) { + return null; + } + return MAP_ADAPTER.toJson(inner); + } catch (IOException | RuntimeException e) { + return null; + } + } + + /** Reads all bytes from {@code in}, restoring its position first if it's a byte stream. */ + private static byte[] readAllBytes(InputStream in) throws IOException { + if (in instanceof ByteArrayInputStream) { + // mark/reset instead of draining, in case the stream is read again elsewhere. + ByteArrayInputStream bais = (ByteArrayInputStream) in; + bais.mark(Integer.MAX_VALUE); + byte[] bytes = new byte[bais.available()]; + int read = bais.read(bytes); + bais.reset(); + return (read == bytes.length) ? bytes : Arrays.copyOf(bytes, Math.max(read, 0)); + } + ByteArrayOutputStream buffer = new ByteArrayOutputStream(); + byte[] chunk = new byte[8192]; + int n; + while ((n = in.read(chunk)) != -1) { + buffer.write(chunk, 0, n); + } + return buffer.toByteArray(); + } + + /** Object adapter that keeps whole numbers as {@link Long} and preserves key order. */ + private static final class NumberPreservingObjectAdapter extends JsonAdapter { + + /** Recursively reads a JSON value, mapping numbers to Long/Double based on their form. */ + @Override + public Object fromJson(JsonReader reader) throws IOException { + switch (reader.peek()) { + case BEGIN_ARRAY: + List list = new ArrayList<>(); + reader.beginArray(); + while (reader.hasNext()) { + list.add(fromJson(reader)); + } + reader.endArray(); + return list; + + case BEGIN_OBJECT: + Map map = new LinkedHashMap<>(); + reader.beginObject(); + while (reader.hasNext()) { + map.put(reader.nextName(), fromJson(reader)); + } + reader.endObject(); + return map; + + case STRING: + return reader.nextString(); + + case NUMBER: + return readNumber(reader); + + case BOOLEAN: + return reader.nextBoolean(); + + case NULL: + return reader.nextNull(); + + default: + throw new IOException("Unexpected JSON token: " + reader.peek()); + } + } + + /** Returns Long for integer literals, Double for fractional/scientific ones. */ + private Object readNumber(JsonReader reader) throws IOException { + String literal = reader.nextString(); // raw text avoids precision loss before parsing + if (literal.indexOf('.') >= 0 || literal.indexOf('e') >= 0 || literal.indexOf('E') >= 0) { + return Double.parseDouble(literal); + } + try { + return Long.parseLong(literal); + } catch (NumberFormatException tooLargeForLong) { + return Double.parseDouble(literal); + } + } + + /** Recursively writes a value produced by {@link #fromJson}. */ + @Override + public void toJson(JsonWriter writer, Object value) throws IOException { + if (value instanceof Map) { + writer.beginObject(); + for (Map.Entry entry : ((Map) value).entrySet()) { + writer.name(String.valueOf(entry.getKey())); + toJson(writer, entry.getValue()); + } + writer.endObject(); + } else if (value instanceof List) { + writer.beginArray(); + for (Object item : (List) value) { + toJson(writer, item); + } + writer.endArray(); + } else if (value instanceof String) { + writer.value((String) value); + } else if (value instanceof Long) { + writer.value(((Long) value).longValue()); + } else if (value instanceof Number) { + writer.value((Number) value); + } else if (value instanceof Boolean) { + writer.value((Boolean) value); + } else if (value == null) { + writer.nullValue(); + } else { + throw new IOException("Unsupported value type: " + value.getClass()); + } + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java new file mode 100644 index 00000000000..4a8629ce54b --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java @@ -0,0 +1,101 @@ +package datadog.trace.lambda; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.concurrent.ThreadLocalRandom; +import org.junit.jupiter.api.Test; + +/** + * Randomized invariant testing for {@link StripInjectedContext}, mirroring the repo's existing + * fuzz-style test pattern (see {@code TagMapFuzzTest}). Generates many random EventBridge-shaped + * payloads instead of enumerating fixed cases by hand, and asserts core invariants hold for every + * one of them. + */ +public final class StripInjectedContextFuzzTest { + + private static final int ITERATIONS = 5_000; + + @Test + void neverThrowsAndAlwaysRemovesDatadogWhenPresent() { + ThreadLocalRandom random = ThreadLocalRandom.current(); + + for (int i = 0; i < ITERATIONS; i++) { + byte[] payload = randomPayload(random); + + byte[] result; + try { + result = StripInjectedContext.stripInternal(payload); + } catch (Throwable t) { + fail( + "stripInternal must never throw, but threw for payload: " + + new String(payload, StandardCharsets.UTF_8), + t); + return; + } + + String resultJson = new String(result, StandardCharsets.UTF_8); + assertFalse( + resultJson.contains("_datadog"), "output must never contain _datadog: " + resultJson); + + // Idempotency: stripping an already-stripped payload must be a no-op. + byte[] second = StripInjectedContext.stripInternal(result); + assertTrue(Arrays.equals(result, second), "stripInternal must be idempotent"); + } + } + + private static byte[] randomPayload(ThreadLocalRandom random) { + StringBuilder sb = new StringBuilder(); + sb.append("{\"detail-type\":\"order.created\",\"detail\":"); + boolean asString = random.nextBoolean(); + boolean includeDatadog = random.nextBoolean(); + + String detailJson = randomDetailObject(random, includeDatadog); + if (asString) { + sb.append('"').append(detailJson.replace("\"", "\\\"")).append('"'); + } else { + sb.append(detailJson); + } + sb.append('}'); + return sb.toString().getBytes(StandardCharsets.UTF_8); + } + + private static String randomDetailObject(ThreadLocalRandom random, boolean includeDatadog) { + StringBuilder sb = new StringBuilder("{"); + int fieldCount = random.nextInt(1, 6); + for (int i = 0; i < fieldCount; i++) { + if (i > 0) { + sb.append(','); + } + sb.append('"').append("field").append(i).append("\":"); + appendRandomValue(sb, random); + } + if (includeDatadog) { + if (fieldCount > 0) { + sb.append(','); + } + sb.append("\"_datadog\":{\"x-datadog-trace-id\":\"").append(random.nextLong()).append("\"}"); + } + sb.append('}'); + return sb.toString(); + } + + private static void appendRandomValue(StringBuilder sb, ThreadLocalRandom random) { + switch (random.nextInt(4)) { + case 0: + sb.append(random.nextInt()); + break; + case 1: + sb.append(random.nextDouble()); + break; + case 2: + sb.append('"').append("value").append(random.nextInt(1000)).append('"'); + break; + default: + sb.append(random.nextBoolean()); + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java new file mode 100644 index 00000000000..6cc374ee66b --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java @@ -0,0 +1,204 @@ +package datadog.trace.lambda; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.api.config.GeneralConfig; +import datadog.trace.test.junit.utils.config.WithConfig; +import datadog.trace.test.util.DDJavaSpecification; +import java.io.ByteArrayInputStream; +import java.nio.charset.StandardCharsets; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link StripInjectedContext}. */ +class StripInjectedContextTest extends DDJavaSpecification { + + private static byte[] bytes(String json) { + return json.getBytes(StandardCharsets.UTF_8); + } + + private static String toUtf8(byte[] bytes) { + return new String(bytes, StandardCharsets.UTF_8); + } + + private static byte[] readAll(ByteArrayInputStream in) { + in.mark(Integer.MAX_VALUE); + byte[] bytes = new byte[in.available()]; + int read = in.read(bytes, 0, bytes.length); + in.reset(); + return read == bytes.length ? bytes : java.util.Arrays.copyOf(bytes, Math.max(read, 0)); + } + + // ============================================================================ + // stripInternal: detail as an object + // ============================================================================ + + @Test + void removesDatadogKeyFromObjectDetail() { + byte[] input = + bytes( + "{\"detail-type\":\"order.created\",\"detail\":{\"orderId\":42," + + "\"_datadog\":{\"x-datadog-trace-id\":\"123\"}}}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(json.contains("\"orderId\":42"), "unrelated fields must survive untouched"); + assertTrue(!json.contains("_datadog"), "the _datadog carrier must be removed"); + } + + // ============================================================================ + // stripInternal: detail as a string-encoded JSON object + // ============================================================================ + + @Test + void removesDatadogKeyFromStringEncodedDetail() { + // "detail" itself is a JSON string containing an escaped JSON object + byte[] input = + bytes( + "{\"detail\":\"{\\\"orderId\\\":42,\\\"_datadog\\\":{\\\"x-datadog-trace-id\\\":\\\"123\\\"}}\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog")); + assertTrue(json.contains("orderId")); + } + + // ============================================================================ + // Fail-open cases: original payload must be returned unchanged + // ============================================================================ + + @Test + void returnsOriginalWhenDatadogKeyIsAbsent() { + byte[] input = bytes("{\"detail\":{\"orderId\":42}}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + @Test + void returnsOriginalWhenDetailFieldIsMissing() { + byte[] input = bytes("{\"detail-type\":\"order.created\",\"_datadog\":{\"trace\":\"1\"}}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + @Test + void returnsOriginalOnMalformedJson() { + byte[] input = bytes("{not valid json, but contains _datadog anyway"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + @Test + void returnsOriginalForNullOrEmptyPayload() { + assertNull(StripInjectedContext.stripInternal(null)); + assertArrayEquals(new byte[0], StripInjectedContext.stripInternal(new byte[0])); + } + + @Test + void returnsOriginalWhenDetailIsNeitherObjectNorString() { + // e.g. detail is a number/array/null — nothing sane to strip from + byte[] input = bytes("{\"detail\":42,\"_datadog\":{\"trace\":\"1\"}}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + // ============================================================================ + // Number type preservation (the whole reason for the custom Moshi adapter) + // ============================================================================ + + @Test + void preservesWholeNumbersAsIntegersAfterStripping() { + byte[] input = + bytes("{\"detail\":{\"orderId\":42,\"quantity\":7," + "\"_datadog\":{\"trace\":\"1\"}}}"); + + String json = toUtf8(StripInjectedContext.stripInternal(input)); + + // Without the custom adapter, Moshi's default would turn these into "42.0"/"7.0" + assertTrue(json.contains("\"orderId\":42")); + assertTrue(json.contains("\"quantity\":7")); + assertTrue(!json.contains(".0")); + } + + @Test + void preservesFractionalNumbersAsDoublesAfterStripping() { + byte[] input = bytes("{\"detail\":{\"price\":19.99,\"_datadog\":{\"trace\":\"1\"}}}"); + + String json = toUtf8(StripInjectedContext.stripInternal(input)); + + assertTrue(json.contains("\"price\":19.99")); + } + + @Test + void preservesLargeIntegersBeyondIntRangeAsLong() { + byte[] input = + bytes("{\"detail\":{\"bigId\":9007199254740993,\"_datadog\":{\"trace\":\"1\"}}}"); + + String json = toUtf8(StripInjectedContext.stripInternal(input)); + + // A double could not represent this value exactly; Long must be used instead. + assertTrue(json.contains("\"bigId\":9007199254740993")); + } + + // ============================================================================ + // strip(): the public entry point that is gated by config + // ============================================================================ + + @Test + void stripReturnsOriginalPayloadWhenConfigIsDisabledByDefault() { + byte[] input = bytes("{\"detail\":{\"_datadog\":{\"trace\":\"1\"}}}"); + + byte[] result = StripInjectedContext.strip(input); + + assertArrayEquals(input, result); + } + + @Test + @WithConfig(key = GeneralConfig.LAMBDA_STRIP_INJECTED_CONTEXT, value = "true") + void stripRemovesDatadogKeyWhenConfigIsEnabled() { + byte[] input = bytes("{\"detail\":{\"orderId\":1,\"_datadog\":{\"trace\":\"1\"}}}"); + + byte[] result = StripInjectedContext.strip(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog")); + } + + // ============================================================================ + // replaceInputStream(): the InputStream-based helper used by the advice + // ============================================================================ + + @Test + void replaceInputStreamReturnsNullForNullInput() { + assertNull(StripInjectedContext.replaceInputStream(null)); + } + + @Test + @WithConfig(key = GeneralConfig.LAMBDA_STRIP_INJECTED_CONTEXT, value = "true") + void replaceInputStreamStripsAndDoesNotConsumeTheOriginalStream() { + String eventJson = "{\"detail\":{\"orderId\":1,\"_datadog\":{\"trace\":\"1\"}}}"; + ByteArrayInputStream original = + new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8)); + + ByteArrayInputStream stripped = StripInjectedContext.replaceInputStream(original); + + assertNotNull(stripped); + String strippedJson = new String(readAll(stripped), StandardCharsets.UTF_8); + assertTrue(!strippedJson.contains("_datadog")); + + // readAllBytes() inside replaceInputStream() must mark/reset, not drain, the original stream + String originalJson = new String(readAll(original), StandardCharsets.UTF_8); + assertTrue(originalJson.contains("_datadog"), "original stream must still be fully readable"); + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index 3471750c347..39a72172d75 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -118,6 +118,7 @@ import static datadog.trace.api.ConfigDefaults.DEFAULT_JMX_FETCH_ENABLED; import static datadog.trace.api.ConfigDefaults.DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_ENABLED; import static datadog.trace.api.ConfigDefaults.DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_LIMIT; +import static datadog.trace.api.ConfigDefaults.DEFAULT_LAMBDA_STRIP_INJECTED_CONTEXT; import static datadog.trace.api.ConfigDefaults.DEFAULT_LLM_OBS_AGENTLESS_ENABLED; import static datadog.trace.api.ConfigDefaults.DEFAULT_LLM_OBS_SAMPLE_RATE; import static datadog.trace.api.ConfigDefaults.DEFAULT_LOGS_INJECTION_ENABLED; @@ -406,6 +407,7 @@ import static datadog.trace.api.config.GeneralConfig.HEALTH_METRICS_STATSD_PORT; import static datadog.trace.api.config.GeneralConfig.INSTRUMENTATION_SOURCE; import static datadog.trace.api.config.GeneralConfig.JDK_SOCKET_ENABLED; +import static datadog.trace.api.config.GeneralConfig.LAMBDA_STRIP_INJECTED_CONTEXT; import static datadog.trace.api.config.GeneralConfig.LOG_LEVEL; import static datadog.trace.api.config.GeneralConfig.PERF_METRICS_ENABLED; import static datadog.trace.api.config.GeneralConfig.PRIMARY_TAG; @@ -1415,6 +1417,7 @@ public static String getHostName() { private final String dogStatsDPath; private final List dogStatsDArgs; private final int dogStatsDPort; + private final boolean lambdaStripInjectedContextEnabled; private String env; private String version; @@ -1777,6 +1780,10 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins sqsInjectDatadogAttributeEnabled = isInjectDatadogAttributeEnabled(DEFAULT_INJECT_DATADOG_ATTRIBUTE, "sqs"); + lambdaStripInjectedContextEnabled = + configProvider.getBoolean( + LAMBDA_STRIP_INJECTED_CONTEXT, DEFAULT_LAMBDA_STRIP_INJECTED_CONTEXT); + spanAttributeSchemaVersion = schemaVersionFromConfig(); peerHostNameEnabled = configProvider.getBoolean(TRACE_PEER_HOSTNAME_ENABLED, true); @@ -5462,6 +5469,10 @@ public boolean isSfnInjectDatadogAttributeEnabled() { return sfnInjectDatadogAttributeEnabled; } + public boolean isLambdaStripInjectedContextEnabled() { + return lambdaStripInjectedContextEnabled; + } + /** * @return A map of tags to be applied only to the local application root span. */ @@ -7209,6 +7220,8 @@ public String toString() { + sqsInjectDatadogAttributeEnabled + ", snsInjectDatadogAttributeEnabled=" + snsInjectDatadogAttributeEnabled + + ", lambdaStripInjectedContextEnabled=" + + lambdaStripInjectedContextEnabled + '}'; } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 540a4a2481a..ddf02ab64f3 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -29,6 +29,7 @@ import datadog.trace.context.NoopTraceScope; import datadog.trace.context.TraceScope; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.io.InputStream; import java.util.Collections; import java.util.List; import java.util.Map; @@ -384,6 +385,8 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera void notifyAppSecEnd(AgentSpan span, Object result); + InputStream stripLambdaInjectedContext(InputStream in); + AgentDataStreamsMonitoring getDataStreamsMonitoring(); String getTraceId(AgentSpan span); @@ -667,6 +670,11 @@ public void notifyExtensionEnd( @Override public void notifyAppSecEnd(AgentSpan span, Object result) {} + @Override + public InputStream stripLambdaInjectedContext(InputStream in) { + return null; + } + @Override public AgentDataStreamsMonitoring getDataStreamsMonitoring() { return NoopDataStreamsMonitoring.INSTANCE; diff --git a/internal-api/src/test/java/datadog/trace/api/ConfigTest.java b/internal-api/src/test/java/datadog/trace/api/ConfigTest.java index bbea22a73fa..777759f9e79 100644 --- a/internal-api/src/test/java/datadog/trace/api/ConfigTest.java +++ b/internal-api/src/test/java/datadog/trace/api/ConfigTest.java @@ -54,6 +54,7 @@ import static datadog.trace.api.config.GeneralConfig.HEALTH_METRICS_STATSD_PORT; import static datadog.trace.api.config.GeneralConfig.INSTRUMENTATION_SOURCE; import static datadog.trace.api.config.GeneralConfig.JDK_SOCKET_ENABLED; +import static datadog.trace.api.config.GeneralConfig.LAMBDA_STRIP_INJECTED_CONTEXT; import static datadog.trace.api.config.GeneralConfig.PERF_METRICS_ENABLED; import static datadog.trace.api.config.GeneralConfig.SERVICE_NAME; import static datadog.trace.api.config.GeneralConfig.SITE; @@ -3010,6 +3011,23 @@ void configInstantiationShouldNotFailIfLlmObsIsEnabledViaEnvVarAndMlAppIsEmpty() assertEquals("test-service", config.getLlmObsMlApp()); } + @Test + void lambdaStripInjectedContextIsDisabledByDefault() throws Exception { + Config config = newConfig(); + + assertFalse(config.isLambdaStripInjectedContextEnabled()); + } + + @Test + void lambdaStripInjectedContextCanBeEnabled() throws Exception { + Properties properties = new Properties(); + properties.setProperty(LAMBDA_STRIP_INJECTED_CONTEXT, "true"); + + Config config = newConfig(withPropertiesOverride(properties)); + + assertTrue(config.isLambdaStripInjectedContextEnabled()); + } + @Test void configInstantiationShouldNotFailIfLlmObsIsEnabledAgentlessDisabledViaSysPropAndMlAppIsSet() throws Exception { diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 45e12e44c85..e43bf780833 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -2356,6 +2356,14 @@ "aliases": [] } ], + "DD_LAMBDA_STRIP_INJECTED_CONTEXT": [ + { + "version": "A", + "type": "boolean", + "default": "false", + "aliases": [] + } + ], "DD_LLMOBS_AGENTLESS_ENABLED": [ { "version": "B", From 1e1de82a1e109490fe50f7670a8895907f4da2f2 Mon Sep 17 00:00:00 2001 From: "eunhye.kim" Date: Tue, 29 Sep 2026 17:06:25 +0900 Subject: [PATCH 2/6] strip _datadog trace context injected for Eventbridge, SNS, SQS --- .../lambda/LambdaHandlerInstrumentation.java | 26 +- .../LambdaHandlerInstrumentationTest.java | 4 +- .../lambda/StripInjectedContextBenchmark.java | 39 ++- .../trace/lambda/StripInjectedContext.java | 305 ++++++++++++++++-- .../lambda/StripInjectedContextFuzzTest.java | 67 +++- .../lambda/StripInjectedContextTest.java | 173 +++++++++- 6 files changed, 560 insertions(+), 54 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index 5c9f8c48793..910a4c06e0d 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.aws.v1.lambda; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -13,6 +14,7 @@ import static net.bytebuddy.asm.Advice.This; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.amazonaws.services.lambda.runtime.Context; import com.amazonaws.services.lambda.runtime.RequestHandler; @@ -31,6 +33,7 @@ import java.io.InputStream; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.bytecode.assign.Assigner; import net.bytebuddy.matcher.ElementMatcher; @AutoService(InstrumenterModule.class) @@ -70,6 +73,10 @@ public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( isMethod() .and(named("handleRequest")) + .and(takesArguments(3)) + // Constrain arg0 to InputStream so an unrelated 3-arg "handleRequest" overload + // with a non-InputStream first parameter isn't matched by mistake. + .and(takesArgument(0, isSubTypeOf(InputStream.class))) .and(takesArgument(2, named("com.amazonaws.services.lambda.runtime.Context"))), getClass().getName() + "$ExtensionCommunicationAdvice"); } @@ -78,24 +85,31 @@ public static class ExtensionCommunicationAdvice { @OnMethodEnter static ContextScope enter( @This final Object that, - @Advice.Argument(value = 0, readOnly = false) InputStream in, + // Bound as Object (not InputStream) so Byte Buddy can always bind this argument for + // writing back, even if the matcher's arg0 constraint is ever loosened or bypassed. + @Advice.Argument(value = 0, readOnly = false, typing = Assigner.Typing.DYNAMIC) Object in, @Advice.Argument(1) final Object out, @Advice.Argument(2) final Context awsContext, @Origin("#m") final String methodName) { - if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) { return null; } + if (!(in instanceof InputStream)) { + // Defensive guard: should never happen given the matcher's arg0 constraint. + return null; + } + InputStream inputStream = (InputStream) in; String lambdaRequestId = awsContext.getAwsRequestId(); - AgentSpanContext lambdaContext = AgentTracer.get().notifyLambdaStart(in, lambdaRequestId); - + AgentSpanContext lambdaContext = + AgentTracer.get().notifyLambdaStart(inputStream, lambdaRequestId); // Skip the strip pass entirely (extra read/parse/re-serialize) unless explicitly enabled. if (Config.get().isLambdaStripInjectedContextEnabled()) { - InputStream stripped = AgentTracer.get().stripLambdaInjectedContext(in); + InputStream stripped = AgentTracer.get().stripLambdaInjectedContext(inputStream); if (stripped != null) { - in = stripped; + inputStream = stripped; } } + in = inputStream; final AgentSpan span; if (null == lambdaContext) { diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java index d46198725be..a5a1563d2d5 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java @@ -645,8 +645,8 @@ void stripsInjectedDatadogContextFromEventBridgeDetail() throws IOException { "{" + "\"detail-type\": \"order.created\"," + "\"detail\": {" - + " \"orderId\": 42," - + " \"_datadog\": {\"x-datadog-trace-id\": \"123\"}" + + " \"orderId\":42," + + " \"_datadog\":{\"x-datadog-trace-id\": \"123\"}" + "}" + "}"; diff --git a/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java b/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java index f0f570319b0..fb33f6ebb02 100644 --- a/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java +++ b/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java @@ -29,16 +29,45 @@ public class StripInjectedContextBenchmark { @Param({"true", "false"}) boolean containsDatadogCarrier; + // Event shapes: "eventbridge", "sqs", "sns" + @Param({"eventbridge", "sqs", "sns"}) + String payloadShape; + byte[] payload; @Setup(Level.Trial) public void setUp() { - String detail = - containsDatadogCarrier - ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":{\"x-datadog-trace-id\":\"123\",\"x-datadog-parent-id\":\"456\"}}" + String carrier = + "{\"x-datadog-trace-id\":\"123\",\"x-datadog-parent-id\":\"456\"}"; + + switch (payloadShape) { + case "sqs": { + // SQS-style: _datadog appears as a direct top-level key. + String body = containsDatadogCarrier + ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":" + carrier + "}" + : "{\"orderId\":42,\"customer\":\"acme\"}"; + payload = body.getBytes(StandardCharsets.UTF_8); + break; + } + case "sns": { + // SNS-style: _datadog is embedded inside a string-encoded "Message" field. + String inner = containsDatadogCarrier + ? "{\\\"orderId\\\":42,\\\"_datadog\\\":" + carrier.replace("\"", "\\\"") + "}" + : "{\\\"orderId\\\":42}"; + String envelope = "{\"Type\":\"Notification\",\"Message\":\"" + inner + "\"}"; + payload = envelope.getBytes(StandardCharsets.UTF_8); + break; + } + default: { + // EventBridge-style: _datadog is inside the "detail" object. + String detail = containsDatadogCarrier + ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":" + carrier + "}" : "{\"orderId\":42,\"customer\":\"acme\"}"; - String envelope = "{\"detail-type\":\"order.created\",\"detail\":" + detail + "}"; - payload = envelope.getBytes(StandardCharsets.UTF_8); + String envelope = "{\"detail-type\":\"order.created\",\"detail\":" + detail + "}"; + payload = envelope.getBytes(StandardCharsets.UTF_8); + break; + } + } } @Benchmark diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java b/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java index 6b9097dbcc8..a3d6bea7a24 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java @@ -11,6 +11,7 @@ import java.io.IOException; import java.io.InputStream; import java.lang.reflect.Type; +import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; @@ -21,17 +22,36 @@ import org.slf4j.LoggerFactory; /** - * Strips the injected "_datadog" trace carrier from a Lambda event's "detail" field before the user - * handler runs, when {@code DD_LAMBDA_STRIP_INJECTED_CONTEXT} is enabled. Fail-open: any error or - * invalid result returns the original payload unchanged. + * Strips the injected "_datadog" trace carrier from Lambda event payloads before the user + * handler runs, when {@code DD_LAMBDA_STRIP_INJECTED_CONTEXT} is enabled. Applies to any + * top-level field value, covering EventBridge, SNS, SQS, Kinesis, and similar event shapes. + * Fail-open: any error or invalid result returns the original payload unchanged. */ public final class StripInjectedContext { private static final Logger log = LoggerFactory.getLogger(StripInjectedContext.class); - private static final String DETAIL_FIELD = "detail"; private static final String DATADOG_CARRIER_KEY = "_datadog"; + /** Raw bytes of "_datadog" for the fast-path substring check. */ + private static final byte[] DATADOG_CARRIER_BYTES = + DATADOG_CARRIER_KEY.getBytes(StandardCharsets.UTF_8); + + /** + * Byte pattern {@code "_datadog":} used to locate object-form carriers in raw JSON bytes. + * The colon is included so a single scan step covers both the key and the separator, + * matching the approach used in the Go tracer implementation. + */ + private static final byte[] KEY_COLON_BYTES = + ("\"" + DATADOG_CARRIER_KEY + "\":").getBytes(StandardCharsets.UTF_8); + + /** + * Byte pattern {@code \"_datadog\"} used to detect carriers that are embedded inside + * string-encoded JSON fields (e.g. an SNS {@code Message} field). + */ + private static final byte[] ESCAPED_KEY_BYTES = + ("\\\"" + DATADOG_CARRIER_KEY + "\\\"").getBytes(StandardCharsets.UTF_8); + // Custom Object adapter preserves Long (not Double) for whole numbers and key order via // LinkedHashMap, so re-serializing the envelope doesn't corrupt untouched fields. private static final Moshi MOSHI = @@ -59,7 +79,7 @@ public static ByteArrayInputStream replaceInputStream(InputStream in) { } } - /** Strips "_datadog" from {@code payload}'s "detail" field when the config is enabled. */ + /** Strips "_datadog" from {@code payload} when the config is enabled. */ public static byte[] strip(byte[] payload) { if (!Config.get().isLambdaStripInjectedContextEnabled()) { return payload; @@ -72,46 +92,262 @@ static byte[] stripInternal(byte[] payload) { if (payload == null || payload.length == 0) { return payload; } - - // Fast path: skip JSON parsing when there's clearly nothing to strip. - String json = new String(payload, StandardCharsets.UTF_8); - if (!json.contains(DATADOG_CARRIER_KEY)) { + // Fast path: skip all processing when "_datadog" bytes are absent entirely. + if (indexOf(payload, 0, DATADOG_CARRIER_BYTES) < 0) { return payload; } + // Pass 1: splice out "_datadog":{...} object carriers via byte-level scanning. + // No JSON parsing or object allocation occurs here. + byte[] result = stripObjectCarriers(payload); + // Pass 2: strip "_datadog" from string-encoded fields (e.g. SNS Message). + // Uses Moshi only when the escaped key pattern is present. + result = stripStringEncodedCarriers(result); + return result; + } - try { - Map envelope = MAP_ADAPTER.fromJson(json); - if (envelope == null || !envelope.containsKey(DETAIL_FIELD)) { - return payload; + // ============================================================================ + // Object carrier stripping: byte-level splice (mirrors Go's stripObjectCarriers) + // ============================================================================ + + /** + * Removes all {@code "_datadog":{...}} object-form carriers from {@code msg} by splicing + * the matching byte ranges out of the array. Returns {@code msg} unchanged if no carriers + * are found. No JSON parsing or heap allocation occurs unless a carrier is actually present. + */ + private static byte[] stripObjectCarriers(byte[] msg) { + ByteArrayOutputStream out = null; + int searchFrom = 0; + int copyFrom = 0; + + while (searchFrom < msg.length) { + int[] carrier = findNextObjectCarrier(msg, searchFrom); + if (carrier == null) { + break; } + int keyStart = carrier[0]; + int valueEnd = carrier[1]; + int[] removal = carrierRemovalRange(msg, keyStart, valueEnd, searchFrom); - Object strippedDetail = stripDetail(envelope.get(DETAIL_FIELD)); - if (strippedDetail == null) { - return payload; // no "_datadog" found in detail + if (out == null) { + // Defer allocation until the first carrier is confirmed. Pre-size to the full + // input length so that multiple removals do not trigger internal buffer growth. + out = new ByteArrayOutputStream(msg.length); } + out.write(msg, copyFrom, removal[0] - copyFrom); + copyFrom = removal[1]; + searchFrom = removal[1]; + } + + if (out == null) { + return msg; + } + out.write(msg, copyFrom, msg.length - copyFrom); + return out.toByteArray(); + } - envelope.put(DETAIL_FIELD, strippedDetail); + /** + * Returns {@code [keyStart, valueEnd]} for the next removable {@code "_datadog":{...}} + * carrier found at or after {@code searchFrom}, or {@code null} if none exists. + * + *

A carrier is removable only when (1) it is an actual JSON object property, not text + * embedded inside a string value, and (2) its value is an object (starts with {@code {}). + */ + private static int[] findNextObjectCarrier(byte[] msg, int searchFrom) { + while (searchFrom < msg.length) { + int keyStart = indexOf(msg, searchFrom, KEY_COLON_BYTES); + if (keyStart < 0) { + return null; + } + if (!isObjectProperty(msg, keyStart)) { + searchFrom = keyStart + 1; + continue; + } + int valueStart = objectValueStart(msg, keyStart + KEY_COLON_BYTES.length); + if (valueStart < 0) { + // Value is not a JSON object (e.g. "_datadog": null); skip this occurrence. + searchFrom = keyStart + 1; + continue; + } + int valueEnd = scanObjectEnd(msg, valueStart); + if (valueEnd < 0) { + // Malformed JSON: opening '{' has no matching '}'; skip safely. + searchFrom = keyStart + 1; + continue; + } + return new int[]{keyStart, valueEnd}; + } + return null; + } + + /** + * Returns {@code true} when the byte immediately before {@code keyStart} (ignoring JSON + * whitespace) is {@code {} or {@code ,}, confirming that this is a real JSON property key + * and not text that happens to appear inside a string value. + */ + private static boolean isObjectProperty(byte[] msg, int keyStart) { + int previous = skipWhitespaceBackward(msg, keyStart - 1); + return previous >= 0 && (msg[previous] == '{' || msg[previous] == ','); + } + + /** + * Returns the index of the opening {@code {} of the carrier value (skipping JSON whitespace + * after the colon), or {@code -1} if the value does not start with {@code {}. + */ + private static int objectValueStart(byte[] msg, int start) { + start = skipWhitespaceForward(msg, start); + return (start < msg.length && msg[start] == '{') ? start : -1; + } + + /** + * Returns the {@code [start, end)} byte range to remove for the carrier whose key begins at + * {@code keyStart} and whose value ends at {@code valueEnd} (exclusive). The range is expanded + * to consume one adjacent comma so the surrounding JSON remains valid. The preceding comma is + * preferred unless it falls before {@code searchFrom} (which belongs to a range already removed + * in an earlier iteration). + */ + private static int[] carrierRemovalRange( + byte[] msg, int keyStart, int valueEnd, int searchFrom) { + int previous = skipWhitespaceBackward(msg, keyStart - 1); + if (previous >= searchFrom && msg[previous] == ',') { + return new int[]{previous, valueEnd}; + } + int next = skipWhitespaceForward(msg, valueEnd); + if (next < msg.length && msg[next] == ',') { + return new int[]{keyStart, next + 1}; + } + // Sole field in its object: remove key and value without any comma. + return new int[]{keyStart, valueEnd}; + } + + /** + * Returns the exclusive end index of the JSON object that opens at {@code start}, tracking + * nested objects and quoted strings (including escape sequences) so that braces inside strings + * are not counted. Returns {@code -1} if no matching closing brace is found. + */ + private static int scanObjectEnd(byte[] msg, int start) { + int depth = 0; + boolean inString = false; + boolean escaped = false; + + for (int i = start; i < msg.length; i++) { + byte c = msg[i]; + + if (inString) { + if (escaped) { + escaped = false; + continue; + } + if (c == '\\') { + escaped = true; + continue; + } + if (c == '"') { + inString = false; + } + continue; + } + + switch (c) { + case '"': inString = true; break; + case '{': depth++; break; + case '}': + depth--; + if (depth == 0) { + return i + 1; + } + break; + default: break; + } + } + return -1; + } + + // ============================================================================ + // String-encoded carrier stripping: Moshi-based (for SNS Message and similar) + // ============================================================================ + + /** + * Removes {@code "_datadog"} from every top-level field whose value is a string-encoded JSON + * object (e.g. an SNS {@code Message} field). Uses Moshi only when {@code \"_datadog\"} is + * actually present in the bytes. Returns {@code msg} unchanged on any error. + */ + private static byte[] stripStringEncodedCarriers(byte[] msg) { + if (indexOf(msg, 0, ESCAPED_KEY_BYTES) < 0) { + return msg; + } + String json = new String(msg, StandardCharsets.UTF_8); + try { + Map envelope = MAP_ADAPTER.fromJson(json); + if (envelope == null) { + return msg; + } + boolean changed = false; + for (Map.Entry entry : envelope.entrySet()) { + if (entry.getValue() instanceof String) { + String stripped = stripFromJsonString((String) entry.getValue()); + if (stripped != null) { + entry.setValue(stripped); + changed = true; + } + } + } + if (!changed) { + return msg; + } return MAP_ADAPTER.toJson(envelope).getBytes(StandardCharsets.UTF_8); } catch (IOException | RuntimeException e) { - // Fail open: stripping must never break the handler invocation. - log.debug("Unable to strip injected _datadog context from Lambda payload", e); - return payload; + log.debug("Unable to strip string-encoded _datadog context from Lambda payload", e); + return msg; } } - /** Removes "_datadog" from an object or string-encoded detail value. Null if unchanged. */ - @SuppressWarnings("unchecked") - private static Object stripDetail(Object detail) { - // Object detail: "detail": {"foo": "bar", "_datadog": {...}} - if (detail instanceof Map) { - Map detailMap = (Map) detail; - return detailMap.remove(DATADOG_CARRIER_KEY) != null ? detailMap : null; + // ============================================================================ + // Byte-level utilities + // ============================================================================ + + /** + * Returns the start index of the first occurrence of {@code needle} in {@code haystack} at or + * after {@code from}, or {@code -1} if not found. + */ + private static int indexOf(byte[] haystack, int from, byte[] needle) { + if (needle.length == 0) { + return from; } - // String-encoded detail: "detail": "{\"foo\":\"bar\",\"_datadog\":{...}}" - if (detail instanceof String) { - return stripFromJsonString((String) detail); + int limit = haystack.length - needle.length; + outer: + for (int i = from; i <= limit; i++) { + for (int j = 0; j < needle.length; j++) { + if (haystack[i + j] != needle[j]) { + continue outer; + } + } + return i; } - return null; + return -1; + } + + /** Returns the first index at or after {@code start} that is not JSON whitespace. */ + private static int skipWhitespaceForward(byte[] msg, int start) { + while (start < msg.length && isJSONWhitespace(msg[start])) { + start++; + } + return start; + } + + /** + * Returns the first index at or before {@code start} that is not JSON whitespace, + * or a negative value if only whitespace precedes. + */ + private static int skipWhitespaceBackward(byte[] msg, int start) { + while (start >= 0 && isJSONWhitespace(msg[start])) { + start--; + } + return start; + } + + /** Returns {@code true} if {@code c} is a JSON whitespace byte (space, tab, CR, or LF). */ + private static boolean isJSONWhitespace(byte c) { + return c == ' ' || c == '\t' || c == '\n' || c == '\r'; } /** Parses, strips "_datadog", and re-encodes a JSON-object-shaped string. Null if unchanged. */ @@ -192,16 +428,17 @@ public Object fromJson(JsonReader reader) throws IOException { } } - /** Returns Long for integer literals, Double for fractional/scientific ones. */ + /** Returns Long for integer literals, BigDecimal for fractional/scientific ones. */ private Object readNumber(JsonReader reader) throws IOException { String literal = reader.nextString(); // raw text avoids precision loss before parsing if (literal.indexOf('.') >= 0 || literal.indexOf('e') >= 0 || literal.indexOf('E') >= 0) { - return Double.parseDouble(literal); + // BigDecimal preserves full precision instead of losing digits like Double would. + return new BigDecimal(literal); } try { return Long.parseLong(literal); } catch (NumberFormatException tooLargeForLong) { - return Double.parseDouble(literal); + return new BigDecimal(literal); } } diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java index 4a8629ce54b..314f305ae84 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java @@ -11,9 +11,9 @@ /** * Randomized invariant testing for {@link StripInjectedContext}, mirroring the repo's existing - * fuzz-style test pattern (see {@code TagMapFuzzTest}). Generates many random EventBridge-shaped - * payloads instead of enumerating fixed cases by hand, and asserts core invariants hold for every - * one of them. + * fuzz-style test pattern (see {@code TagMapFuzzTest}). Generates many random payloads in + * EventBridge, SQS, and SNS shapes instead of enumerating fixed cases by hand, and asserts + * core invariants hold for every one of them. */ public final class StripInjectedContextFuzzTest { @@ -48,6 +48,18 @@ void neverThrowsAndAlwaysRemovesDatadogWhenPresent() { } private static byte[] randomPayload(ThreadLocalRandom random) { + switch (random.nextInt(3)) { + case 0: + return randomEventBridgePayload(random); + case 1: + return randomTopLevelCarrierPayload(random); + default: + return randomStringEncodedFieldPayload(random); + } + } + + // EventBridge-shaped: {"detail-type": ..., "detail": {...}} or {"detail": "..."} + private static byte[] randomEventBridgePayload(ThreadLocalRandom random) { StringBuilder sb = new StringBuilder(); sb.append("{\"detail-type\":\"order.created\",\"detail\":"); boolean asString = random.nextBoolean(); @@ -63,6 +75,45 @@ private static byte[] randomPayload(ThreadLocalRandom random) { return sb.toString().getBytes(StandardCharsets.UTF_8); } + // SQS/generic-shaped: _datadog appears directly as a top-level key. + // Randomly places _datadog as the sole field, the first field, a middle field, + // or the last field so that all four comma-placement cases in carrierRemovalRange + // are exercised across the fuzz iterations. + private static byte[] randomTopLevelCarrierPayload(ThreadLocalRandom random) { + boolean includeDatadog = random.nextBoolean(); + // Allow 0 regular fields so the sole-field case ({"_datadog":{...}}) is also generated. + int fieldCount = random.nextInt(0, 4); + // Choose a random insertion position for the _datadog key among the regular fields. + // Position 0 = first, fieldCount = last, values in between = middle. + int datadogPos = includeDatadog ? random.nextInt(0, fieldCount + 1) : fieldCount + 1; + + StringBuilder sb = new StringBuilder("{"); + boolean needsComma = false; + int regularIdx = 0; + + for (int slot = 0; slot <= fieldCount; slot++) { + if (slot == datadogPos) { + // Insert _datadog at the chosen position. + if (needsComma) sb.append(','); + sb.append("\"_datadog\":{\"x-datadog-trace-id\":\"") + .append(random.nextLong()) + .append("\"}"); + needsComma = true; + } + if (regularIdx < fieldCount) { + // Insert one regular field at this slot. + if (needsComma) sb.append(','); + sb.append("\"field").append(regularIdx).append("\":"); + appendRandomValue(sb, random); + needsComma = true; + regularIdx++; + } + } + + sb.append('}'); + return sb.toString().getBytes(StandardCharsets.UTF_8); + } + private static String randomDetailObject(ThreadLocalRandom random, boolean includeDatadog) { StringBuilder sb = new StringBuilder("{"); int fieldCount = random.nextInt(1, 6); @@ -83,6 +134,16 @@ private static String randomDetailObject(ThreadLocalRandom random, boolean inclu return sb.toString(); } + // SNS-shaped: a top-level string field contains a string-encoded JSON object with _datadog. + private static byte[] randomStringEncodedFieldPayload(ThreadLocalRandom random) { + boolean includeDatadog = random.nextBoolean(); + String innerJson = randomDetailObject(random, includeDatadog); + String escaped = innerJson.replace("\"", "\\\""); + String payload = + "{\"Type\":\"Notification\",\"Message\":\"" + escaped + "\",\"Subject\":\"test\"}"; + return payload.getBytes(StandardCharsets.UTF_8); + } + private static void appendRandomValue(StringBuilder sb, ThreadLocalRandom random) { switch (random.nextInt(4)) { case 0: diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java index 6cc374ee66b..237136291dd 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java @@ -81,12 +81,15 @@ void returnsOriginalWhenDatadogKeyIsAbsent() { } @Test - void returnsOriginalWhenDetailFieldIsMissing() { + void removesTopLevelDatadogKey() { + // _datadog appearing directly as a top-level key is removed regardless of event shape. byte[] input = bytes("{\"detail-type\":\"order.created\",\"_datadog\":{\"trace\":\"1\"}}"); byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); - assertArrayEquals(input, result); + assertTrue(!json.contains("_datadog"), "_datadog at the top level must be stripped"); + assertTrue(json.contains("detail-type"), "unrelated top-level fields must survive"); } @Test @@ -106,8 +109,8 @@ void returnsOriginalForNullOrEmptyPayload() { @Test void returnsOriginalWhenDetailIsNeitherObjectNorString() { - // e.g. detail is a number/array/null — nothing sane to strip from - byte[] input = bytes("{\"detail\":42,\"_datadog\":{\"trace\":\"1\"}}"); + // e.g. detail is a number — nothing to strip from, so the fast path returns unchanged. + byte[] input = bytes("{\"detail\":42}"); byte[] result = StripInjectedContext.stripInternal(input); @@ -175,6 +178,168 @@ void stripRemovesDatadogKeyWhenConfigIsEnabled() { assertTrue(!json.contains("_datadog")); } + // ============================================================================ + // Non-EventBridge payloads: direct top-level _datadog key (SQS, Kinesis, etc.) + // ============================================================================ + + @Test + void removesTopLevelDatadogKeyFromSqsStylePayload() { + // SQS-style: _datadog appears as a direct top-level property alongside business data. + byte[] input = + bytes("{\"orderId\":\"abc-123\",\"amount\":99," + + "\"_datadog\":{\"x-datadog-trace-id\":\"456\",\"x-datadog-parent-id\":\"789\"}}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog carrier must be removed"); + assertTrue(json.contains("\"orderId\":\"abc-123\""), "orderId must survive untouched"); + assertTrue(json.contains("\"amount\":99"), "amount must survive untouched"); + } + + @Test + void returnsOriginalWhenTopLevelDatadogKeyIsAbsent() { + // No _datadog anywhere: fast path should return the original byte array unchanged. + byte[] input = bytes("{\"orderId\":\"abc-123\",\"amount\":99}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + // ============================================================================ + // Non-EventBridge payloads: string-encoded JSON fields (SNS, etc.) + // ============================================================================ + + @Test + void removesDatadogKeyFromStringEncodedNonDetailField() { + // SNS-style: the "Message" field is a string-encoded JSON object containing _datadog. + byte[] input = + bytes( + "{\"Type\":\"Notification\"," + + "\"Message\":\"{\\\"hello\\\":\\\"world\\\"," + + "\\\"_datadog\\\":{\\\"x-datadog-trace-id\\\":\\\"123\\\"}}\"," + + "\"Subject\":\"test-event\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog must be stripped from string-encoded Message"); + assertTrue(json.contains("hello"), "business data inside Message must survive"); + assertTrue(json.contains("Subject"), "other top-level fields must survive"); + } + + @Test + void returnsOriginalWhenStringEncodedFieldContainsNoDatadogKey() { + // A string-encoded JSON field that doesn't carry _datadog must be left completely unchanged. + byte[] input = + bytes("{\"Type\":\"Notification\",\"Message\":\"{\\\"hello\\\":\\\"world\\\"}\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + @Test + void handlesMultipleTopLevelFieldsWithDatadogCarrier() { + // Both a direct top-level _datadog key and a string-encoded field containing _datadog + // are present at the same time; both must be stripped. + byte[] input = + bytes( + "{\"_datadog\":{\"trace\":\"1\"}," + + "\"Message\":\"{\\\"key\\\":\\\"val\\\"," + + "\\\"_datadog\\\":{\\\"trace\\\":\\\"2\\\"}}\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "all _datadog carriers must be stripped"); + assertTrue(json.contains("key"), "business data inside Message must survive"); + } + + // ============================================================================ + // Byte-level splice edge cases: comma placement around the removed carrier + // ============================================================================ + + @Test + void removesDatadogWhenItIsTheSoleField() { + // carrierRemovalRange: no preceding comma, no following comma -> remove key+value only. + // Result must be a valid empty object, not a dangling comma. + byte[] input = bytes("{\"_datadog\":{\"trace\":\"1\"}}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog must be removed"); + assertTrue(json.equals("{}"), "result must be a valid empty object, got: " + json); + } + + @Test + void removesDatadogWhenItIsTheFirstField() { + // carrierRemovalRange: no preceding comma but a following comma exists -> remove trailing comma. + // If the trailing comma is left in, the result is invalid JSON: {"other":"val"} vs {,"other":"val"}. + byte[] input = bytes("{\"_datadog\":{\"trace\":\"1\"},\"other\":\"val\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog must be removed"); + assertTrue(json.contains("\"other\":\"val\""), "following field must survive"); + assertTrue(!json.startsWith("{,"), "result must not start with a dangling comma"); + } + + @Test + void removesDatadogWhenItIsTheMiddleField() { + // carrierRemovalRange: both a preceding comma and a following comma exist. + // The preceding comma is consumed (preferred), so the remaining fields stay adjacent. + byte[] input = + bytes("{\"first\":\"a\",\"_datadog\":{\"trace\":\"1\"},\"last\":\"b\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog must be removed"); + assertTrue(json.contains("\"first\":\"a\""), "first field must survive"); + assertTrue(json.contains("\"last\":\"b\""), "last field must survive"); + assertTrue(!json.contains(",,"), "result must not contain a double comma"); + } + + @Test + void doesNotRemoveDatadogKeyWhenValueIsNotAnObject() { + // The byte-level splice only removes "_datadog":{...} (object value form). + // A string, null, or number value is not a carrier and must be left untouched. + byte[] input = bytes("{\"_datadog\":\"not-a-carrier\",\"other\":1}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + @Test + void doesNotMatchDatadogTextEmbeddedInsideAStringValue() { + // A string value that happens to contain the text "_datadog:" must not trigger removal. + // isObjectProperty() guards against this by checking the preceding byte is '{' or ','. + byte[] input = bytes("{\"note\":\"the _datadog: key is for tracing\",\"other\":1}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result); + } + + @Test + void handlesWhitespacePaddedJson() { + // skipWhitespaceForward and skipWhitespaceBackward must tolerate spaces, tabs, and newlines + // around the key, colon, and value so that pretty-printed payloads are handled correctly. + byte[] input = + bytes("{ \"_datadog\": { \"trace\" : \"1\" } , \"other\" : 1 }"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog must be removed from whitespace-padded JSON"); + assertTrue(json.contains("\"other\""), "unrelated field must survive"); + } + // ============================================================================ // replaceInputStream(): the InputStream-based helper used by the advice // ============================================================================ From ab28019c76ce42f3fb4eb3d6f5e276d4dbca2cbb Mon Sep 17 00:00:00 2001 From: "eunhye.kim" Date: Fri, 2 Oct 2026 16:41:00 +0900 Subject: [PATCH 3/6] Move StripInjectedContext into aws-java-lambda-handler instrumentation --- .../aws-java-lambda-handler-1.2/build.gradle | 2 + .../java/StripInjectedContextBenchmark.java | 82 +++ .../lambda/LambdaHandlerInstrumentation.java | 4 +- .../aws/v1/lambda/StripInjectedContext.java | 266 ++++++++++ .../java}/StripInjectedContextFuzzTest.java | 8 +- .../test/java}/StripInjectedContextTest.java | 185 ++++++- .../lambda/StripInjectedContextBenchmark.java | 77 --- .../java/datadog/trace/core/CoreTracer.java | 7 - .../trace/lambda/StripInjectedContext.java | 476 ------------------ .../instrumentation/api/AgentTracer.java | 8 - 10 files changed, 528 insertions(+), 587 deletions(-) create mode 100644 dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/jmh/java/StripInjectedContextBenchmark.java create mode 100644 dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java rename {dd-trace-core/src/test/java/datadog/trace/lambda => dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java}/StripInjectedContextFuzzTest.java (97%) rename {dd-trace-core/src/test/java/datadog/trace/lambda => dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java}/StripInjectedContextTest.java (67%) delete mode 100644 dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java delete mode 100644 dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle index b0c49f3eb4f..db5f54dbc63 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle @@ -1,5 +1,6 @@ plugins { id 'dd-trace-java.module.instrumentation' + id 'dd-trace-java.jmh-conventions' } muzzle { @@ -15,4 +16,5 @@ addTestSuiteForDir('latestDepTest', 'test') dependencies { compileOnly group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1' testImplementation group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1' + testImplementation project(':utils:test-utils') } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/jmh/java/StripInjectedContextBenchmark.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/jmh/java/StripInjectedContextBenchmark.java new file mode 100644 index 00000000000..110501322cf --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/jmh/java/StripInjectedContextBenchmark.java @@ -0,0 +1,82 @@ +package datadog.trace.instrumentation.aws.v1.lambda; + +import static java.util.concurrent.TimeUnit.MICROSECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; + +import java.nio.charset.StandardCharsets; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@State(Scope.Benchmark) +@Warmup(iterations = 3, time = 5, timeUnit = SECONDS) +@Measurement(iterations = 3, time = 5, timeUnit = SECONDS) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(MICROSECONDS) +@Fork(value = 1) +public class StripInjectedContextBenchmark { + + @Param({"true", "false"}) + boolean containsDatadogCarrier; + + // Event shapes: "eventbridge", "sqs", "sns" + @Param({"eventbridge", "sqs", "sns"}) + String payloadShape; + + byte[] payload; + + @Setup(Level.Trial) + public void setUp() { + String carrier = "{\"x-datadog-trace-id\":\"123\",\"x-datadog-parent-id\":\"456\"}"; + + switch (payloadShape) { + case "sqs": + { + // SQS-style: _datadog appears as a direct top-level key. + String body = + containsDatadogCarrier + ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":" + carrier + "}" + : "{\"orderId\":42,\"customer\":\"acme\"}"; + payload = body.getBytes(StandardCharsets.UTF_8); + break; + } + case "sns": + { + // SNS-style: _datadog is embedded inside a string-encoded "Message" field. + String inner = + containsDatadogCarrier + ? "{\\\"orderId\\\":42,\\\"_datadog\\\":" + carrier.replace("\"", "\\\"") + "}" + : "{\\\"orderId\\\":42}"; + String envelope = "{\"Type\":\"Notification\",\"Message\":\"" + inner + "\"}"; + payload = envelope.getBytes(StandardCharsets.UTF_8); + break; + } + default: + { + // EventBridge-style: _datadog is inside the "detail" object. + String detail = + containsDatadogCarrier + ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":" + carrier + "}" + : "{\"orderId\":42,\"customer\":\"acme\"}"; + String envelope = "{\"detail-type\":\"order.created\",\"detail\":" + detail + "}"; + payload = envelope.getBytes(StandardCharsets.UTF_8); + break; + } + } + } + + @Benchmark + public void stripInternal(Blackhole blackhole) { + blackhole.consume(StripInjectedContext.stripInternal(payload)); + } +} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index 910a4c06e0d..33f86d2935c 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.aws.v1.lambda; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; -import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -13,6 +12,7 @@ import static net.bytebuddy.asm.Advice.Origin; import static net.bytebuddy.asm.Advice.This; import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -104,7 +104,7 @@ static ContextScope enter( AgentTracer.get().notifyLambdaStart(inputStream, lambdaRequestId); // Skip the strip pass entirely (extra read/parse/re-serialize) unless explicitly enabled. if (Config.get().isLambdaStripInjectedContextEnabled()) { - InputStream stripped = AgentTracer.get().stripLambdaInjectedContext(inputStream); + InputStream stripped = StripInjectedContext.replaceInputStream(inputStream); if (stripped != null) { inputStream = stripped; } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java new file mode 100644 index 00000000000..f283a4be0db --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java @@ -0,0 +1,266 @@ +package datadog.trace.instrumentation.aws.v1.lambda; + +import datadog.json.JsonMapper; +import datadog.json.JsonReader; +import datadog.trace.api.Config; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Removes {@code "_datadog":{...}} trace-context carriers injected by Datadog into Lambda events + * (EventBridge, SNS, SQS, ...) before the customer handler reads them. + * + *

Carriers are cut out of the JSON text instead of parsing and re-serializing the payload, so + * numbers, key order and formatting of customer data are kept. Carriers inside string-encoded JSON + * (an SNS {@code Message}, an SQS {@code body}) are handled by decoding that string, stripping it + * recursively up to {@link #MAX_DEPTH} levels, and re-encoding it. A {@code _datadog} key whose + * value is not an object is customer data and is kept. Fail-open: on any error the original payload + * is used. + */ +public final class StripInjectedContext { + + private static final Logger log = LoggerFactory.getLogger(StripInjectedContext.class); + + /** SNS-to-SQS fan-out needs 2 levels of string-encoded JSON; one more is a safety margin. */ + static final int MAX_DEPTH = 3; + + private static final String CARRIER_KEY = "_datadog"; + private static final String QUOTED_CARRIER_KEY = "\"" + CARRIER_KEY + "\""; + private static final byte[] CARRIER_KEY_BYTES = CARRIER_KEY.getBytes(StandardCharsets.UTF_8); + + /** Returned by {@link #carrierRange} when the carrier object never closes (malformed JSON). */ + private static final int[] UNCLOSED = new int[0]; + + private StripInjectedContext() {} + + /** Reads {@code in} fully, strips "_datadog", and returns a new stream over the result. */ + public static ByteArrayInputStream replaceInputStream(InputStream in) { + if (in == null) { + return null; + } + byte[] original; + try { + original = readAllBytes(in); + } catch (IOException e) { + log.debug("Unable to read Lambda payload for _datadog stripping", e); + return null; + } + try { + return new ByteArrayInputStream(strip(original)); + } catch (RuntimeException e) { + log.debug("Unable to strip _datadog from Lambda payload", e); + return new ByteArrayInputStream(original); + } + } + + /** Strips "_datadog" from {@code payload} when the config is enabled. */ + public static byte[] strip(byte[] payload) { + if (!Config.get().isLambdaStripInjectedContextEnabled()) { + return payload; + } + return stripInternal(payload); + } + + /** Unconditional strip logic, package-private so tests can call it directly. */ + static byte[] stripInternal(byte[] payload) { + if (payload == null || !containsCarrierKey(payload)) { + return payload; + } + String json = new String(payload, StandardCharsets.UTF_8); + String stripped = stripJson(json, 0); + return stripped.equals(json) ? payload : stripped.getBytes(StandardCharsets.UTF_8); + } + + /** + * Walks the string tokens of {@code json} once. A {@code "_datadog"} key with an object value is + * cut out; any other string containing {@code _datadog} is treated as encoded JSON and stripped + * recursively. Returns {@code json} itself when nothing changed. + */ + private static String stripJson(String json, int depth) { + StringBuilder out = null; + int copyFrom = 0; + int nextKey = -1; + int start = json.indexOf('"'); + while (start >= 0) { + if (nextKey < start) { + nextKey = json.indexOf(CARRIER_KEY, start); + if (nextKey < 0) { + break; + } + } + int end = stringEnd(json, start); + if (end < 0) { + break; + } + int[] carrier = carrierRange(json, start, end, copyFrom); + if (carrier == UNCLOSED) { + break; + } + if (carrier != null) { + if (out == null) { + out = new StringBuilder(json.length()); + } + out.append(json, copyFrom, carrier[0]); + copyFrom = carrier[1]; + end = carrier[1]; + } else if (nextKey < end && depth < MAX_DEPTH) { + String replacement = stripEncodedJson(json.substring(start, end), depth); + if (replacement != null) { + if (out == null) { + out = new StringBuilder(json.length()); + } + out.append(json, copyFrom, start).append(replacement); + copyFrom = end; + } + } + start = json.indexOf('"', end); + } + if (out == null) { + return json; + } + return out.append(json, copyFrom, json.length()).toString(); + } + + /** + * If the string token {@code [keyStart, keyEnd)} is a {@code "_datadog"} key with an object + * value, returns the {@code [start, end)} range to cut, including one adjacent comma so the JSON + * stays valid. A preceding comma before {@code copyFrom} was already cut and is not reused. + * Returns {@link #UNCLOSED} if the value object never closes. + */ + private static int[] carrierRange(String json, int keyStart, int keyEnd, int copyFrom) { + if (keyEnd - keyStart != QUOTED_CARRIER_KEY.length() + || !json.startsWith(QUOTED_CARRIER_KEY, keyStart)) { + return null; + } + int colon = skipWhitespace(json, keyEnd); + if (colon >= json.length() || json.charAt(colon) != ':') { + return null; + } + int valueStart = skipWhitespace(json, colon + 1); + if (valueStart >= json.length() || json.charAt(valueStart) != '{') { + return null; + } + int valueEnd = objectEnd(json, valueStart); + if (valueEnd < 0) { + return UNCLOSED; + } + int before = skipWhitespaceBackward(json, keyStart - 1); + if (before >= copyFrom && json.charAt(before) == ',') { + return new int[] {before, valueEnd}; + } + int after = skipWhitespace(json, valueEnd); + if (after < json.length() && json.charAt(after) == ',') { + return new int[] {keyStart, after + 1}; + } + return new int[] {keyStart, valueEnd}; + } + + /** + * Decodes the JSON string literal {@code literal}, strips it if it holds a JSON object, and + * returns the re-encoded literal, or {@code null} if nothing changed. + */ + private static String stripEncodedJson(String literal, int depth) { + String decoded; + try { + decoded = new JsonReader(literal).nextString(); + } catch (IOException e) { + return null; + } + int first = skipWhitespace(decoded, 0); + if (first >= decoded.length() || decoded.charAt(first) != '{') { + return null; + } + String stripped = stripJson(decoded, depth + 1); + return stripped.equals(decoded) ? null : JsonMapper.toJson(stripped); + } + + /** Returns the index just past the closing quote of the string opening at {@code start}. */ + private static int stringEnd(String json, int start) { + for (int i = start + 1; i < json.length(); i++) { + char c = json.charAt(i); + if (c == '\\') { + i++; + } else if (c == '"') { + return i + 1; + } + } + return -1; + } + + /** Returns the index just past the closing brace of the object opening at {@code start}. */ + private static int objectEnd(String json, int start) { + int depth = 0; + for (int i = start; i < json.length(); i++) { + char c = json.charAt(i); + if (c == '"') { + int end = stringEnd(json, i); + if (end < 0) { + return -1; + } + i = end - 1; + } else if (c == '{') { + depth++; + } else if (c == '}' && --depth == 0) { + return i + 1; + } + } + return -1; + } + + private static int skipWhitespace(String json, int i) { + while (i < json.length() && isWhitespace(json.charAt(i))) { + i++; + } + return i; + } + + private static int skipWhitespaceBackward(String json, int i) { + while (i >= 0 && isWhitespace(json.charAt(i))) { + i--; + } + return i; + } + + private static boolean isWhitespace(char c) { + return c == ' ' || c == '\t' || c == '\n' || c == '\r'; + } + + private static boolean containsCarrierKey(byte[] payload) { + outer: + for (int i = 0; i <= payload.length - CARRIER_KEY_BYTES.length; i++) { + for (int j = 0; j < CARRIER_KEY_BYTES.length; j++) { + if (payload[i + j] != CARRIER_KEY_BYTES[j]) { + continue outer; + } + } + return true; + } + return false; + } + + /** Reads all bytes from {@code in}, restoring its position first if it's a byte stream. */ + private static byte[] readAllBytes(InputStream in) throws IOException { + if (in instanceof ByteArrayInputStream) { + // mark/reset instead of draining, in case the stream is read again elsewhere. + ByteArrayInputStream bais = (ByteArrayInputStream) in; + bais.mark(Integer.MAX_VALUE); + byte[] bytes = new byte[bais.available()]; + int read = bais.read(bytes); + bais.reset(); + return (read == bytes.length) ? bytes : Arrays.copyOf(bytes, Math.max(read, 0)); + } + ByteArrayOutputStream buffer = new ByteArrayOutputStream(); + byte[] chunk = new byte[8192]; + int n; + while ((n = in.read(chunk)) != -1) { + buffer.write(chunk, 0, n); + } + return buffer.toByteArray(); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java similarity index 97% rename from dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java rename to dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java index 314f305ae84..c17f2b3c75f 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextFuzzTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java @@ -1,4 +1,4 @@ -package datadog.trace.lambda; +package datadog.trace.instrumentation.aws.v1.lambda; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -12,8 +12,8 @@ /** * Randomized invariant testing for {@link StripInjectedContext}, mirroring the repo's existing * fuzz-style test pattern (see {@code TagMapFuzzTest}). Generates many random payloads in - * EventBridge, SQS, and SNS shapes instead of enumerating fixed cases by hand, and asserts - * core invariants hold for every one of them. + * EventBridge, SQS, and SNS shapes instead of enumerating fixed cases by hand, and asserts core + * invariants hold for every one of them. */ public final class StripInjectedContextFuzzTest { @@ -77,7 +77,7 @@ private static byte[] randomEventBridgePayload(ThreadLocalRandom random) { // SQS/generic-shaped: _datadog appears directly as a top-level key. // Randomly places _datadog as the sole field, the first field, a middle field, - // or the last field so that all four comma-placement cases in carrierRemovalRange + // or the last field so that all four comma-placement cases in carrierRange // are exercised across the fuzz iterations. private static byte[] randomTopLevelCarrierPayload(ThreadLocalRandom random) { boolean includeDatadog = random.nextBoolean(); diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java similarity index 67% rename from dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java rename to dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java index 237136291dd..dd131a2fd50 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/StripInjectedContextTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java @@ -1,8 +1,9 @@ -package datadog.trace.lambda; +package datadog.trace.instrumentation.aws.v1.lambda; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; import datadog.trace.api.config.GeneralConfig; @@ -10,6 +11,7 @@ import datadog.trace.test.util.DDJavaSpecification; import java.io.ByteArrayInputStream; import java.nio.charset.StandardCharsets; +import java.time.Duration; import org.junit.jupiter.api.Test; /** Unit tests for {@link StripInjectedContext}. */ @@ -186,8 +188,9 @@ void stripRemovesDatadogKeyWhenConfigIsEnabled() { void removesTopLevelDatadogKeyFromSqsStylePayload() { // SQS-style: _datadog appears as a direct top-level property alongside business data. byte[] input = - bytes("{\"orderId\":\"abc-123\",\"amount\":99," - + "\"_datadog\":{\"x-datadog-trace-id\":\"456\",\"x-datadog-parent-id\":\"789\"}}"); + bytes( + "{\"orderId\":\"abc-123\",\"amount\":99," + + "\"_datadog\":{\"x-datadog-trace-id\":\"456\",\"x-datadog-parent-id\":\"789\"}}"); byte[] result = StripInjectedContext.stripInternal(input); String json = toUtf8(result); @@ -263,7 +266,7 @@ void handlesMultipleTopLevelFieldsWithDatadogCarrier() { @Test void removesDatadogWhenItIsTheSoleField() { - // carrierRemovalRange: no preceding comma, no following comma -> remove key+value only. + // carrierRange: no preceding comma, no following comma -> remove key+value only. // Result must be a valid empty object, not a dangling comma. byte[] input = bytes("{\"_datadog\":{\"trace\":\"1\"}}"); @@ -276,8 +279,10 @@ void removesDatadogWhenItIsTheSoleField() { @Test void removesDatadogWhenItIsTheFirstField() { - // carrierRemovalRange: no preceding comma but a following comma exists -> remove trailing comma. - // If the trailing comma is left in, the result is invalid JSON: {"other":"val"} vs {,"other":"val"}. + // carrierRange: no preceding comma but a following comma exists -> remove trailing + // comma. + // If the trailing comma is left in, the result is invalid JSON: {"other":"val"} vs + // {,"other":"val"}. byte[] input = bytes("{\"_datadog\":{\"trace\":\"1\"},\"other\":\"val\"}"); byte[] result = StripInjectedContext.stripInternal(input); @@ -290,10 +295,9 @@ void removesDatadogWhenItIsTheFirstField() { @Test void removesDatadogWhenItIsTheMiddleField() { - // carrierRemovalRange: both a preceding comma and a following comma exist. + // carrierRange: both a preceding comma and a following comma exist. // The preceding comma is consumed (preferred), so the remaining fields stay adjacent. - byte[] input = - bytes("{\"first\":\"a\",\"_datadog\":{\"trace\":\"1\"},\"last\":\"b\"}"); + byte[] input = bytes("{\"first\":\"a\",\"_datadog\":{\"trace\":\"1\"},\"last\":\"b\"}"); byte[] result = StripInjectedContext.stripInternal(input); String json = toUtf8(result); @@ -318,7 +322,7 @@ void doesNotRemoveDatadogKeyWhenValueIsNotAnObject() { @Test void doesNotMatchDatadogTextEmbeddedInsideAStringValue() { // A string value that happens to contain the text "_datadog:" must not trigger removal. - // isObjectProperty() guards against this by checking the preceding byte is '{' or ','. + // String tokens are skipped as a whole, so text inside a value is never treated as a key. byte[] input = bytes("{\"note\":\"the _datadog: key is for tracing\",\"other\":1}"); byte[] result = StripInjectedContext.stripInternal(input); @@ -328,10 +332,9 @@ void doesNotMatchDatadogTextEmbeddedInsideAStringValue() { @Test void handlesWhitespacePaddedJson() { - // skipWhitespaceForward and skipWhitespaceBackward must tolerate spaces, tabs, and newlines + // skipWhitespace and skipWhitespaceBackward must tolerate spaces, tabs, and newlines // around the key, colon, and value so that pretty-printed payloads are handled correctly. - byte[] input = - bytes("{ \"_datadog\": { \"trace\" : \"1\" } , \"other\" : 1 }"); + byte[] input = bytes("{ \"_datadog\": { \"trace\" : \"1\" } , \"other\" : 1 }"); byte[] result = StripInjectedContext.stripInternal(input); String json = toUtf8(result); @@ -366,4 +369,160 @@ void replaceInputStreamStripsAndDoesNotConsumeTheOriginalStream() { String originalJson = new String(readAll(original), StandardCharsets.UTF_8); assertTrue(originalJson.contains("_datadog"), "original stream must still be fully readable"); } + + @Test + void doesNotRemoveNonObjectDatadogValueInsideStringEncodedField() { + // Regression test for review feedback: a string-encoded business payload (e.g. SQS body) + // whose "_datadog" field holds a plain string, not an object, must survive untouched. + // Only object-shaped propagation carriers ("_datadog":{...}) are carriers; anything else + // is customer data that happens to reuse the same key name. + byte[] input = + bytes( + "{\"Records\":[{\"body\":\"{\\\"_datadog\\\":\\\"customer-value\\\"," + + "\\\"order\\\":\\\"foo\\\"}\"}]}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals( + input, result, "non-object _datadog value inside a string-encoded field must be preserved"); + } + + @Test + void removesDatadogFromNestedSqsRecordsBody() { + // SQS Records[*].body is itself a string-encoded JSON object. + byte[] input = + bytes( + "{\"Records\":[{\"body\":\"{\\\"orderId\\\":1," + + "\\\"_datadog\\\":{\\\"trace\\\":\\\"1\\\"}}\"}]}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog inside Records[].body must be stripped"); + assertTrue(json.contains("orderId")); + } + + @Test + void removesDatadogFromDoubleEncodedSnsInsideSqsBody() { + // SQS body containing an SNS envelope whose Message is itself string-encoded JSON. + String snsMessage = + "{\\\\\\\"orderId\\\\\\\":1,\\\\\\\"_datadog\\\\\\\":" + + "{\\\\\\\"trace\\\\\\\":\\\\\\\"1\\\\\\\"}}"; + byte[] input = + bytes("{\"Records\":[{\"body\":\"{\\\"Message\\\":\\\"" + snsMessage + "\\\"}\"}]}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "double-encoded _datadog must be stripped"); + assertTrue(json.contains("orderId")); + } + + @Test + void removesObjectCarrierWithWhitespaceBeforeColon() { + byte[] input = bytes("{\"_datadog\" : {\"trace\":\"1\"},\"other\":1}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "whitespace before colon must not block removal"); + assertTrue(json.contains("\"other\":1")); + } + + @Test + void doesNotRecurseBeyondMaxDepth() { + // The carrier is wrapped one level deeper than MAX_DEPTH allows. The final layer is never + // decoded, + // so no plain "_datadog" object key is ever seen. + int wraps = StripInjectedContext.MAX_DEPTH + 1; + String payload = "{\"orderId\":1,\"_datadog\":{\"trace\":\"1\"}}"; + for (int i = 0; i < wraps; i++) { + payload = "{\"detail\":" + jsonQuote(payload) + "}"; + } + + byte[] input = bytes(payload); + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result, "carrier nested beyond MAX_DEPTH must be returned unchanged"); + } + + /** Encodes {@code raw} as a JSON string literal (adds surrounding quotes and escapes). */ + private static String jsonQuote(String raw) { + StringBuilder sb = new StringBuilder("\""); + for (int i = 0; i < raw.length(); i++) { + char c = raw.charAt(i); + if (c == '"' || c == '\\') { + sb.append('\\'); + } + sb.append(c); + } + return sb.append('"').toString(); + } + + @Test + void stripsCarrierNestedExactlyAtMaxDepth() { + int wraps = StripInjectedContext.MAX_DEPTH; + String payload = "{\"orderId\":1,\"_datadog\":{\"trace\":\"1\"}}"; + for (int i = 0; i < wraps; i++) { + payload = "{\"detail\":" + jsonQuote(payload) + "}"; + } + + byte[] result = StripInjectedContext.stripInternal(bytes(payload)); + String json = toUtf8(result); + + assertTrue( + !json.contains("_datadog"), "carrier nested exactly at MAX_DEPTH must still be stripped"); + assertTrue(json.contains("orderId")); + } + + @Test + void preservesSurrogatePairEmojiInStringEncodedBusinessData() { + // Business text containing an astral-plane character encoded as a surrogate pair of unicode + // escapes (e.g. an emoji) must survive when it sits in the same string-encoded field as a + // genuine _datadog carrier that gets stripped. + byte[] input = + bytes( + "{\"Message\":\"{\\\"note\\\":\\\"Thanks \\uD83D\\uDE00\\\"," + + "\\\"_datadog\\\":{\\\"trace\\\":\\\"1\\\"}}\"}"); + + byte[] result = StripInjectedContext.stripInternal(input); + String json = toUtf8(result); + + assertTrue(!json.contains("_datadog"), "_datadog must be stripped"); + assertTrue( + json.contains("Thanks \\uD83D\\uDE00"), + "surrogate-pair emoji in business data must survive intact"); + } + + @Test + void doesNotThrowOnSignedHexInUnicodeEscape() { + byte[] input = bytes("{\"Message\":\"{\\\"n\\\":\\\"\\u-001\\\",\\\"_datadog\\\":{}}\"}"); + + String json = toUtf8(StripInjectedContext.stripInternal(input)); + + assertTrue(!json.contains("_datadog"), "_datadog must be stripped"); + } + + @Test + void preservesLoneSurrogateWhileStrippingCarrier() { + byte[] input = bytes("{\"Message\":\"{\\\"n\\\":\\\"x\\uD83D y\\\",\\\"_datadog\\\":{}}\"}"); + + String json = toUtf8(StripInjectedContext.stripInternal(input)); + + assertTrue(!json.contains("_datadog"), "_datadog must be stripped"); + assertTrue(json.contains("x\\uD83D y"), "a lone surrogate must not be turned into '?'"); + } + + @Test + void unclosedCarriersDoNotCauseQuadraticScan() { + StringBuilder body = new StringBuilder(); + for (int i = 0; i < 20_000; i++) { + body.append("{\\\"_datadog\\\":{"); + } + byte[] input = bytes("{\"Records\":[{\"body\":\"" + body + "\"}]}"); + + assertTimeoutPreemptively( + Duration.ofSeconds(1), + () -> assertArrayEquals(input, StripInjectedContext.stripInternal(input))); + } } diff --git a/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java b/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java deleted file mode 100644 index fb33f6ebb02..00000000000 --- a/dd-trace-core/src/jmh/java/datadog/trace/lambda/StripInjectedContextBenchmark.java +++ /dev/null @@ -1,77 +0,0 @@ -package datadog.trace.lambda; - -import static java.util.concurrent.TimeUnit.MICROSECONDS; -import static java.util.concurrent.TimeUnit.SECONDS; - -import java.nio.charset.StandardCharsets; -import org.openjdk.jmh.annotations.Benchmark; -import org.openjdk.jmh.annotations.BenchmarkMode; -import org.openjdk.jmh.annotations.Fork; -import org.openjdk.jmh.annotations.Level; -import org.openjdk.jmh.annotations.Measurement; -import org.openjdk.jmh.annotations.Mode; -import org.openjdk.jmh.annotations.OutputTimeUnit; -import org.openjdk.jmh.annotations.Param; -import org.openjdk.jmh.annotations.Scope; -import org.openjdk.jmh.annotations.Setup; -import org.openjdk.jmh.annotations.State; -import org.openjdk.jmh.annotations.Warmup; -import org.openjdk.jmh.infra.Blackhole; - -@State(Scope.Benchmark) -@Warmup(iterations = 3, time = 5, timeUnit = SECONDS) -@Measurement(iterations = 3, time = 5, timeUnit = SECONDS) -@BenchmarkMode(Mode.AverageTime) -@OutputTimeUnit(MICROSECONDS) -@Fork(value = 1) -public class StripInjectedContextBenchmark { - - @Param({"true", "false"}) - boolean containsDatadogCarrier; - - // Event shapes: "eventbridge", "sqs", "sns" - @Param({"eventbridge", "sqs", "sns"}) - String payloadShape; - - byte[] payload; - - @Setup(Level.Trial) - public void setUp() { - String carrier = - "{\"x-datadog-trace-id\":\"123\",\"x-datadog-parent-id\":\"456\"}"; - - switch (payloadShape) { - case "sqs": { - // SQS-style: _datadog appears as a direct top-level key. - String body = containsDatadogCarrier - ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":" + carrier + "}" - : "{\"orderId\":42,\"customer\":\"acme\"}"; - payload = body.getBytes(StandardCharsets.UTF_8); - break; - } - case "sns": { - // SNS-style: _datadog is embedded inside a string-encoded "Message" field. - String inner = containsDatadogCarrier - ? "{\\\"orderId\\\":42,\\\"_datadog\\\":" + carrier.replace("\"", "\\\"") + "}" - : "{\\\"orderId\\\":42}"; - String envelope = "{\"Type\":\"Notification\",\"Message\":\"" + inner + "\"}"; - payload = envelope.getBytes(StandardCharsets.UTF_8); - break; - } - default: { - // EventBridge-style: _datadog is inside the "detail" object. - String detail = containsDatadogCarrier - ? "{\"orderId\":42,\"customer\":\"acme\",\"_datadog\":" + carrier + "}" - : "{\"orderId\":42,\"customer\":\"acme\"}"; - String envelope = "{\"detail-type\":\"order.created\",\"detail\":" + detail + "}"; - payload = envelope.getBytes(StandardCharsets.UTF_8); - break; - } - } - } - - @Benchmark - public void stripInternal(Blackhole blackhole) { - blackhole.consume(StripInjectedContext.stripInternal(payload)); - } -} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 22d6f60e4ef..9ed3f34aceb 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -119,10 +119,8 @@ import datadog.trace.core.traceinterceptor.LatencyTraceInterceptor; import datadog.trace.lambda.LambdaAppSecHandler; import datadog.trace.lambda.LambdaHandler; -import datadog.trace.lambda.StripInjectedContext; import datadog.trace.util.AgentTaskScheduler; import java.io.IOException; -import java.io.InputStream; import java.lang.ref.WeakReference; import java.util.ArrayList; import java.util.Collection; @@ -1368,11 +1366,6 @@ public void notifyAppSecEnd(AgentSpan span, Object result) { } } - @Override - public InputStream stripLambdaInjectedContext(InputStream in) { - return StripInjectedContext.replaceInputStream(in); - } - @Override public AgentDataStreamsMonitoring getDataStreamsMonitoring() { return dataStreamsMonitoring; diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java b/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java deleted file mode 100644 index a3d6bea7a24..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/StripInjectedContext.java +++ /dev/null @@ -1,476 +0,0 @@ -package datadog.trace.lambda; - -import com.squareup.moshi.JsonAdapter; -import com.squareup.moshi.JsonReader; -import com.squareup.moshi.JsonWriter; -import com.squareup.moshi.Moshi; -import com.squareup.moshi.Types; -import datadog.trace.api.Config; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.InputStream; -import java.lang.reflect.Type; -import java.math.BigDecimal; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Strips the injected "_datadog" trace carrier from Lambda event payloads before the user - * handler runs, when {@code DD_LAMBDA_STRIP_INJECTED_CONTEXT} is enabled. Applies to any - * top-level field value, covering EventBridge, SNS, SQS, Kinesis, and similar event shapes. - * Fail-open: any error or invalid result returns the original payload unchanged. - */ -public final class StripInjectedContext { - - private static final Logger log = LoggerFactory.getLogger(StripInjectedContext.class); - - private static final String DATADOG_CARRIER_KEY = "_datadog"; - - /** Raw bytes of "_datadog" for the fast-path substring check. */ - private static final byte[] DATADOG_CARRIER_BYTES = - DATADOG_CARRIER_KEY.getBytes(StandardCharsets.UTF_8); - - /** - * Byte pattern {@code "_datadog":} used to locate object-form carriers in raw JSON bytes. - * The colon is included so a single scan step covers both the key and the separator, - * matching the approach used in the Go tracer implementation. - */ - private static final byte[] KEY_COLON_BYTES = - ("\"" + DATADOG_CARRIER_KEY + "\":").getBytes(StandardCharsets.UTF_8); - - /** - * Byte pattern {@code \"_datadog\"} used to detect carriers that are embedded inside - * string-encoded JSON fields (e.g. an SNS {@code Message} field). - */ - private static final byte[] ESCAPED_KEY_BYTES = - ("\\\"" + DATADOG_CARRIER_KEY + "\\\"").getBytes(StandardCharsets.UTF_8); - - // Custom Object adapter preserves Long (not Double) for whole numbers and key order via - // LinkedHashMap, so re-serializing the envelope doesn't corrupt untouched fields. - private static final Moshi MOSHI = - new Moshi.Builder().add(Object.class, new NumberPreservingObjectAdapter()).build(); - - private static final Type STRING_OBJECT_MAP_TYPE = - Types.newParameterizedType(Map.class, String.class, Object.class); - - @SuppressWarnings("unchecked") - private static final JsonAdapter> MAP_ADAPTER = - (JsonAdapter>) (JsonAdapter) MOSHI.adapter(STRING_OBJECT_MAP_TYPE); - - private StripInjectedContext() {} - - /** Reads {@code in} fully, strips "_datadog", and returns a new stream over the result. */ - public static ByteArrayInputStream replaceInputStream(InputStream in) { - if (in == null) { - return null; - } - try { - return new ByteArrayInputStream(strip(readAllBytes(in))); - } catch (IOException e) { - log.debug("Unable to read Lambda payload for _datadog stripping", e); - return null; - } - } - - /** Strips "_datadog" from {@code payload} when the config is enabled. */ - public static byte[] strip(byte[] payload) { - if (!Config.get().isLambdaStripInjectedContextEnabled()) { - return payload; - } - return stripInternal(payload); - } - - /** Unconditional strip logic, package-private so tests can call it directly. */ - static byte[] stripInternal(byte[] payload) { - if (payload == null || payload.length == 0) { - return payload; - } - // Fast path: skip all processing when "_datadog" bytes are absent entirely. - if (indexOf(payload, 0, DATADOG_CARRIER_BYTES) < 0) { - return payload; - } - // Pass 1: splice out "_datadog":{...} object carriers via byte-level scanning. - // No JSON parsing or object allocation occurs here. - byte[] result = stripObjectCarriers(payload); - // Pass 2: strip "_datadog" from string-encoded fields (e.g. SNS Message). - // Uses Moshi only when the escaped key pattern is present. - result = stripStringEncodedCarriers(result); - return result; - } - - // ============================================================================ - // Object carrier stripping: byte-level splice (mirrors Go's stripObjectCarriers) - // ============================================================================ - - /** - * Removes all {@code "_datadog":{...}} object-form carriers from {@code msg} by splicing - * the matching byte ranges out of the array. Returns {@code msg} unchanged if no carriers - * are found. No JSON parsing or heap allocation occurs unless a carrier is actually present. - */ - private static byte[] stripObjectCarriers(byte[] msg) { - ByteArrayOutputStream out = null; - int searchFrom = 0; - int copyFrom = 0; - - while (searchFrom < msg.length) { - int[] carrier = findNextObjectCarrier(msg, searchFrom); - if (carrier == null) { - break; - } - int keyStart = carrier[0]; - int valueEnd = carrier[1]; - int[] removal = carrierRemovalRange(msg, keyStart, valueEnd, searchFrom); - - if (out == null) { - // Defer allocation until the first carrier is confirmed. Pre-size to the full - // input length so that multiple removals do not trigger internal buffer growth. - out = new ByteArrayOutputStream(msg.length); - } - out.write(msg, copyFrom, removal[0] - copyFrom); - copyFrom = removal[1]; - searchFrom = removal[1]; - } - - if (out == null) { - return msg; - } - out.write(msg, copyFrom, msg.length - copyFrom); - return out.toByteArray(); - } - - /** - * Returns {@code [keyStart, valueEnd]} for the next removable {@code "_datadog":{...}} - * carrier found at or after {@code searchFrom}, or {@code null} if none exists. - * - *

A carrier is removable only when (1) it is an actual JSON object property, not text - * embedded inside a string value, and (2) its value is an object (starts with {@code {}). - */ - private static int[] findNextObjectCarrier(byte[] msg, int searchFrom) { - while (searchFrom < msg.length) { - int keyStart = indexOf(msg, searchFrom, KEY_COLON_BYTES); - if (keyStart < 0) { - return null; - } - if (!isObjectProperty(msg, keyStart)) { - searchFrom = keyStart + 1; - continue; - } - int valueStart = objectValueStart(msg, keyStart + KEY_COLON_BYTES.length); - if (valueStart < 0) { - // Value is not a JSON object (e.g. "_datadog": null); skip this occurrence. - searchFrom = keyStart + 1; - continue; - } - int valueEnd = scanObjectEnd(msg, valueStart); - if (valueEnd < 0) { - // Malformed JSON: opening '{' has no matching '}'; skip safely. - searchFrom = keyStart + 1; - continue; - } - return new int[]{keyStart, valueEnd}; - } - return null; - } - - /** - * Returns {@code true} when the byte immediately before {@code keyStart} (ignoring JSON - * whitespace) is {@code {} or {@code ,}, confirming that this is a real JSON property key - * and not text that happens to appear inside a string value. - */ - private static boolean isObjectProperty(byte[] msg, int keyStart) { - int previous = skipWhitespaceBackward(msg, keyStart - 1); - return previous >= 0 && (msg[previous] == '{' || msg[previous] == ','); - } - - /** - * Returns the index of the opening {@code {} of the carrier value (skipping JSON whitespace - * after the colon), or {@code -1} if the value does not start with {@code {}. - */ - private static int objectValueStart(byte[] msg, int start) { - start = skipWhitespaceForward(msg, start); - return (start < msg.length && msg[start] == '{') ? start : -1; - } - - /** - * Returns the {@code [start, end)} byte range to remove for the carrier whose key begins at - * {@code keyStart} and whose value ends at {@code valueEnd} (exclusive). The range is expanded - * to consume one adjacent comma so the surrounding JSON remains valid. The preceding comma is - * preferred unless it falls before {@code searchFrom} (which belongs to a range already removed - * in an earlier iteration). - */ - private static int[] carrierRemovalRange( - byte[] msg, int keyStart, int valueEnd, int searchFrom) { - int previous = skipWhitespaceBackward(msg, keyStart - 1); - if (previous >= searchFrom && msg[previous] == ',') { - return new int[]{previous, valueEnd}; - } - int next = skipWhitespaceForward(msg, valueEnd); - if (next < msg.length && msg[next] == ',') { - return new int[]{keyStart, next + 1}; - } - // Sole field in its object: remove key and value without any comma. - return new int[]{keyStart, valueEnd}; - } - - /** - * Returns the exclusive end index of the JSON object that opens at {@code start}, tracking - * nested objects and quoted strings (including escape sequences) so that braces inside strings - * are not counted. Returns {@code -1} if no matching closing brace is found. - */ - private static int scanObjectEnd(byte[] msg, int start) { - int depth = 0; - boolean inString = false; - boolean escaped = false; - - for (int i = start; i < msg.length; i++) { - byte c = msg[i]; - - if (inString) { - if (escaped) { - escaped = false; - continue; - } - if (c == '\\') { - escaped = true; - continue; - } - if (c == '"') { - inString = false; - } - continue; - } - - switch (c) { - case '"': inString = true; break; - case '{': depth++; break; - case '}': - depth--; - if (depth == 0) { - return i + 1; - } - break; - default: break; - } - } - return -1; - } - - // ============================================================================ - // String-encoded carrier stripping: Moshi-based (for SNS Message and similar) - // ============================================================================ - - /** - * Removes {@code "_datadog"} from every top-level field whose value is a string-encoded JSON - * object (e.g. an SNS {@code Message} field). Uses Moshi only when {@code \"_datadog\"} is - * actually present in the bytes. Returns {@code msg} unchanged on any error. - */ - private static byte[] stripStringEncodedCarriers(byte[] msg) { - if (indexOf(msg, 0, ESCAPED_KEY_BYTES) < 0) { - return msg; - } - String json = new String(msg, StandardCharsets.UTF_8); - try { - Map envelope = MAP_ADAPTER.fromJson(json); - if (envelope == null) { - return msg; - } - boolean changed = false; - for (Map.Entry entry : envelope.entrySet()) { - if (entry.getValue() instanceof String) { - String stripped = stripFromJsonString((String) entry.getValue()); - if (stripped != null) { - entry.setValue(stripped); - changed = true; - } - } - } - if (!changed) { - return msg; - } - return MAP_ADAPTER.toJson(envelope).getBytes(StandardCharsets.UTF_8); - } catch (IOException | RuntimeException e) { - log.debug("Unable to strip string-encoded _datadog context from Lambda payload", e); - return msg; - } - } - - // ============================================================================ - // Byte-level utilities - // ============================================================================ - - /** - * Returns the start index of the first occurrence of {@code needle} in {@code haystack} at or - * after {@code from}, or {@code -1} if not found. - */ - private static int indexOf(byte[] haystack, int from, byte[] needle) { - if (needle.length == 0) { - return from; - } - int limit = haystack.length - needle.length; - outer: - for (int i = from; i <= limit; i++) { - for (int j = 0; j < needle.length; j++) { - if (haystack[i + j] != needle[j]) { - continue outer; - } - } - return i; - } - return -1; - } - - /** Returns the first index at or after {@code start} that is not JSON whitespace. */ - private static int skipWhitespaceForward(byte[] msg, int start) { - while (start < msg.length && isJSONWhitespace(msg[start])) { - start++; - } - return start; - } - - /** - * Returns the first index at or before {@code start} that is not JSON whitespace, - * or a negative value if only whitespace precedes. - */ - private static int skipWhitespaceBackward(byte[] msg, int start) { - while (start >= 0 && isJSONWhitespace(msg[start])) { - start--; - } - return start; - } - - /** Returns {@code true} if {@code c} is a JSON whitespace byte (space, tab, CR, or LF). */ - private static boolean isJSONWhitespace(byte c) { - return c == ' ' || c == '\t' || c == '\n' || c == '\r'; - } - - /** Parses, strips "_datadog", and re-encodes a JSON-object-shaped string. Null if unchanged. */ - private static String stripFromJsonString(String value) { - if (!value.contains(DATADOG_CARRIER_KEY)) { - return null; - } - try { - Map inner = MAP_ADAPTER.fromJson(value); - if (inner == null || inner.remove(DATADOG_CARRIER_KEY) == null) { - return null; - } - return MAP_ADAPTER.toJson(inner); - } catch (IOException | RuntimeException e) { - return null; - } - } - - /** Reads all bytes from {@code in}, restoring its position first if it's a byte stream. */ - private static byte[] readAllBytes(InputStream in) throws IOException { - if (in instanceof ByteArrayInputStream) { - // mark/reset instead of draining, in case the stream is read again elsewhere. - ByteArrayInputStream bais = (ByteArrayInputStream) in; - bais.mark(Integer.MAX_VALUE); - byte[] bytes = new byte[bais.available()]; - int read = bais.read(bytes); - bais.reset(); - return (read == bytes.length) ? bytes : Arrays.copyOf(bytes, Math.max(read, 0)); - } - ByteArrayOutputStream buffer = new ByteArrayOutputStream(); - byte[] chunk = new byte[8192]; - int n; - while ((n = in.read(chunk)) != -1) { - buffer.write(chunk, 0, n); - } - return buffer.toByteArray(); - } - - /** Object adapter that keeps whole numbers as {@link Long} and preserves key order. */ - private static final class NumberPreservingObjectAdapter extends JsonAdapter { - - /** Recursively reads a JSON value, mapping numbers to Long/Double based on their form. */ - @Override - public Object fromJson(JsonReader reader) throws IOException { - switch (reader.peek()) { - case BEGIN_ARRAY: - List list = new ArrayList<>(); - reader.beginArray(); - while (reader.hasNext()) { - list.add(fromJson(reader)); - } - reader.endArray(); - return list; - - case BEGIN_OBJECT: - Map map = new LinkedHashMap<>(); - reader.beginObject(); - while (reader.hasNext()) { - map.put(reader.nextName(), fromJson(reader)); - } - reader.endObject(); - return map; - - case STRING: - return reader.nextString(); - - case NUMBER: - return readNumber(reader); - - case BOOLEAN: - return reader.nextBoolean(); - - case NULL: - return reader.nextNull(); - - default: - throw new IOException("Unexpected JSON token: " + reader.peek()); - } - } - - /** Returns Long for integer literals, BigDecimal for fractional/scientific ones. */ - private Object readNumber(JsonReader reader) throws IOException { - String literal = reader.nextString(); // raw text avoids precision loss before parsing - if (literal.indexOf('.') >= 0 || literal.indexOf('e') >= 0 || literal.indexOf('E') >= 0) { - // BigDecimal preserves full precision instead of losing digits like Double would. - return new BigDecimal(literal); - } - try { - return Long.parseLong(literal); - } catch (NumberFormatException tooLargeForLong) { - return new BigDecimal(literal); - } - } - - /** Recursively writes a value produced by {@link #fromJson}. */ - @Override - public void toJson(JsonWriter writer, Object value) throws IOException { - if (value instanceof Map) { - writer.beginObject(); - for (Map.Entry entry : ((Map) value).entrySet()) { - writer.name(String.valueOf(entry.getKey())); - toJson(writer, entry.getValue()); - } - writer.endObject(); - } else if (value instanceof List) { - writer.beginArray(); - for (Object item : (List) value) { - toJson(writer, item); - } - writer.endArray(); - } else if (value instanceof String) { - writer.value((String) value); - } else if (value instanceof Long) { - writer.value(((Long) value).longValue()); - } else if (value instanceof Number) { - writer.value((Number) value); - } else if (value instanceof Boolean) { - writer.value((Boolean) value); - } else if (value == null) { - writer.nullValue(); - } else { - throw new IOException("Unsupported value type: " + value.getClass()); - } - } - } -} diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index ddf02ab64f3..540a4a2481a 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -29,7 +29,6 @@ import datadog.trace.context.NoopTraceScope; import datadog.trace.context.TraceScope; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; -import java.io.InputStream; import java.util.Collections; import java.util.List; import java.util.Map; @@ -385,8 +384,6 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera void notifyAppSecEnd(AgentSpan span, Object result); - InputStream stripLambdaInjectedContext(InputStream in); - AgentDataStreamsMonitoring getDataStreamsMonitoring(); String getTraceId(AgentSpan span); @@ -670,11 +667,6 @@ public void notifyExtensionEnd( @Override public void notifyAppSecEnd(AgentSpan span, Object result) {} - @Override - public InputStream stripLambdaInjectedContext(InputStream in) { - return null; - } - @Override public AgentDataStreamsMonitoring getDataStreamsMonitoring() { return NoopDataStreamsMonitoring.INSTANCE; From ad5f85298afc5f4756250cbd5cb8549e4264d1cd Mon Sep 17 00:00:00 2001 From: "eunhye.kim" Date: Tue, 13 Oct 2026 14:21:00 +0900 Subject: [PATCH 4/6] Make StripInjectedContext fuzz test reproducible with FUZZ_SEED and add Go parity cases --- .../aws-java-lambda-handler-1.2/build.gradle | 1 + .../gradle.lockfile | 173 ++++++------ .../java/StripInjectedContextFuzzTest.java | 260 ++++++++++++++++-- .../test/java/StripInjectedContextTest.java | 71 +++++ 4 files changed, 403 insertions(+), 102 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle index db5f54dbc63..a6bde8bce05 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/build.gradle @@ -17,4 +17,5 @@ dependencies { compileOnly group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1' testImplementation group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1' testImplementation project(':utils:test-utils') + testImplementation libs.jackson.databind } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/gradle.lockfile b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/gradle.lockfile index 4e36b31be18..9b11bd80363 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/gradle.lockfile +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/gradle.lockfile @@ -2,69 +2,75 @@ # Manual edits can break the build and are not advised. # This file is expected to be part of source control. # To regenerate this file, run: ./gradlew :dd-java-agent:instrumentation:aws-java:aws-java-lambda-handler-1.2:dependencies --write-locks -cafe.cryptography:curve25519-elisabeth:0.1.0=latestDepTestRuntimeClasspath,testRuntimeClasspath -cafe.cryptography:ed25519-elisabeth:0.1.0=latestDepTestRuntimeClasspath,testRuntimeClasspath -ch.qos.logback:logback-classic:1.2.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -ch.qos.logback:logback-core:1.2.13=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.amazonaws:aws-lambda-java-core:1.2.1=compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.blogspot.mydailyjava:weak-lock-free:0.17=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq.okhttp3:okhttp:3.12.15=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq.okio:okio:1.17.6=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:dd-instrument-java:0.0.6=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:dd-javac-plugin-client:0.2.2=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -com.datadoghq:java-dogstatsd-client:4.4.5=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.datadoghq:sketches-java:0.8.3=latestDepTestRuntimeClasspath,testRuntimeClasspath +cafe.cryptography:curve25519-elisabeth:0.1.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +cafe.cryptography:ed25519-elisabeth:0.1.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +ch.qos.logback:logback-classic:1.2.13=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +ch.qos.logback:logback-core:1.2.13=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.amazonaws:aws-lambda-java-core:1.2.1=compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.blogspot.mydailyjava:weak-lock-free:0.17=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq.okhttp3:okhttp:3.12.15=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq.okio:okio:1.17.6=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:dd-instrument-java:0.0.6=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:dd-javac-plugin-client:0.2.2=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.datadoghq:java-dogstatsd-client:4.4.5=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.datadoghq:sketches-java:0.8.3=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.fasterxml.jackson.core:jackson-annotations:2.20=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.fasterxml.jackson.core:jackson-core:2.20.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.fasterxml.jackson.core:jackson-databind:2.20.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.fasterxml.jackson:jackson-bom:2.20.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath com.github.javaparser:javaparser-core:3.25.6=codenarc -com.github.jnr:jffi:1.3.15=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-a64asm:1.0.0=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-constants:0.10.4=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-enxio:0.32.20=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-ffi:2.2.19=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-posix:3.1.22=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-unixsocket:0.38.25=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.jnr:jnr-x86asm:1.0.2=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.github.spotbugs:spotbugs-annotations:4.10.4=compileClasspath,spotbugs +com.github.jnr:jffi:1.3.15=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-a64asm:1.0.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-constants:0.10.4=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-enxio:0.32.20=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-ffi:2.2.19=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-posix:3.1.22=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-unixsocket:0.38.25=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.jnr:jnr-x86asm:1.0.2=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.github.spotbugs:spotbugs-annotations:4.10.4=compileClasspath,jmhCompileClasspath,spotbugs com.github.spotbugs:spotbugs:4.10.4=spotbugs com.github.stephenc.jcip:jcip-annotations:1.0-1=spotbugs -com.google.auto.service:auto-service-annotations:1.1.1=annotationProcessor,compileClasspath,latestDepTestAnnotationProcessor,latestDepTestCompileClasspath,testAnnotationProcessor,testCompileClasspath +com.google.auto.service:auto-service-annotations:1.1.1=annotationProcessor,compileClasspath,jmhCompileClasspath,latestDepTestAnnotationProcessor,latestDepTestCompileClasspath,testAnnotationProcessor,testCompileClasspath com.google.auto.service:auto-service:1.1.1=annotationProcessor,latestDepTestAnnotationProcessor,testAnnotationProcessor com.google.auto:auto-common:1.2.1=annotationProcessor,latestDepTestAnnotationProcessor,testAnnotationProcessor -com.google.code.findbugs:jsr305:3.0.2=annotationProcessor,compileClasspath,latestDepTestAnnotationProcessor,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,spotbugs,testAnnotationProcessor,testCompileClasspath,testRuntimeClasspath +com.google.code.findbugs:jsr305:3.0.2=annotationProcessor,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestAnnotationProcessor,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,spotbugs,testAnnotationProcessor,testCompileClasspath,testRuntimeClasspath com.google.code.gson:gson:2.14.0=spotbugs com.google.errorprone:error_prone_annotations:2.18.0=annotationProcessor,latestDepTestAnnotationProcessor,testAnnotationProcessor -com.google.errorprone:error_prone_annotations:2.47.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.google.errorprone:error_prone_annotations:2.47.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath com.google.errorprone:error_prone_annotations:2.48.0=spotbugs com.google.guava:failureaccess:1.0.1=annotationProcessor,latestDepTestAnnotationProcessor,testAnnotationProcessor -com.google.guava:failureaccess:1.0.3=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.google.guava:failureaccess:1.0.3=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath com.google.guava:guava:32.0.1-jre=annotationProcessor,latestDepTestAnnotationProcessor,testAnnotationProcessor -com.google.guava:guava:33.6.0-jre=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.google.guava:listenablefuture:9999.0-empty-to-avoid-conflict-with-guava=annotationProcessor,latestDepTestAnnotationProcessor,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testAnnotationProcessor,testCompileClasspath,testRuntimeClasspath +com.google.guava:guava:33.6.0-jre=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.google.guava:listenablefuture:9999.0-empty-to-avoid-conflict-with-guava=annotationProcessor,jmhRuntimeClasspath,latestDepTestAnnotationProcessor,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testAnnotationProcessor,testCompileClasspath,testRuntimeClasspath com.google.j2objc:j2objc-annotations:2.8=annotationProcessor,latestDepTestAnnotationProcessor,testAnnotationProcessor -com.google.j2objc:j2objc-annotations:3.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.google.re2j:re2j:1.8=latestDepTestRuntimeClasspath,testRuntimeClasspath -com.squareup.moshi:moshi:1.11.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.squareup.okhttp3:logging-interceptor:3.12.12=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.squareup.okhttp3:okhttp:3.12.12=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -com.squareup.okio:okio:1.17.5=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.google.j2objc:j2objc-annotations:3.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.google.re2j:re2j:1.8=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +com.squareup.moshi:moshi:1.11.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.squareup.okhttp3:logging-interceptor:3.12.12=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.squareup.okhttp3:okhttp:3.12.12=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +com.squareup.okio:okio:1.17.5=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath com.thoughtworks.qdox:qdox:1.12.1=codenarc -commons-fileupload:commons-fileupload:1.5=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -commons-io:commons-io:2.11.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +commons-fileupload:commons-fileupload:1.5=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +commons-io:commons-io:2.11.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath commons-io:commons-io:2.21.0=spotbugs -de.thetaphi:forbiddenapis:3.11=compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -io.leangen.geantyref:geantyref:1.3.16=latestDepTestRuntimeClasspath,testRuntimeClasspath -io.sqreen:libsqreen:17.5.0=latestDepTestRuntimeClasspath,testRuntimeClasspath -javax.servlet:javax.servlet-api:3.1.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +de.thetaphi:forbiddenapis:3.11=compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +io.leangen.geantyref:geantyref:1.3.16=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +io.sqreen:libsqreen:17.5.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +javax.servlet:javax.servlet-api:3.1.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath jaxen:jaxen:2.0.6=spotbugs -junit:junit:4.13.2=latestDepTestRuntimeClasspath,testRuntimeClasspath -net.bytebuddy:byte-buddy-agent:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -net.bytebuddy:byte-buddy:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -net.java.dev.jna:jna-platform:5.8.0=latestDepTestRuntimeClasspath,testRuntimeClasspath -net.java.dev.jna:jna:5.8.0=latestDepTestRuntimeClasspath,testRuntimeClasspath +junit:junit:4.13.2=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +net.bytebuddy:byte-buddy-agent:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +net.bytebuddy:byte-buddy:1.18.12=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +net.java.dev.jna:jna-platform:5.8.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +net.java.dev.jna:jna:5.8.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +net.sf.jopt-simple:jopt-simple:5.0.4=jmh,jmhCompileClasspath,jmhRuntimeClasspath net.sf.saxon:Saxon-HE:12.10=spotbugs org.apache.ant:ant-antlr:1.10.14=codenarc org.apache.ant:ant-junit:1.10.14=codenarc org.apache.bcel:bcel:6.12.0=spotbugs org.apache.commons:commons-lang3:3.20.0=spotbugs +org.apache.commons:commons-math3:3.2=jmh,jmhCompileClasspath,jmhRuntimeClasspath org.apache.commons:commons-text:1.15.0=spotbugs org.apache.logging.log4j:log4j-api:2.26.1=spotbugs org.apache.logging.log4j:log4j-core:2.26.1=spotbugs @@ -74,52 +80,57 @@ org.codehaus.groovy:groovy-ant:3.0.23=codenarc org.codehaus.groovy:groovy-docgenerator:3.0.23=codenarc org.codehaus.groovy:groovy-groovydoc:3.0.23=codenarc org.codehaus.groovy:groovy-json:3.0.23=codenarc -org.codehaus.groovy:groovy-json:3.0.25=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.codehaus.groovy:groovy-json:3.0.25=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.codehaus.groovy:groovy-templates:3.0.23=codenarc org.codehaus.groovy:groovy-xml:3.0.23=codenarc org.codehaus.groovy:groovy:3.0.23=codenarc -org.codehaus.groovy:groovy:3.0.25=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.codehaus.groovy:groovy:3.0.25=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.codenarc:CodeNarc:3.7.0=codenarc org.dom4j:dom4j:2.2.0=spotbugs org.gmetrics:GMetrics:2.1.0=codenarc -org.hamcrest:hamcrest-core:1.3=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.hamcrest:hamcrest:3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.jctools:jctools-core-jdk11:4.0.6=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.jctools:jctools-core:4.0.6=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.jspecify:jspecify:1.0.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.junit.jupiter:junit-jupiter-api:5.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.junit.jupiter:junit-jupiter-engine:5.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.junit.jupiter:junit-jupiter-params:5.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.junit.jupiter:junit-jupiter:5.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.junit.platform:junit-platform-commons:1.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.junit.platform:junit-platform-engine:1.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.junit.platform:junit-platform-launcher:1.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.junit.platform:junit-platform-runner:1.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.junit.platform:junit-platform-suite-api:1.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.junit.platform:junit-platform-suite-commons:1.14.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.junit:junit-bom:5.14.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.mockito:mockito-core:4.4.0=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.objenesis:objenesis:3.3=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.opentest4j:opentest4j:1.3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.hamcrest:hamcrest-core:1.3=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.hamcrest:hamcrest:3.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.jctools:jctools-core-jdk11:4.0.6=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.jctools:jctools-core:4.0.6=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.jspecify:jspecify:1.0.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.junit.jupiter:junit-jupiter-api:5.14.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.junit.jupiter:junit-jupiter-engine:5.14.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.junit.jupiter:junit-jupiter-params:5.14.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.junit.jupiter:junit-jupiter:5.14.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.junit.platform:junit-platform-commons:1.14.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.junit.platform:junit-platform-engine:1.14.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.junit.platform:junit-platform-launcher:1.14.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.junit.platform:junit-platform-runner:1.14.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.junit.platform:junit-platform-suite-api:1.14.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.junit.platform:junit-platform-suite-commons:1.14.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.junit:junit-bom:5.14.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.mockito:mockito-core:4.4.0=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.objenesis:objenesis:3.3=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.openjdk.jmh:jmh-core:1.36=jmh,jmhCompileClasspath,jmhRuntimeClasspath +org.openjdk.jmh:jmh-generator-asm:1.36=jmh,jmhCompileClasspath,jmhRuntimeClasspath +org.openjdk.jmh:jmh-generator-bytecode:1.36=jmh,jmhCompileClasspath,jmhRuntimeClasspath +org.openjdk.jmh:jmh-generator-reflection:1.36=jmh,jmhCompileClasspath,jmhRuntimeClasspath +org.opentest4j:opentest4j:1.3.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.ow2.asm:asm-analysis:9.10.1=spotbugs -org.ow2.asm:asm-analysis:9.7.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.ow2.asm:asm-commons:9.10.1=latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath -org.ow2.asm:asm-tree:9.10.1=latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath +org.ow2.asm:asm-analysis:9.7.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.ow2.asm:asm-commons:9.10.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath +org.ow2.asm:asm-tree:9.10.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,spotbugs,testRuntimeClasspath org.ow2.asm:asm-util:9.10.1=spotbugs -org.ow2.asm:asm-util:9.7.1=latestDepTestRuntimeClasspath,testRuntimeClasspath -org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath -org.slf4j:jcl-over-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.slf4j:jul-to-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.slf4j:log4j-over-slf4j:1.7.30=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.slf4j:slf4j-api:1.7.30=buildTimeInstrumentationPlugin,compileClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath -org.slf4j:slf4j-api:1.7.32=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.ow2.asm:asm-util:9.7.1=jmhRuntimeClasspath,latestDepTestRuntimeClasspath,testRuntimeClasspath +org.ow2.asm:asm:9.0=jmh +org.ow2.asm:asm:9.10.1=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath +org.slf4j:jcl-over-slf4j:1.7.30=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.slf4j:jul-to-slf4j:1.7.30=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.slf4j:log4j-over-slf4j:1.7.30=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.slf4j:slf4j-api:1.7.30=buildTimeInstrumentationPlugin,compileClasspath,jmhCompileClasspath,muzzleBootstrap,muzzleTooling,runtimeClasspath +org.slf4j:slf4j-api:1.7.32=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.slf4j:slf4j-api:2.0.17=spotbugsSlf4j org.slf4j:slf4j-api:2.0.18=spotbugs org.slf4j:slf4j-simple:2.0.17=spotbugsSlf4j -org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath -org.spockframework:spock-bom:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.spockframework:spock-core:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.tabletest:tabletest-junit:1.2.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath -org.tabletest:tabletest-parser:1.2.1=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,jmhRuntimeClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath +org.spockframework:spock-bom:2.4-groovy-3.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.spockframework:spock-core:2.4-groovy-3.0=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.tabletest:tabletest-junit:1.2.2=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath +org.tabletest:tabletest-parser:1.2.1=jmhRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath org.xmlresolver:xmlresolver:5.3.3=spotbugs -empty=spotbugsPlugins +empty=jmhAnnotationProcessor,spotbugsPlugins diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java index c17f2b3c75f..3b665d14d00 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextFuzzTest.java @@ -1,13 +1,23 @@ package datadog.trace.instrumentation.aws.v1.lambda; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.Arrays; -import java.util.concurrent.ThreadLocalRandom; +import java.util.Random; +import java.util.stream.Stream; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; /** * Randomized invariant testing for {@link StripInjectedContext}, mirroring the repo's existing @@ -17,37 +27,159 @@ */ public final class StripInjectedContextFuzzTest { - private static final int ITERATIONS = 5_000; + private static final long SEED = envLong("FUZZ_SEED", System.nanoTime()); + private static final int ITERATIONS = (int) envLong("FUZZ_ITERATIONS", 5_000); + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static final String[] SEEDS = { + "{\"detail\":{\"_datadog\":{\"x-datadog-trace-id\":\"1\"},\"foo\":\"bar\"}}", + "{\"detail\":\"{\\\"foo\\\":\\\"bar\\\",\\\"_datadog\\\":{\\\"k\\\":\\\"v\\\"}}\"}", + "{\"foo\":\"a,}b\",\"_datadog\":{\"x\":1}}", + "{\"detail\":\"{\\\"foo\\\":\\\"bar\\\",\\\"_datadog\\\":{\\\"k\\\":\\\"v}z\\\"}}\"}", + "{\"_datadog\":{},\"_datadog\":{}}", + "not even json", + "{}", + "" + }; + + /** Inputs that once failed. Paste the input from a failure message here so it always runs. */ + private static final String[] REGRESSIONS = {}; + + @BeforeAll + static void printSeed() { + System.out.println( + "StripInjectedContextFuzzTest FUZZ_SEED=" + SEED + " FUZZ_ITERATIONS=" + ITERATIONS); + } + + static Stream corpus() { + return Stream.concat(Arrays.stream(SEEDS), Arrays.stream(REGRESSIONS)); + } + + @ParameterizedTest + @MethodSource("corpus") + void seedCorpusHoldsInvariants(String input) { + checkInvariants(input, ""); + } + + private static final String[] FRAGMENTS = { + "{", + "}", + "[", + "]", + "\"", + ",", + ":", + "\\", + " ", + "\\\"", + "_datadog", + "\"_datadog\":", + "\"_datadog\":{}", + "\\u00e9", + "1" + }; + + private static final String[] STRING_PARTS = { + "a", "}", "{", ",", "\\\"", "\\\\", ":", " ", "_datadog", "\\u00e9" + }; + + // A "_datadog" key inside the "other" value would itself be stripped, so keys there avoid it. + private static final String[] OTHER_KEY_PARTS = { + "a", "}", "{", ",", "\\\"", "\\\\", ":", " ", "x_datadog", "\\u00e9" + }; + + private static void checkInvariants(String input, String where) { + byte[] in = input.getBytes(StandardCharsets.UTF_8); + byte[] out; + try { + out = StripInjectedContext.stripInternal(in); + } catch (Throwable t) { + fail(where + "stripInternal threw for in=" + input, t); + return; + } + // Malformed input is only guaranteed not to throw: cutting a carrier out of it can splice + // the surrounding bytes into a new carrier, so it is not necessarily idempotent. + if (isValidJson(in)) { + assertTrue( + isValidJson(out), where + "valid JSON became invalid: in=" + input + " out=" + utf8(out)); + assertArrayEquals( + out, StripInjectedContext.stripInternal(out), where + "not idempotent: in=" + input); + } + } + + private static String where(int iteration) { + return "[FUZZ_SEED=" + SEED + " iteration=" + iteration + "] "; + } + + private static long envLong(String name, long defaultValue) { + String value = System.getenv(name); + return value == null || value.isEmpty() ? defaultValue : Long.parseLong(value); + } @Test void neverThrowsAndAlwaysRemovesDatadogWhenPresent() { - ThreadLocalRandom random = ThreadLocalRandom.current(); + Random random = new Random(SEED); for (int i = 0; i < ITERATIONS; i++) { byte[] payload = randomPayload(random); + String input = utf8(payload); byte[] result; try { result = StripInjectedContext.stripInternal(payload); } catch (Throwable t) { - fail( - "stripInternal must never throw, but threw for payload: " - + new String(payload, StandardCharsets.UTF_8), - t); + fail(where(i) + "stripInternal threw for in=" + input, t); return; } - String resultJson = new String(result, StandardCharsets.UTF_8); + String resultJson = utf8(result); assertFalse( - resultJson.contains("_datadog"), "output must never contain _datadog: " + resultJson); + resultJson.contains("_datadog"), + where(i) + "output contains _datadog: in=" + input + " out=" + resultJson); + assertArrayEquals( + result, + StripInjectedContext.stripInternal(result), + where(i) + "not idempotent: in=" + input); + } + } + + @Test + void mutatedSeedsKeepValidJsonValidAndAreIdempotent() { + Random random = new Random(SEED); - // Idempotency: stripping an already-stripped payload must be a no-op. - byte[] second = StripInjectedContext.stripInternal(result); - assertTrue(Arrays.equals(result, second), "stripInternal must be idempotent"); + for (int i = 0; i < ITERATIONS; i++) { + checkInvariants(mutate(SEEDS[random.nextInt(SEEDS.length)], random), where(i)); } } - private static byte[] randomPayload(ThreadLocalRandom random) { + @Test + void arbitraryCarrierIsAlwaysRemovedAndOtherFieldSurvives() throws IOException { + Random random = new Random(SEED); + + for (int i = 0; i < ITERATIONS; i++) { + String carrier = randomObject(random, 0, STRING_PARTS); + String other = randomValue(random, 0, OTHER_KEY_PARTS); + String payload = "{\"_datadog\":" + carrier + ",\"other\":" + other + "}"; + + byte[] out = StripInjectedContext.stripInternal(payload.getBytes(StandardCharsets.UTF_8)); + + JsonNode result; + try { + result = MAPPER.readTree(out); + } catch (IOException e) { + fail(where(i) + "output is not valid JSON: in=" + payload + " out=" + utf8(out), e); + return; + } + assertFalse( + result.has("_datadog"), + where(i) + "carrier was not removed: in=" + payload + " out=" + utf8(out)); + assertEquals( + MAPPER.readTree(other), result.get("other"), where(i) + "other changed: in=" + payload); + } + } + + private static byte[] randomPayload(Random random) { switch (random.nextInt(3)) { case 0: return randomEventBridgePayload(random); @@ -59,7 +191,7 @@ private static byte[] randomPayload(ThreadLocalRandom random) { } // EventBridge-shaped: {"detail-type": ..., "detail": {...}} or {"detail": "..."} - private static byte[] randomEventBridgePayload(ThreadLocalRandom random) { + private static byte[] randomEventBridgePayload(Random random) { StringBuilder sb = new StringBuilder(); sb.append("{\"detail-type\":\"order.created\",\"detail\":"); boolean asString = random.nextBoolean(); @@ -79,13 +211,13 @@ private static byte[] randomEventBridgePayload(ThreadLocalRandom random) { // Randomly places _datadog as the sole field, the first field, a middle field, // or the last field so that all four comma-placement cases in carrierRange // are exercised across the fuzz iterations. - private static byte[] randomTopLevelCarrierPayload(ThreadLocalRandom random) { + private static byte[] randomTopLevelCarrierPayload(Random random) { boolean includeDatadog = random.nextBoolean(); // Allow 0 regular fields so the sole-field case ({"_datadog":{...}}) is also generated. - int fieldCount = random.nextInt(0, 4); + int fieldCount = random.nextInt(4); // Choose a random insertion position for the _datadog key among the regular fields. // Position 0 = first, fieldCount = last, values in between = middle. - int datadogPos = includeDatadog ? random.nextInt(0, fieldCount + 1) : fieldCount + 1; + int datadogPos = includeDatadog ? random.nextInt(fieldCount + 1) : fieldCount + 1; StringBuilder sb = new StringBuilder("{"); boolean needsComma = false; @@ -114,9 +246,9 @@ private static byte[] randomTopLevelCarrierPayload(ThreadLocalRandom random) { return sb.toString().getBytes(StandardCharsets.UTF_8); } - private static String randomDetailObject(ThreadLocalRandom random, boolean includeDatadog) { + private static String randomDetailObject(Random random, boolean includeDatadog) { StringBuilder sb = new StringBuilder("{"); - int fieldCount = random.nextInt(1, 6); + int fieldCount = 1 + random.nextInt(5); for (int i = 0; i < fieldCount; i++) { if (i > 0) { sb.append(','); @@ -135,7 +267,7 @@ private static String randomDetailObject(ThreadLocalRandom random, boolean inclu } // SNS-shaped: a top-level string field contains a string-encoded JSON object with _datadog. - private static byte[] randomStringEncodedFieldPayload(ThreadLocalRandom random) { + private static byte[] randomStringEncodedFieldPayload(Random random) { boolean includeDatadog = random.nextBoolean(); String innerJson = randomDetailObject(random, includeDatadog); String escaped = innerJson.replace("\"", "\\\""); @@ -144,7 +276,7 @@ private static byte[] randomStringEncodedFieldPayload(ThreadLocalRandom random) return payload.getBytes(StandardCharsets.UTF_8); } - private static void appendRandomValue(StringBuilder sb, ThreadLocalRandom random) { + private static void appendRandomValue(StringBuilder sb, Random random) { switch (random.nextInt(4)) { case 0: sb.append(random.nextInt()); @@ -159,4 +291,90 @@ private static void appendRandomValue(StringBuilder sb, ThreadLocalRandom random sb.append(random.nextBoolean()); } } + + /** Applies up to three random insert, delete, or replace edits using JSON-significant text. */ + private static String mutate(String seed, Random random) { + StringBuilder sb = new StringBuilder(seed); + int edits = random.nextInt(4); + for (int e = 0; e < edits; e++) { + int pos = random.nextInt(sb.length() + 1); + switch (random.nextInt(3)) { + case 0: + sb.insert(pos, FRAGMENTS[random.nextInt(FRAGMENTS.length)]); + break; + case 1: + if (pos < sb.length()) { + sb.deleteCharAt(pos); + } + break; + default: + if (pos < sb.length()) { + sb.setCharAt(pos, FRAGMENTS[random.nextInt(FRAGMENTS.length)].charAt(0)); + } + } + } + return sb.toString(); + } + + private static String randomValue(Random random, int depth, String[] keyParts) { + switch (random.nextInt(depth < 3 ? 7 : 5)) { + case 0: + return Long.toString(random.nextLong()); + case 1: + return random.nextBoolean() ? "true" : "null"; + case 2: + return randomString(random, STRING_PARTS); + case 3: + return "1.50"; + case 4: + return "9007199254740993"; + case 5: + return randomObject(random, depth + 1, keyParts); + default: + StringBuilder sb = new StringBuilder("["); + for (int k = 0, n = random.nextInt(3); k < n; k++) { + sb.append(k > 0 ? "," : "").append(randomValue(random, depth + 1, keyParts)); + } + return sb.append(']').toString(); + } + } + + private static String randomObject(Random random, int depth, String[] keyParts) { + StringBuilder sb = new StringBuilder("{"); + for (int k = 0, n = random.nextInt(4); k < n; k++) { + sb.append(k > 0 ? "," : "") + .append(randomString(random, keyParts)) + .append(':') + .append(randomValue(random, depth + 1, keyParts)); + } + return sb.append('}').toString(); + } + + private static String randomString(Random random, String[] parts) { + StringBuilder sb = new StringBuilder("\""); + for (int k = 0, n = random.nextInt(5); k < n; k++) { + sb.append(parts[random.nextInt(parts.length)]); + } + return sb.append('"').toString(); + } + + /** Strict check: exactly one JSON value with nothing but whitespace after it. */ + private static boolean isValidJson(byte[] json) { + if (json.length == 0) { + return false; + } + try (JsonParser parser = MAPPER.getFactory().createParser(json)) { + if (parser.nextToken() == null) { + return false; + } + parser.skipChildren(); + return parser.nextToken() == null; + } catch (IOException e) { + return false; + } + } + + private static String utf8(byte[] bytes) { + return new String(bytes, StandardCharsets.UTF_8); + } } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java index dd131a2fd50..8d06c30dd97 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.aws.v1.lambda; import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; @@ -12,7 +13,11 @@ import java.io.ByteArrayInputStream; import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.util.stream.Stream; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; /** Unit tests for {@link StripInjectedContext}. */ class StripInjectedContextTest extends DDJavaSpecification { @@ -525,4 +530,70 @@ void unclosedCarriersDoNotCauseQuadraticScan() { Duration.ofSeconds(1), () -> assertArrayEquals(input, StripInjectedContext.stripInternal(input))); } + + /** + * Keeps behavior aligned with the dd-trace-go implementation. The only intended difference is + * that whitespace before the colon does not prevent removal. + */ + static Stream sharedCasesWithGo() { + return Stream.of( + Arguments.of( + "no datadog key", "{\"detail\":{\"foo\":\"bar\"}}", "{\"detail\":{\"foo\":\"bar\"}}"), + Arguments.of("empty object", "{}", "{}"), + Arguments.of("empty input", "", ""), + Arguments.of("invalid JSON", "not even json", "not even json"), + Arguments.of( + "SQS event", + "{\"Records\":[{\"body\":\"hello\"}]}", + "{\"Records\":[{\"body\":\"hello\"}]}"), + Arguments.of( + "object detail alone", "{\"detail\":{\"_datadog\":{\"x\":\"y\"}}}", "{\"detail\":{}}"), + Arguments.of( + "object detail at beginning", + "{\"detail\":{\"_datadog\":{\"x\":\"y\"},\"foo\":\"bar\"}}", + "{\"detail\":{\"foo\":\"bar\"}}"), + Arguments.of( + "object detail at end", + "{\"detail\":{\"foo\":\"bar\",\"_datadog\":{\"x\":\"y\"}}}", + "{\"detail\":{\"foo\":\"bar\"}}"), + Arguments.of( + "object detail in middle", + "{\"detail\":{\"foo\":\"bar\",\"_datadog\":{\"x\":\"y\"},\"baz\":\"qux\"}}", + "{\"detail\":{\"foo\":\"bar\",\"baz\":\"qux\"}}"), + Arguments.of("adjacent carriers", "{\"_datadog\":{},\"_datadog\":{}}", "{}"), + Arguments.of( + "value with comma and brace", + "{\"foo\":\"a,}b\",\"_datadog\":{\"x\":1}}", + "{\"foo\":\"a,}b\"}"), + Arguments.of("_datadog value null", "{\"_datadog\":null}", "{\"_datadog\":null}"), + Arguments.of("_datadog value bool", "{\"_datadog\":true}", "{\"_datadog\":true}"), + Arguments.of("_datadog value number", "{\"_datadog\":42}", "{\"_datadog\":42}"), + Arguments.of("_datadog value string", "{\"_datadog\":\"str\"}", "{\"_datadog\":\"str\"}"), + Arguments.of("_datadog value array", "{\"_datadog\":[1,2,3]}", "{\"_datadog\":[1,2,3]}"), + Arguments.of( + "similar key prefix", "{\"x_datadog\":{\"a\":\"b\"}}", "{\"x_datadog\":{\"a\":\"b\"}}"), + Arguments.of( + "similar key suffix", "{\"_datadogx\":{\"a\":\"b\"}}", "{\"_datadogx\":{\"a\":\"b\"}}"), + Arguments.of("whitespace before colon", "{\"_datadog\" :{\"x\":\"y\"}}", "{}"), + Arguments.of( + "even backslashes in value", + "{\"foo\":\"bar\\\\\\\\\",\"_datadog\":{\"x\":\"y\"}}", + "{\"foo\":\"bar\\\\\\\\\"}"), + Arguments.of( + "odd backslashes in carrier value", "{\"_datadog\":{\"x\":\"val\\\\\\\"\"}}", "{}"), + Arguments.of( + "string detail", + "{\"detail\":\"{\\\"foo\\\":\\\"bar\\\",\\\"_datadog\\\":{\\\"k\\\":\\\"v\\\"}}\"}", + "{\"detail\":\"{\\\"foo\\\":\\\"bar\\\"}\"}"), + Arguments.of( + "string detail with brace in carrier", + "{\"detail\":\"{\\\"foo\\\":\\\"bar\\\",\\\"_datadog\\\":{\\\"k\\\":\\\"v}z\\\"}}\"}", + "{\"detail\":\"{\\\"foo\\\":\\\"bar\\\"}\"}")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("sharedCasesWithGo") + void matchesExpectedOutput(String name, String input, String expected) { + assertEquals(expected, toUtf8(StripInjectedContext.stripInternal(bytes(input)))); + } } From 7fd141b4611477e4854977e13f9f32b68621ed05 Mon Sep 17 00:00:00 2001 From: "eunhye.kim" Date: Tue, 13 Oct 2026 15:44:24 +0900 Subject: [PATCH 5/6] Preserve keys, unsupported streams, invalid UTF-8 and free text when stripping _datadog --- .../aws/v1/lambda/StripInjectedContext.java | 75 +++++++------- .../test/java/StripInjectedContextTest.java | 97 +++++++++++++++++++ 2 files changed, 135 insertions(+), 37 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java index f283a4be0db..0042a30c673 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/StripInjectedContext.java @@ -4,7 +4,6 @@ import datadog.json.JsonReader; import datadog.trace.api.Config; import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.nio.charset.StandardCharsets; @@ -18,10 +17,10 @@ * *

Carriers are cut out of the JSON text instead of parsing and re-serializing the payload, so * numbers, key order and formatting of customer data are kept. Carriers inside string-encoded JSON - * (an SNS {@code Message}, an SQS {@code body}) are handled by decoding that string, stripping it - * recursively up to {@link #MAX_DEPTH} levels, and re-encoding it. A {@code _datadog} key whose - * value is not an object is customer data and is kept. Fail-open: on any error the original payload - * is used. + * values (an SNS {@code Message}, an SQS {@code body}) are handled by decoding that string, + * stripping it recursively up to {@link #MAX_DEPTH} levels, and re-encoding it. Object keys are + * never rewritten, and a {@code _datadog} key whose value is not an object is customer data and is + * kept. Fail-open: on any error the original payload is used. */ public final class StripInjectedContext { @@ -39,18 +38,16 @@ public final class StripInjectedContext { private StripInjectedContext() {} - /** Reads {@code in} fully, strips "_datadog", and returns a new stream over the result. */ + /** + * Returns a stream over the stripped payload, or {@code null} to keep {@code in}. Like the other + * Lambda event consumers, only {@link ByteArrayInputStream} is supported, because it can be read + * with mark/reset without consuming the customer's input. + */ public static ByteArrayInputStream replaceInputStream(InputStream in) { - if (in == null) { - return null; - } - byte[] original; - try { - original = readAllBytes(in); - } catch (IOException e) { - log.debug("Unable to read Lambda payload for _datadog stripping", e); + if (!(in instanceof ByteArrayInputStream)) { return null; } + byte[] original = readAllBytes((ByteArrayInputStream) in); try { return new ByteArrayInputStream(strip(original)); } catch (RuntimeException e) { @@ -73,14 +70,18 @@ static byte[] stripInternal(byte[] payload) { return payload; } String json = new String(payload, StandardCharsets.UTF_8); + // Invalid UTF-8 decodes to U+FFFD and would be re-encoded as such, so keep the payload as is. + if (json.indexOf('\uFFFD') >= 0) { + return payload; + } String stripped = stripJson(json, 0); return stripped.equals(json) ? payload : stripped.getBytes(StandardCharsets.UTF_8); } /** * Walks the string tokens of {@code json} once. A {@code "_datadog"} key with an object value is - * cut out; any other string containing {@code _datadog} is treated as encoded JSON and stripped - * recursively. Returns {@code json} itself when nothing changed. + * cut out; any other string value containing {@code _datadog} is treated as encoded JSON and + * stripped recursively. Returns {@code json} itself when nothing changed. */ private static String stripJson(String json, int depth) { StringBuilder out = null; @@ -109,7 +110,7 @@ private static String stripJson(String json, int depth) { out.append(json, copyFrom, carrier[0]); copyFrom = carrier[1]; end = carrier[1]; - } else if (nextKey < end && depth < MAX_DEPTH) { + } else if (nextKey < end && depth < MAX_DEPTH && !isKey(json, end)) { String replacement = stripEncodedJson(json.substring(start, end), depth); if (replacement != null) { if (out == null) { @@ -162,8 +163,9 @@ private static int[] carrierRange(String json, int keyStart, int keyEnd, int cop } /** - * Decodes the JSON string literal {@code literal}, strips it if it holds a JSON object, and - * returns the re-encoded literal, or {@code null} if nothing changed. + * Decodes the JSON string literal {@code literal}, strips it if it looks like a JSON object + * (starts with '{' and ends with '}'), and returns the re-encoded literal, or {@code null} if + * nothing changed. */ private static String stripEncodedJson(String literal, int depth) { String decoded; @@ -176,6 +178,10 @@ private static String stripEncodedJson(String literal, int depth) { if (first >= decoded.length() || decoded.charAt(first) != '{') { return null; } + int last = skipWhitespaceBackward(decoded, decoded.length() - 1); + if (decoded.charAt(last) != '}') { + return null; + } String stripped = stripJson(decoded, depth + 1); return stripped.equals(decoded) ? null : JsonMapper.toJson(stripped); } @@ -213,6 +219,12 @@ private static int objectEnd(String json, int start) { return -1; } + /** Object keys are customer property names and are never rewritten. */ + private static boolean isKey(String json, int stringEnd) { + int next = skipWhitespace(json, stringEnd); + return next < json.length() && json.charAt(next) == ':'; + } + private static int skipWhitespace(String json, int i) { while (i < json.length() && isWhitespace(json.charAt(i))) { i++; @@ -244,23 +256,12 @@ private static boolean containsCarrierKey(byte[] payload) { return false; } - /** Reads all bytes from {@code in}, restoring its position first if it's a byte stream. */ - private static byte[] readAllBytes(InputStream in) throws IOException { - if (in instanceof ByteArrayInputStream) { - // mark/reset instead of draining, in case the stream is read again elsewhere. - ByteArrayInputStream bais = (ByteArrayInputStream) in; - bais.mark(Integer.MAX_VALUE); - byte[] bytes = new byte[bais.available()]; - int read = bais.read(bytes); - bais.reset(); - return (read == bytes.length) ? bytes : Arrays.copyOf(bytes, Math.max(read, 0)); - } - ByteArrayOutputStream buffer = new ByteArrayOutputStream(); - byte[] chunk = new byte[8192]; - int n; - while ((n = in.read(chunk)) != -1) { - buffer.write(chunk, 0, n); - } - return buffer.toByteArray(); + /** Reads all bytes from {@code in} and restores its position, so the stream is not consumed. */ + private static byte[] readAllBytes(ByteArrayInputStream in) { + in.mark(Integer.MAX_VALUE); + byte[] bytes = new byte[in.available()]; + int read = in.read(bytes, 0, bytes.length); + in.reset(); + return read == bytes.length ? bytes : Arrays.copyOf(bytes, Math.max(read, 0)); } } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java index 8d06c30dd97..837af133305 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/StripInjectedContextTest.java @@ -11,8 +11,11 @@ import datadog.trace.test.junit.utils.config.WithConfig; import datadog.trace.test.util.DDJavaSpecification; import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -375,6 +378,100 @@ void replaceInputStreamStripsAndDoesNotConsumeTheOriginalStream() { assertTrue(originalJson.contains("_datadog"), "original stream must still be fully readable"); } + @Test + @WithConfig(key = GeneralConfig.LAMBDA_STRIP_INJECTED_CONTEXT, value = "true") + void replaceInputStreamDoesNotReadUnsupportedStreams() { + byte[] payload = bytes("{\"a\":1,\"_datadog\":{\"t\":\"1\"}}"); + AtomicInteger reads = new AtomicInteger(); + InputStream failsAfterPrefix = + new InputStream() { + private int pos; + + @Override + public int read() throws IOException { + reads.incrementAndGet(); + if (pos >= 5) { + throw new IOException("stream failed after a prefix"); + } + return payload[pos++]; + } + }; + + assertNull(StripInjectedContext.replaceInputStream(failsAfterPrefix)); + assertEquals(0, reads.get(), "an unsupported stream must be passed through unread"); + } + + @Test + void doesNotRewriteObjectKeysThatLookLikeEncodedJson() { + // {"{\"_datadog\":{}}":1,"{}":2} must not become {"{}":1,"{}":2} + byte[] input = bytes("{\"{\\\"_datadog\\\":{}}\":1,\"{}\":2}"); + + byte[] result = StripInjectedContext.stripInternal(input); + + assertArrayEquals(input, result, "object keys must never be rewritten: " + toUtf8(result)); + } + + @Test + void keepsPayloadWithInvalidUtf8Untouched() { + // byte 6 is the 'X' value; 0xFF is never valid in UTF-8 + byte[] input = bytes("{\"a\":\"X\",\"_datadog\":{}}"); + input[6] = (byte) 0xFF; + + assertArrayEquals(input, StripInjectedContext.stripInternal(input)); + } + + @Test + void keepsPayloadWithLiteralReplacementCharacterUntouched() { + // A raw U+FFFD can't be told apart from a decoded invalid byte, so the payload is kept as is. + byte[] input = bytes("{\"a\":\"\uFFFD\",\"_datadog\":{}}"); + + assertArrayEquals(input, StripInjectedContext.stripInternal(input)); + } + + @Test + void stillStripsPayloadWithEscapedReplacementCharacter() { + String result = + toUtf8(StripInjectedContext.stripInternal(bytes("{\"a\":\"\\uFFFD\",\"_datadog\":{}}"))); + + assertEquals("{\"a\":\"\\uFFFD\"}", result); + } + + @Test + void stillStripsPayloadWithValidMultiByteUtf8() { + String result = + toUtf8( + StripInjectedContext.stripInternal(bytes("{\"a\":\"\u00e9\u4e2d\",\"_datadog\":{}}"))); + + assertEquals("{\"a\":\"\u00e9\u4e2d\"}", result); + } + + @Test + void doesNotTreatFreeTextStartingWithBraceAsEncodedJson() { + byte[] input = bytes("{\"note\":\"{ pasted \\\"_datadog\\\":{x} here\"}"); + + assertArrayEquals(input, StripInjectedContext.stripInternal(input)); + } + + @Test + void stillStripsEncodedJsonWithSurroundingWhitespace() { + String result = + toUtf8( + StripInjectedContext.stripInternal( + bytes("{\"m\":\" {\\\"a\\\":1,\\\"_datadog\\\":{}} \"}"))); + + assertEquals("{\"m\":\" {\\\"a\\\":1} \"}", result); + } + + @Test + void stripsEncodedValueButKeepsEncodedLookingKeyInSameObject() { + String key = "\"{\\\"_datadog\\\":{}}\""; + byte[] input = bytes("{" + key + ":\"{\\\"a\\\":1,\\\"_datadog\\\":{\\\"t\\\":\\\"1\\\"}}\"}"); + + String result = toUtf8(StripInjectedContext.stripInternal(input)); + + assertEquals("{" + key + ":\"{\\\"a\\\":1}\"}", result); + } + @Test void doesNotRemoveNonObjectDatadogValueInsideStringEncodedField() { // Regression test for review feedback: a string-encoded business payload (e.g. SQS body) From 3347b6b38ba8bc4c46938aea49347268ec1f9b0b Mon Sep 17 00:00:00 2001 From: "eunhye.kim" Date: Thu, 15 Oct 2026 09:58:34 +0900 Subject: [PATCH 6/6] add DD_TRACE_LAMBDA_STRIP_INJECTED_CONTEXT as an alias --- .../java/datadog/trace/api/config/GeneralConfig.java | 2 ++ .../src/main/java/datadog/trace/api/Config.java | 5 ++++- .../src/test/java/datadog/trace/api/ConfigTest.java | 11 +++++++++++ metadata/supported-configurations.json | 2 +- 4 files changed, 18 insertions(+), 2 deletions(-) diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java index a17dd65e580..b3f05eba9be 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java @@ -141,6 +141,8 @@ public final class GeneralConfig { public static final String APP_LOGS_COLLECTION_ENABLED = "app.logs.collection.enabled"; public static final String LAMBDA_STRIP_INJECTED_CONTEXT = "lambda.strip.injected.context"; + public static final String TRACE_LAMBDA_STRIP_INJECTED_CONTEXT = + "trace.lambda.strip.injected.context"; private GeneralConfig() {} } diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index 39a72172d75..377b5c12359 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -440,6 +440,7 @@ import static datadog.trace.api.config.GeneralConfig.TRACER_METRICS_MAX_AGGREGATES; import static datadog.trace.api.config.GeneralConfig.TRACER_METRICS_MAX_PENDING; import static datadog.trace.api.config.GeneralConfig.TRACE_DEBUG; +import static datadog.trace.api.config.GeneralConfig.TRACE_LAMBDA_STRIP_INJECTED_CONTEXT; import static datadog.trace.api.config.GeneralConfig.TRACE_LOG_LEVEL; import static datadog.trace.api.config.GeneralConfig.TRACE_OTEL_SEMANTICS_ENABLED; import static datadog.trace.api.config.GeneralConfig.TRACE_STATS_ADDITIONAL_TAGS; @@ -1782,7 +1783,9 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins lambdaStripInjectedContextEnabled = configProvider.getBoolean( - LAMBDA_STRIP_INJECTED_CONTEXT, DEFAULT_LAMBDA_STRIP_INJECTED_CONTEXT); + LAMBDA_STRIP_INJECTED_CONTEXT, + DEFAULT_LAMBDA_STRIP_INJECTED_CONTEXT, + TRACE_LAMBDA_STRIP_INJECTED_CONTEXT); spanAttributeSchemaVersion = schemaVersionFromConfig(); diff --git a/internal-api/src/test/java/datadog/trace/api/ConfigTest.java b/internal-api/src/test/java/datadog/trace/api/ConfigTest.java index 777759f9e79..0e8d255ca45 100644 --- a/internal-api/src/test/java/datadog/trace/api/ConfigTest.java +++ b/internal-api/src/test/java/datadog/trace/api/ConfigTest.java @@ -62,6 +62,7 @@ import static datadog.trace.api.config.GeneralConfig.SSI_INJECTION_FORCE; import static datadog.trace.api.config.GeneralConfig.TAGS; import static datadog.trace.api.config.GeneralConfig.TRACER_METRICS_IGNORED_RESOURCES; +import static datadog.trace.api.config.GeneralConfig.TRACE_LAMBDA_STRIP_INJECTED_CONTEXT; import static datadog.trace.api.config.GeneralConfig.TRACE_OTEL_SEMANTICS_ENABLED; import static datadog.trace.api.config.GeneralConfig.VERSION; import static datadog.trace.api.config.JmxFetchConfig.JMX_FETCH_CHECK_PERIOD; @@ -3028,6 +3029,16 @@ void lambdaStripInjectedContextCanBeEnabled() throws Exception { assertTrue(config.isLambdaStripInjectedContextEnabled()); } + @Test + void lambdaStripInjectedContextCanBeEnabledWithTraceAlias() throws Exception { + Properties properties = new Properties(); + properties.setProperty(TRACE_LAMBDA_STRIP_INJECTED_CONTEXT, "true"); + + Config config = newConfig(withPropertiesOverride(properties)); + + assertTrue(config.isLambdaStripInjectedContextEnabled()); + } + @Test void configInstantiationShouldNotFailIfLlmObsIsEnabledAgentlessDisabledViaSysPropAndMlAppIsSet() throws Exception { diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index e43bf780833..766f951d374 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -2361,7 +2361,7 @@ "version": "A", "type": "boolean", "default": "false", - "aliases": [] + "aliases": ["DD_TRACE_LAMBDA_STRIP_INJECTED_CONTEXT"] } ], "DD_LLMOBS_AGENTLESS_ENABLED": [