Strip _datadog trace context injected for Eventbridge, SNS, SQS

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,6 @@ 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')
+  testImplementation libs.jackson.databind
 }

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));
+  }
+}

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java

@@ -12,23 +12,28 @@
 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;
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.google.auto.service.AutoService;
 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;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 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.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumenterModule.class)
@@ -68,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");
   }
@@ -76,16 +85,32 @@ public static class ExtensionCommunicationAdvice {
     @OnMethodEnter
     static ContextScope enter(
         @This final Object that,
-        @Advice.Argument(0) final Object 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 = StripInjectedContext.replaceInputStream(inputStream);
+        if (stripped != null) {
+          inputStream = stripped;
+        }
+      }
+      in = inputStream;
+
       final AgentSpan span;
       if (null == lambdaContext) {
         span = startSpan("java-aws-sdk", INVOCATION_SPAN_NAME);