From 15a04a1b2cd8fe6c66fda8d810b984e233d9a383 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:06:18 +0100 Subject: [PATCH 01/14] Support instrumentations that make structural changes and need to tolerate already-loaded target classes Co-Authored-By: Claude Sonnet 5 --- .../tooling/CombiningTransformerBuilder.java | 58 ++++++++++++++++--- .../trace/agent/tooling/MatchRecorder.java | 58 +++++++++++++++++++ .../trace/agent/tooling/Instrumenter.java | 6 ++ 3 files changed, 113 insertions(+), 9 deletions(-) diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java index 0529488fda7..6c2dd3e8a6d 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java @@ -152,6 +152,16 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati muzzle = new MuzzleCheck(module, instrumentationId); } + /** Allocates a fresh transformation id not known at build-time, growing storage as needed. */ + private int allocateRuntimeTransformationId() { + int transformationId = nextRuntimeTransformationId++; + if (transformers.length <= transformationId) { + int newLen = Math.max(transformationId + 1, transformers.length + (transformers.length >> 1)); + transformers = Arrays.copyOf(transformers, newLen); + } + return transformationId; + } + /** Builds a type-specific transformer, controlled by one or more matchers. */ private void buildTypeInstrumentation(Instrumenter member) { @@ -162,10 +172,7 @@ private void buildTypeInstrumentation(Instrumenter member) { if (transformationId < 0) { // this is a non-indexed transformation configured at runtime, e.g. "dd.trace.methods" // allocate a distinct runtime id to each extra transformation for matching purposes - transformationId = nextRuntimeTransformationId++; - if (transformers.length <= transformationId) { - transformers = Arrays.copyOf(transformers, transformationId + 1); - } + transformationId = allocateRuntimeTransformationId(); } buildTypeMatcher(member, transformationId); @@ -222,30 +229,63 @@ private void buildTypeMatcher(Instrumenter member, int transformationId) { } matchers.add(new MatchRecorder.NarrowLocation(transformationId, muzzle)); + + // preserve structural change, unless we're going to split it out in buildTypeAdvice + if (member instanceof Instrumenter.WithStructuralChange + && !(member instanceof Instrumenter.HasMethodAdvice)) { + addStructuralNarrowing((Instrumenter.WithStructuralChange) member, transformationId); + } } private void buildTypeAdvice(Instrumenter member, int transformationId) { - if (null != helperTransformer) { advice.add(helperTransformer); } - if (null != contextRequestRewriter) { registerContextStoreInjection(member, contextStore); // rewrite context store access to call FieldBackedContextStores with assigned store-id advice.add(contextRequestRewriter); } - if (member instanceof Instrumenter.HasTypeAdvice) { + // a structural change may contain method advice that works with and without the change + // we split out the structural change to its own transformation so it can't accidentally + // turn off the method advice when it's skipped + boolean splitOutStructuralChange = + member instanceof Instrumenter.WithStructuralChange + && member instanceof Instrumenter.HasMethodAdvice; + + if (member instanceof Instrumenter.HasTypeAdvice && !splitOutStructuralChange) { ((Instrumenter.HasTypeAdvice) member).typeAdvice(this); } if (member instanceof Instrumenter.HasMethodAdvice) { ((Instrumenter.HasMethodAdvice) member).methodAdvice(this); } + finishAdviceStack(transformationId); + + if (splitOutStructuralChange) { + Instrumenter.WithStructuralChange structuralChange = + (Instrumenter.WithStructuralChange) member; + + // reuse the type match already computed for transformationId instead of rebuilding it + int structuralTransformationId = allocateRuntimeTransformationId(); + matchers.add(new MatchRecorder.CopyMatch(transformationId, structuralTransformationId)); + addStructuralNarrowing(structuralChange, structuralTransformationId); + structuralChange.typeAdvice(this); + finishAdviceStack(structuralTransformationId); + } + } - // record the advice collected for this transformationId - transformers[transformationId] = new AdviceStack(advice); + /** Narrows a transformation away from already-loaded types missing the structural marker. */ + private void addStructuralNarrowing( + Instrumenter.WithStructuralChange member, int transformationId) { + matchers.add( + new MatchRecorder.PreserveLoadedStructure( + transformationId, member.structuralChangeMarker())); + } + /** Records the advice collected so far as the stack for this transformationId. */ + private void finishAdviceStack(int transformationId) { + transformers[transformationId] = new AdviceStack(advice); advice.clear(); // reset for next transformationId } diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java index 8e6ed90e991..4b3c9186d88 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java @@ -140,6 +140,64 @@ public void record( } } + /** Narrows the current match to avoid structural changes on already-loaded classes. */ + static final class PreserveLoadedStructure extends MatchRecorder { + private final Class structuralChangeMarker; + + PreserveLoadedStructure(int id, Class structuralChangeMarker) { + super(id); + this.structuralChangeMarker = structuralChangeMarker; + } + + @Override + public void record( + TypeDescription type, + ClassLoader classLoader, + Class classBeingRedefined, + BitSet matches) { + // don't transform loaded classes unless they declare the marker (we must retransform those) + if (matches.get(id) && null != classBeingRedefined && !declaresMarker(classBeingRedefined)) { + matches.clear(id); + } + } + + // the marker is added to the exact type we structurally changed before it was loaded; only + // that type needs retransforming, not its sub-types, so this must not use isAssignableFrom + private boolean declaresMarker(Class classBeingRedefined) { + for (Class intf : classBeingRedefined.getInterfaces()) { + if (structuralChangeMarker.equals(intf)) { + return true; + } + } + return false; + } + } + + /** Copies an already-computed match into another id, avoiding a redundant type match. */ + static final class CopyMatch extends MatchRecorder { + private final int fromId; + + CopyMatch(int fromId, int toId) { + super(toId); + if (toId <= fromId) { + // matchers are evaluated in 'id' order; we can only copy from ids strictly before this id + throw new IllegalArgumentException("fromId " + fromId + " must be before toId " + toId); + } + this.fromId = fromId; + } + + @Override + public void record( + TypeDescription type, + ClassLoader classLoader, + Class classBeingRedefined, + BitSet matches) { + if (matches.get(fromId)) { + matches.set(id); + } + } + } + /** Narrows the current match to eliminate incompatible class-loaders. */ static final class NarrowLocation extends MatchRecorder { private final ElementMatcher matcher; diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java index b146121d273..cb3b7e16cdc 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java @@ -81,6 +81,12 @@ interface WithTypeStructure { ElementMatcher structureMatcher(); } + /** Instrumentation whose type advice adds fields, methods, or interfaces to the original type. */ + interface WithStructuralChange extends HasTypeAdvice { + /** The marker interface added by the structural change, used to detect already-loaded types. */ + Class structuralChangeMarker(); + } + /** Instrumentation that provides advice which affects the whole type. */ interface HasTypeAdvice extends Instrumenter { /** From 735fc2e84c406d7e3584ae2e926b02722913c89f Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:09:39 +0100 Subject: [PATCH 02/14] Update TaskUnwrappingInstrumentation to use WithStructuralChange Co-Authored-By: Claude Sonnet 5 --- .../java/concurrent/TaskUnwrappingInstrumentation.java | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java index 01e810a9875..88fd9aeae55 100644 --- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java @@ -8,10 +8,13 @@ import datadog.trace.agent.tooling.bytebuddy.profiling.UnwrappingVisitor; import datadog.trace.api.config.ProfilingConfig; import datadog.trace.bootstrap.config.provider.ConfigProvider; +import datadog.trace.bootstrap.instrumentation.api.TaskWrapper; @AutoService(InstrumenterModule.class) public class TaskUnwrappingInstrumentation extends InstrumenterModule.Profiling - implements Instrumenter.ForKnownTypes, Instrumenter.HasTypeAdvice { + implements Instrumenter.ForKnownTypes, + Instrumenter.HasTypeAdvice, + Instrumenter.WithStructuralChange { public TaskUnwrappingInstrumentation() { super(EXECUTOR_INSTRUMENTATION_NAME, "task-unwrapping"); } @@ -85,4 +88,9 @@ public String[] knownMatchingTypes() { } return types; } + + @Override + public Class structuralChangeMarker() { + return TaskWrapper.class; + } } From f77bda7d7396f21f2d80dc7c6e244fda9ea18bd0 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:14:52 +0100 Subject: [PATCH 03/14] Update IAST Taintable instrumentations to use WithStructuralChange Co-Authored-By: Claude Sonnet 5 --- .../trace/agent/tooling/AgentInstaller.java | 12 ---- .../agent/tooling/InstrumenterModule.java | 13 ++++ ...TaintableRedefinitionStrategyListener.java | 62 ----------------- .../bytebuddy/iast/TaintableVisitor.java | 19 +----- .../bytebuddy/iast/MrReturnAdvice.java | 10 --- .../iast/TaintableVisitorTest.groovy | 68 ------------------- .../bytebuddy/LoadedTaintableClass.java | 8 --- .../iast/AkkaIastTestWebServer.groovy | 5 +- .../akkahttp/iast/IastAkkaTest.groovy | 14 ++-- .../iast/MakeTaintableInstrumentation.java | 2 +- ...IastHttpUriRequestBaseInstrumentation.java | 2 +- .../IastHttpMethodBaseInstrumentation.java | 2 +- .../ByteBufInputStreamInstrumentation.java | 2 +- .../buffer/ByteBufInstrumentation.java | 2 +- .../okhttp2/IastHttpUrlInstrumentation.java | 2 +- .../okhttp3/IastHttpUrlInstrumentation.java | 2 +- .../iast/MakeTaintableInstrumentation.java | 2 +- .../servlet/http/CookieInstrumentation.java | 2 +- .../vertx_3_4/core/BufferInstrumentation.java | 2 +- ...CaseInsensitiveHeadersInstrumentation.java | 2 +- .../core/HeadersAdaptorInstrumentation.java | 2 +- .../server/CookieImplInstrumentation.java | 2 +- .../core/VertxHttpHeadersInstrumentation.java | 2 +- .../vertx_4_0/core/BufferInstrumentation.java | 2 +- .../core/MultiMapInstrumentation.java | 2 +- .../server/CookieImplInstrumentation.java | 2 +- 26 files changed, 47 insertions(+), 198 deletions(-) delete mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableRedefinitionStrategyListener.java delete mode 100644 dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/MrReturnAdvice.java delete mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/LoadedTaintableClass.java diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 321b726d63b..d331759d2df 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java @@ -8,7 +8,6 @@ import datadog.environment.SystemProperties; import datadog.instrument.fieldinject.GlobalObjectStore; import datadog.trace.agent.tooling.bytebuddy.SharedTypePools; -import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener; import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers; import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers; import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade; @@ -166,7 +165,6 @@ public static ClassFileTransformer installBytebuddyAgent( .with(AgentStrategies.transformerDecorator()) .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) .with(AgentStrategies.rediscoveryStrategy()) - .with(redefinitionStrategyListener(enabledSystems)) .with(AgentStrategies.locationStrategy()) .with(AgentStrategies.poolStrategy()) .with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY) @@ -183,7 +181,6 @@ public static ClassFileTransformer installBytebuddyAgent( agentBuilder .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) .with(AgentStrategies.rediscoveryStrategy()) - .with(redefinitionStrategyListener(enabledSystems)) .with(new RedefinitionLoggingListener()) .with(new TransformLoggingListener()); } @@ -398,15 +395,6 @@ private static void temporaryOverride(String key, String value, BooleanSupplier } } - private static AgentBuilder.RedefinitionStrategy.Listener redefinitionStrategyListener( - final Set enabledSystems) { - if (enabledSystems.contains(InstrumenterModule.TargetSystem.IAST)) { - return TaintableRedefinitionStrategyListener.INSTANCE; - } else { - return AgentBuilder.RedefinitionStrategy.Listener.NoOp.INSTANCE; - } - } - static class RedefinitionLoggingListener implements AgentBuilder.RedefinitionStrategy.Listener { private static final Logger log = LoggerFactory.getLogger(RedefinitionLoggingListener.class); diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java index d2abbc265e5..659caffe7bb 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java @@ -13,6 +13,7 @@ import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.ProductActivation; import datadog.trace.api.config.ProfilingConfig; +import datadog.trace.api.iast.Taintable; import datadog.trace.bootstrap.config.provider.ConfigProvider; import datadog.trace.util.Strings; import de.thetaphi.forbiddenapis.SuppressForbidden; @@ -308,6 +309,18 @@ protected boolean isOptOutEnabled() { } } + /** Parent class for IAST instrumentations that restructure classes to add {@link Taintable}. */ + public abstract static class TaintableIast extends Iast implements WithStructuralChange { + public TaintableIast(String instrumentationName, String... additionalNames) { + super(instrumentationName, additionalNames); + } + + @Override + public final Class structuralChangeMarker() { + return Taintable.class; + } + } + /** Parent class for all USM related instrumentations */ public abstract static class Usm extends InstrumenterModule { public Usm(String instrumentationName, String... additionalNames) { diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableRedefinitionStrategyListener.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableRedefinitionStrategyListener.java deleted file mode 100644 index d8b8284146d..00000000000 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableRedefinitionStrategyListener.java +++ /dev/null @@ -1,62 +0,0 @@ -package datadog.trace.agent.tooling.bytebuddy.iast; - -import java.util.Collections; -import java.util.List; -import java.util.Map; -import javax.annotation.Nonnull; -import net.bytebuddy.agent.builder.AgentBuilder; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * {@link TaintableVisitor} redefines the structure of a class by adding interfaces and fields, - * meaning that it cannot be applied to already loaded classes. - * - *

This listener will disable the visitor to prevent a failure with the whole redefinition batch. - */ -public final class TaintableRedefinitionStrategyListener - extends AgentBuilder.RedefinitionStrategy.Listener.Adapter { - - private static final Logger LOGGER = - LoggerFactory.getLogger(TaintableRedefinitionStrategyListener.class); - private static final boolean DEBUG = LOGGER.isDebugEnabled(); - - public static final TaintableRedefinitionStrategyListener INSTANCE = - new TaintableRedefinitionStrategyListener(); - - private TaintableRedefinitionStrategyListener() {} - - @Override - @Nonnull - public Iterable>> onError( - final int index, - @Nonnull final List> batch, - @Nonnull final Throwable throwable, - @Nonnull final List> types) { - if (TaintableVisitor.ENABLED) { - if (DEBUG) { - LOGGER.debug( - "Exception while retransforming with the visitor in batch {}, disabling it", index); - } - TaintableVisitor.ENABLED = false; - return Collections.singletonList(batch); - } else { - if (DEBUG) { - LOGGER.debug( - "Exception while retransforming after disabling the visitor in batch {}, classes won't be instrumented", - index); - } - return Collections.emptyList(); - } - } - - @Override - public void onComplete( - final int amount, final List> types, final Map>, Throwable> failures) { - if (DEBUG) { - if (!TaintableVisitor.ENABLED) { - LOGGER.debug("Retransforming succeeded with a disabled visitor"); - } - } - } -} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java index 1d644ba7580..595b60a6e9b 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java @@ -22,7 +22,6 @@ public class TaintableVisitor implements AsmVisitorWrapper { public static volatile boolean DEBUG = false; - static volatile boolean ENABLED = true; private static final String INTERFACE_NAME = "datadog/trace/api/iast/Taintable"; private static final String SOURCE_CLASS_NAME = "L" + INTERFACE_NAME + "$Source;"; @@ -56,21 +55,9 @@ public ClassVisitor wrap( final MethodList methods, final int writerFlags, final int readerFlags) { - if (ENABLED) { - return types.contains(instrumentedType.getName()) - ? new AddTaintableInterfaceVisitor(classVisitor) - : classVisitor; - } else { - return NoOp.INSTANCE.wrap( - instrumentedType, - classVisitor, - implementationContext, - typePool, - fields, - methods, - writerFlags, - readerFlags); - } + return types.contains(instrumentedType.getName()) + ? new AddTaintableInterfaceVisitor(classVisitor) + : classVisitor; } private static class AddTaintableInterfaceVisitor extends ClassVisitor { diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/MrReturnAdvice.java b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/MrReturnAdvice.java deleted file mode 100644 index 88872a50f41..00000000000 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/MrReturnAdvice.java +++ /dev/null @@ -1,10 +0,0 @@ -package datadog.trace.agent.tooling.bytebuddy.iast; - -import net.bytebuddy.asm.Advice; - -public class MrReturnAdvice { - @Advice.OnMethodExit - public static void sayHello(@Advice.Return(readOnly = false) String result) { - result = "Hello Mr!"; - } -} diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitorTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitorTest.groovy index 88a51e3ab35..23e9f941abf 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitorTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitorTest.groovy @@ -1,33 +1,12 @@ package datadog.trace.agent.tooling.bytebuddy.iast -import datadog.trace.agent.tooling.bytebuddy.LoadedTaintableClass import datadog.trace.api.iast.Taintable import datadog.trace.test.util.DDSpecification import net.bytebuddy.ByteBuddy -import net.bytebuddy.agent.ByteBuddyAgent -import net.bytebuddy.agent.builder.AgentBuilder import net.bytebuddy.description.modifier.Visibility -import net.bytebuddy.description.type.TypeDescription -import net.bytebuddy.dynamic.DynamicType -import net.bytebuddy.utility.JavaModule -import net.bytebuddy.utility.nullability.MaybeNull - -import java.security.ProtectionDomain - -import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named class TaintableVisitorTest extends DDSpecification { - private boolean wasEnabled - - void setup() { - wasEnabled = TaintableVisitor.ENABLED - } - - void cleanup() { - TaintableVisitor.ENABLED = wasEnabled - } - void 'test taintable visitor'() { given: final className = 'datadog.trace.agent.tooling.bytebuddy.iast.TaintableTest' @@ -100,51 +79,4 @@ class TaintableVisitorTest extends DDSpecification { then: 1 * source.getOrigin() } - - void 'test taintable visitor with already loaded class'() { - given: - final instance = new LoadedTaintableClass() - final listener = Mock(AgentBuilder.RedefinitionStrategy.Listener) - - when: - final result = instance.sayHello() - - then: - result == 'Hello!' - - when: - new AgentBuilder.Default() - .disableClassFormatChanges() - .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) - .with(TaintableRedefinitionStrategyListener.INSTANCE) - .with(listener) - .type(named(LoadedTaintableClass.name)) - .transform(new AgentBuilder.Transformer.ForAdvice().advice(named('sayHello'), 'datadog.trace.agent.tooling.bytebuddy.iast.MrReturnAdvice')) - .transform(new AgentBuilder.Transformer() { - @Override - DynamicType.Builder transform(DynamicType.Builder builder, - TypeDescription typeDescription, - @MaybeNull ClassLoader classLoader, - @MaybeNull JavaModule module, - ProtectionDomain protectionDomain) { - return builder.visit(new TaintableVisitor(LoadedTaintableClass.name)) - } - }) - .installOn(ByteBuddyAgent.instrumentation) - - then: - final modifiedResult = instance.sayHello() - - then: - modifiedResult == 'Hello Mr!' - // failing initial batch - 1 * listener.onBatch(0, { List> list -> list.contains(LoadedTaintableClass) }, _) - 1 * listener.onError(0, { List> list -> list.contains(LoadedTaintableClass) }, _, _) >> [] - - // successful batch after disabling the visitor - 1 * listener.onBatch(1, { List> list -> list.contains(LoadedTaintableClass) }, _) - - // finally two batches where executed - 1 * listener.onComplete(2, { List> list -> list.contains(LoadedTaintableClass) }, _) - } } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/LoadedTaintableClass.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/LoadedTaintableClass.java deleted file mode 100644 index a851f988502..00000000000 --- a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/LoadedTaintableClass.java +++ /dev/null @@ -1,8 +0,0 @@ -package datadog.trace.agent.tooling.bytebuddy; - -public class LoadedTaintableClass { - - public String sayHello() { - return "Hello!"; - } -} diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/AkkaIastTestWebServer.groovy b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/AkkaIastTestWebServer.groovy index 161d443937f..5389581c452 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/AkkaIastTestWebServer.groovy +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/AkkaIastTestWebServer.groovy @@ -54,7 +54,10 @@ class AkkaIastTestWebServer extends AllDirectives implements Closeable { path(segment('path').slash(segment())) { var1 -> get { - complete("IAST: ${t(var1)}") + extractRequest { + request -> + complete("IAST: ${t(var1)} ${t(request)}") + } } }, path('query') { diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/IastAkkaTest.groovy b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/IastAkkaTest.groovy index 5d5486264ec..36710e2ad65 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/IastAkkaTest.groovy +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/iastTest/groovy/datadog/trace/instrumentation/akkahttp/iast/IastAkkaTest.groovy @@ -40,7 +40,8 @@ class IastAkkaTest extends IastRequestTestRunner { then: response.code() == 200 - response.body().string() == 'IAST: myValue (tainted)' + // verify that HttpRequest is now also tainted via the Taintable interface + response.body().string() =~ /^IAST: myValue \(tainted\) HttpRequest\(.*\) \(tainted\)$/ when: def toc = finReqTaintedObjects @@ -50,9 +51,14 @@ class IastAkkaTest extends IastRequestTestRunner { value 'myValue' range 0, 7, source(SourceTypes.REQUEST_PATH_PARAMETER, null, 'myValue') } - // After migrating from JUnit 4 to 5, the IAST instrumentation scope has expanded so that the following are tainted: - // - Accept-Encoding, Connection, Host, HttpRequest, User-Agent, RequestContext, Timeout-Access, Remote-Address, myValue - toc.size() == 9 + // The following are always tainted: + // - Accept-Encoding, Connection, Host, User-Agent, Timeout-Access, Remote-Address, myValue (7 entries) + // HttpRequest and its HttpEntity are also always tainted, but the mechanism depends on whether their + // classes were already loaded when the akka-http IAST instrumentation was applied: if already loaded, + // PreserveLoadedStructure can't inject the Taintable interface into them (the JVM can't add an interface + // via retransformation), so they fall back to being tracked in the tainted objects map instead - adding + // up to 2 extra entries. Whether that happens depends on the akka-http version and class-loading order. + toc.size() >= 7 && toc.size() <= 9 } void 'cookie — #variant variant'() { diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/MakeTaintableInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/MakeTaintableInstrumentation.java index a5608c48110..c2898fc5afc 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/MakeTaintableInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/MakeTaintableInstrumentation.java @@ -6,7 +6,7 @@ import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; @AutoService(InstrumenterModule.class) -public class MakeTaintableInstrumentation extends InstrumenterModule.Iast +public class MakeTaintableInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForKnownTypes, Instrumenter.HasTypeAdvice { public MakeTaintableInstrumentation() { super("akka-http"); diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/IastHttpUriRequestBaseInstrumentation.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/IastHttpUriRequestBaseInstrumentation.java index 12fece51ef2..41483b459d3 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/IastHttpUriRequestBaseInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/IastHttpUriRequestBaseInstrumentation.java @@ -15,7 +15,7 @@ import org.apache.hc.client5.http.classic.methods.HttpUriRequestBase; @AutoService(InstrumenterModule.class) -public class IastHttpUriRequestBaseInstrumentation extends InstrumenterModule.Iast +public class IastHttpUriRequestBaseInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java b/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java index fc6f09dfb11..9594b7b78e8 100644 --- a/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java +++ b/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java @@ -14,7 +14,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class IastHttpMethodBaseInstrumentation extends InstrumenterModule.Iast +public class IastHttpMethodBaseInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java index 596dd5f5523..b240ae1a1d8 100644 --- a/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java @@ -16,7 +16,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class ByteBufInputStreamInstrumentation extends InstrumenterModule.Iast +public class ByteBufInputStreamInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInstrumentation.java index 266e9c633db..0b7da02e4c6 100644 --- a/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-buffer-4.0/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInstrumentation.java @@ -6,7 +6,7 @@ import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; @AutoService(InstrumenterModule.class) -public class ByteBufInstrumentation extends InstrumenterModule.Iast +public class ByteBufInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java index 96fbcb373f8..7b8b0ee2879 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java @@ -17,7 +17,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class IastHttpUrlInstrumentation extends InstrumenterModule.Iast +public class IastHttpUrlInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java index a718db91704..f56eaf349c3 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java @@ -17,7 +17,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class IastHttpUrlInstrumentation extends InstrumenterModule.Iast +public class IastHttpUrlInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/MakeTaintableInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/MakeTaintableInstrumentation.java index fca33f5c437..f65d6e19371 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/MakeTaintableInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/MakeTaintableInstrumentation.java @@ -6,7 +6,7 @@ import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; @AutoService(InstrumenterModule.class) -public class MakeTaintableInstrumentation extends InstrumenterModule.Iast +public class MakeTaintableInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForKnownTypes, Instrumenter.HasTypeAdvice { public MakeTaintableInstrumentation() { super("pekko-http"); diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java index 0fd197c7af4..ea2334882bc 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java @@ -20,7 +20,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class CookieInstrumentation extends InstrumenterModule.Iast +public class CookieInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java index ae7ef68c05d..e0718103cdc 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java @@ -20,7 +20,7 @@ /** Propagation is way easier in io.vertx.core.buffer.impl.BufferImpl than in io.netty.Buffer */ @AutoService(InstrumenterModule.class) -public class BufferInstrumentation extends InstrumenterModule.Iast +public class BufferInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/CaseInsensitiveHeadersInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/CaseInsensitiveHeadersInstrumentation.java index c663a957d3e..f84f6d00908 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/CaseInsensitiveHeadersInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/CaseInsensitiveHeadersInstrumentation.java @@ -31,7 +31,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class CaseInsensitiveHeadersInstrumentation extends InstrumenterModule.Iast +public class CaseInsensitiveHeadersInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HeadersAdaptorInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HeadersAdaptorInstrumentation.java index 90aae45ae86..8109f448ed7 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HeadersAdaptorInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HeadersAdaptorInstrumentation.java @@ -29,7 +29,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class HeadersAdaptorInstrumentation extends InstrumenterModule.Iast +public class HeadersAdaptorInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForKnownTypes, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/CookieImplInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/CookieImplInstrumentation.java index 04f46384a4a..d34c4950ccb 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/CookieImplInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/CookieImplInstrumentation.java @@ -24,7 +24,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class CookieImplInstrumentation extends InstrumenterModule.Iast +public class CookieImplInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.5/src/main/java/datadog/trace/instrumentation/vertx_3_5/core/VertxHttpHeadersInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.5/src/main/java/datadog/trace/instrumentation/vertx_3_5/core/VertxHttpHeadersInstrumentation.java index c6dfcc37c37..538380a2c2f 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.5/src/main/java/datadog/trace/instrumentation/vertx_3_5/core/VertxHttpHeadersInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.5/src/main/java/datadog/trace/instrumentation/vertx_3_5/core/VertxHttpHeadersInstrumentation.java @@ -28,7 +28,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class VertxHttpHeadersInstrumentation extends InstrumenterModule.Iast +public class VertxHttpHeadersInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java index 1f0bbe17613..851335fc7f7 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java @@ -19,7 +19,7 @@ /** Propagation is way easier in io.vertx.core.buffer.impl.BufferImpl than in io.netty.Buffer */ @AutoService(InstrumenterModule.class) -public class BufferInstrumentation extends InstrumenterModule.Iast +public class BufferInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/MultiMapInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/MultiMapInstrumentation.java index 75cf8e1e83c..4ba4c21b522 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/MultiMapInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/MultiMapInstrumentation.java @@ -29,7 +29,7 @@ import net.bytebuddy.description.method.MethodDescription; import net.bytebuddy.matcher.ElementMatcher; -public abstract class MultiMapInstrumentation extends InstrumenterModule.Iast +public abstract class MultiMapInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { private final String className = MultiMapInstrumentation.class.getName(); diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/CookieImplInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/CookieImplInstrumentation.java index 9fd9d4540d5..ce09f351d12 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/CookieImplInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/CookieImplInstrumentation.java @@ -23,7 +23,7 @@ import net.bytebuddy.asm.Advice; @AutoService(InstrumenterModule.class) -public class CookieImplInstrumentation extends InstrumenterModule.Iast +public class CookieImplInstrumentation extends InstrumenterModule.TaintableIast implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { From 8f75a216ff3cfbe04d05193e1c1dbba86e69e020 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:15:34 +0100 Subject: [PATCH 04/14] Fix bug in TaintedMap where it wouldn't deduplicate the last element This in turn revealed a gap in KafkaIastDeserializerTest when checking pass-through byte-buffers --- .../com/datadog/iast/taint/TaintedMap.java | 7 ++-- .../iast/KafkaIastDeserializerTest.groovy | 34 +++++++++++++------ 2 files changed, 29 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedMap.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedMap.java index 6e2de584b2a..63f77c7ca28 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedMap.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedMap.java @@ -194,12 +194,15 @@ public void put(final @Nonnull TaintedObject entry) { entry.generation = generation; } else { int bucketSize = 1; - TaintedObject next; - while ((next = next(cur)) != null) { + while (true) { if (cur.positiveHashCode == entry.positiveHashCode && cur.get() == entry.get()) { // Duplicate, exit early. return; } + final TaintedObject next = next(cur); + if (next == null) { + break; + } bucketSize++; cur = next; } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy index 3b2c0c41f22..40ced18009e 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy @@ -18,6 +18,9 @@ import static org.hamcrest.core.IsEqual.equalTo class KafkaIastDeserializerTest extends IastRequestTestRunner { + private static final String PAYLOAD_STRING = "Hello World!" + private static final int PAYLOAD_LENGTH = PAYLOAD_STRING.bytes.length + private static final int BUFF_OFFSET = 10 void 'test string deserializer: #test'() { @@ -27,7 +30,7 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { InstrumentationBridge.registerIastModule(propagationModule) and: - final payload = "Hello World!".bytes + final payload = PAYLOAD_STRING.bytes final deserializer = new StringDeserializer() when: @@ -39,8 +42,8 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { then: final to = finReqTaintedObjects to.hasTaintedObject { - value('Hello World!') - range(0, 12, source(origin)) + value(PAYLOAD_STRING) + range(0, PAYLOAD_LENGTH, source(origin)) } where: @@ -54,7 +57,7 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { InstrumentationBridge.registerIastModule(propagationModule) and: - final payload = "Hello World!".bytes + final payload = PAYLOAD_STRING.bytes final deserializer = new ByteArrayDeserializer() when: @@ -81,7 +84,7 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { InstrumentationBridge.registerIastModule(propagationModule) and: - final payload = "Hello World!".bytes + final payload = PAYLOAD_STRING.bytes final deserializer = new ByteBufferDeserializer() when: @@ -94,7 +97,7 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { final to = finReqTaintedObjects to.hasTaintedObject { value(instanceOf(ByteBuffer)) - range(0, Integer.MAX_VALUE, source(origin)) + range(test.method.start, test.method.length, source(origin)) } where: @@ -145,19 +148,19 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { } enum Method { - DEFAULT{ + DEFAULT(0, Integer.MAX_VALUE){ @Override T deserialize(Deserializer deserializer, String topic, byte[] payload) { return deserializer.deserialize(topic, payload) } }, - WITH_HEADERS{ + WITH_HEADERS(0, Integer.MAX_VALUE){ @Override T deserialize(Deserializer deserializer, String topic, byte[] payload) { return deserializer.deserialize(topic, new RecordHeaders(), payload) } }, - WITH_BYTE_BUFFER{ + WITH_BYTE_BUFFER(0, PAYLOAD_LENGTH){ @SuppressWarnings('GroovyAssignabilityCheck') @Override T deserialize(Deserializer deserializer, String topic, byte[] payload) { @@ -167,7 +170,7 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { return deserializer.deserialize(topic, new RecordHeaders(), buffer) } }, - WITH_BYTE_BUFFER_OFFSET{ + WITH_BYTE_BUFFER_OFFSET(BUFF_OFFSET, PAYLOAD_LENGTH){ @SuppressWarnings('GroovyAssignabilityCheck') @Override T deserialize(Deserializer deserializer, String topic, byte[] payload) { @@ -177,6 +180,17 @@ class KafkaIastDeserializerTest extends IastRequestTestRunner { } } + // start/length of the taint range expected on the ByteBuffer result of deserialization; + // a pass-through deserializer keeps the precise range tainted beforehand, otherwise a + // brand-new ByteBuffer is wrapped around the result and gets a fresh, full-object range + final int start + final int length + + Method(int start, int length) { + this.start = start + this.length = length + } + abstract T deserialize(Deserializer deserializer, String topic, byte[] payload) } From 2251a6cc4ec814653b9e8f5aa2cebe033b3211ce Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:17:36 +0100 Subject: [PATCH 05/14] Update Vert.x RequestImplInstrumentation to use WithStructuralChange Co-Authored-By: Claude Sonnet 5 --- .../RequestImplInstrumentation.java | 75 ++++++++++++------- 1 file changed, 48 insertions(+), 27 deletions(-) diff --git a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java index 8ef3502d257..da1c8285963 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java @@ -22,7 +22,8 @@ public class RequestImplInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice, - Instrumenter.HasMethodAdvice { + Instrumenter.HasMethodAdvice, + Instrumenter.WithStructuralChange { public RequestImplInstrumentation() { super("vertx", "vertx-redis-client"); } @@ -43,6 +44,11 @@ public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice(none(), packageName + ".RequestImplMuzzle"); } + @Override + public Class structuralChangeMarker() { + return Cloneable.class; + } + // This Transformer will add the Cloneable interface to RequestImpl, as well // as a clone method that calls the protected shallow clone method in Object public static class RequestImplVisitorWrapper implements AsmVisitorWrapper { @@ -68,6 +74,8 @@ public ClassVisitor wrap( int readerFlags) { return new ClassVisitor(Opcodes.ASM7, classVisitor) { + private boolean addCloneable = true; + @Override public void visit( int version, @@ -76,38 +84,51 @@ public void visit( String signature, String superName, String[] interfaces) { - // Add the Cloneable interface - if (null == interfaces) { - interfaces = new String[1]; - } else { - interfaces = Arrays.copyOf(interfaces, interfaces.length + 1); + // Add the Cloneable interface, unless it's already there (e.g. reapplying this + // change while retransforming a class that was already modified on initial load) + if (interfaces != null) { + for (String iface : interfaces) { + if ("java/lang/Cloneable".equals(iface)) { + addCloneable = false; + break; + } + } + } + if (addCloneable) { + if (null == interfaces) { + interfaces = new String[1]; + } else { + interfaces = Arrays.copyOf(interfaces, interfaces.length + 1); + } + interfaces[interfaces.length - 1] = "java/lang/Cloneable"; } - interfaces[interfaces.length - 1] = "java/lang/Cloneable"; cv.visit(version, access, name, signature, superName, interfaces); } @Override public void visitEnd() { - // Add a clone method that calls the protected shallow clone method in Object - // - // public Object clone() throws CloneNotSupportedException { - // return super.clone(); // Object is the super class - // } - // - final MethodVisitor mv = - cv.visitMethod( - Opcodes.ACC_PUBLIC, - "clone", - "()Ljava/lang/Object;", - null, - new String[] {"java/lang/CloneNotSupportedException"}); - mv.visitCode(); - mv.visitIntInsn(Opcodes.ALOAD, 0); - mv.visitMethodInsn( - Opcodes.INVOKESPECIAL, "java/lang/Object", "clone", "()Ljava/lang/Object;", false); - mv.visitInsn(Opcodes.ARETURN); - mv.visitMaxs(0, 0); - mv.visitEnd(); + if (addCloneable) { + // Add a clone method that calls the protected shallow clone method in Object + // + // public Object clone() throws CloneNotSupportedException { + // return super.clone(); // Object is the super class + // } + // + final MethodVisitor mv = + cv.visitMethod( + Opcodes.ACC_PUBLIC, + "clone", + "()Ljava/lang/Object;", + null, + new String[] {"java/lang/CloneNotSupportedException"}); + mv.visitCode(); + mv.visitIntInsn(Opcodes.ALOAD, 0); + mv.visitMethodInsn( + Opcodes.INVOKESPECIAL, "java/lang/Object", "clone", "()Ljava/lang/Object;", false); + mv.visitInsn(Opcodes.ARETURN); + mv.visitMaxs(0, 0); + mv.visitEnd(); + } cv.visitEnd(); } From a0dddd67cd2d6bc8f15c9f0f6b6079ade814619c Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:45:32 +0100 Subject: [PATCH 06/14] Use consistent pattern when adding a new interface via advice --- .../bytebuddy/iast/TaintableVisitor.java | 42 ++++++------------- .../profiling/UnwrappingVisitor.java | 12 +----- .../RequestImplInstrumentation.java | 24 ++++------- .../datadog/trace/util/CollectionUtils.java | 17 ++++++++ .../trace/util/CollectionUtilsTest.groovy | 22 ++++++++++ 5 files changed, 62 insertions(+), 55 deletions(-) diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java index 595b60a6e9b..72e34b610e5 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java @@ -1,10 +1,10 @@ package datadog.trace.agent.tooling.bytebuddy.iast; import datadog.trace.api.iast.Taintable; +import datadog.trace.util.CollectionUtils; import java.util.Arrays; import java.util.HashSet; import java.util.Set; -import javax.annotation.Nullable; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.field.FieldDescription; import net.bytebuddy.description.field.FieldList; @@ -23,8 +23,8 @@ public class TaintableVisitor implements AsmVisitorWrapper { public static volatile boolean DEBUG = false; - private static final String INTERFACE_NAME = "datadog/trace/api/iast/Taintable"; - private static final String SOURCE_CLASS_NAME = "L" + INTERFACE_NAME + "$Source;"; + private static final String TAINTABLE = "datadog/trace/api/iast/Taintable"; + private static final String SOURCE_CLASS_NAME = "L" + TAINTABLE + "$Source;"; private static final String FIELD_NAME = "$$DD$source"; private static final String GETTER_NAME = "$$DD$getSource"; private static final String SETTER_NAME = "$$DD$setSource"; @@ -64,7 +64,7 @@ private static class AddTaintableInterfaceVisitor extends ClassVisitor { private String owner; - private boolean addTaintable = true; + private boolean addTaintable = false; protected AddTaintableInterfaceVisitor(final ClassVisitor classVisitor) { super(OpenedClassReader.ASM_API, classVisitor); @@ -75,25 +75,18 @@ public void visit( final int version, final int access, final String name, - final String signature, + String signature, final String superName, - final String[] interfaces) { + String[] interfaces) { owner = name; - if (interfaces != null) { - for (final String iface : interfaces) { - if (INTERFACE_NAME.equals(iface)) { - addTaintable = false; - break; - } + if (interfaces == null || !Arrays.asList(interfaces).contains(TAINTABLE)) { + interfaces = CollectionUtils.append(interfaces, TAINTABLE); + if (signature != null) { + signature += 'L' + TAINTABLE + ';'; } + addTaintable = true; } - super.visit( - version, - access, - name, - signature, - superName, - addTaintable ? addInterface(interfaces) : interfaces); + super.visit(version, access, name, signature, superName, interfaces); } @Override @@ -109,17 +102,6 @@ public void visitEnd() { } } - private String[] addInterface(@Nullable final String[] interfaces) { - if (interfaces == null || interfaces.length == 0) { - return new String[] {INTERFACE_NAME}; - } else { - final String[] newInterfaces = new String[interfaces.length + 1]; - System.arraycopy(interfaces, 0, newInterfaces, 0, interfaces.length); - newInterfaces[newInterfaces.length - 1] = INTERFACE_NAME; - return newInterfaces; - } - } - private void addField() { final FieldVisitor fv = cv.visitField( diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java index 1c8e07ecaf4..4f24df67bdf 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java @@ -1,5 +1,6 @@ package datadog.trace.agent.tooling.bytebuddy.profiling; +import datadog.trace.util.CollectionUtils; import java.util.Arrays; import java.util.HashMap; import java.util.Map; @@ -82,7 +83,7 @@ public void visit( String superName, String[] interfaces) { if (interfaces == null || !Arrays.asList(interfaces).contains(TASK_WRAPPER)) { - interfaces = append(interfaces, TASK_WRAPPER); + interfaces = CollectionUtils.append(interfaces, TASK_WRAPPER); if (signature != null) { signature += 'L' + TASK_WRAPPER + ';'; } @@ -91,15 +92,6 @@ public void visit( super.visit(version, access, name, signature, superName, interfaces); } - private static String[] append(String[] strings, String toAppend) { - if (strings == null || strings.length == 0) { - return new String[] {toAppend}; - } - String[] appended = Arrays.copyOf(strings, strings.length + 1); - appended[strings.length] = toAppend; - return appended; - } - @Override public FieldVisitor visitField( int access, String name, String descriptor, String signature, Object value) { diff --git a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java index da1c8285963..83e2397bf76 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java @@ -5,6 +5,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.util.CollectionUtils; import java.util.Arrays; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.field.FieldDescription; @@ -74,7 +75,9 @@ public ClassVisitor wrap( int readerFlags) { return new ClassVisitor(Opcodes.ASM7, classVisitor) { - private boolean addCloneable = true; + private static final String CLONEABLE = "java/lang/Cloneable"; + + private boolean addCloneable = false; @Override public void visit( @@ -86,21 +89,12 @@ public void visit( String[] interfaces) { // Add the Cloneable interface, unless it's already there (e.g. reapplying this // change while retransforming a class that was already modified on initial load) - if (interfaces != null) { - for (String iface : interfaces) { - if ("java/lang/Cloneable".equals(iface)) { - addCloneable = false; - break; - } - } - } - if (addCloneable) { - if (null == interfaces) { - interfaces = new String[1]; - } else { - interfaces = Arrays.copyOf(interfaces, interfaces.length + 1); + if (interfaces == null || !Arrays.asList(interfaces).contains(CLONEABLE)) { + interfaces = CollectionUtils.append(interfaces, CLONEABLE); + if (signature != null) { + signature += 'L' + CLONEABLE + ';'; } - interfaces[interfaces.length - 1] = "java/lang/Cloneable"; + addCloneable = true; } cv.visit(version, access, name, signature, superName, interfaces); } diff --git a/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java b/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java index 00878c3f77e..ad032186350 100644 --- a/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java +++ b/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java @@ -5,12 +5,15 @@ import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; +import java.lang.reflect.Array; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; +import javax.annotation.Nullable; public final class CollectionUtils { @@ -87,4 +90,18 @@ private static MethodHandle findCopyOf(Class clazz, Class arg) { } return null; } + + /** Appends an element to an array, growing it as needed. */ + public static T[] append(@Nullable T[] array, T toAppend) { + T[] appended; + if (array == null) { + //noinspection unchecked + appended = (T[]) Array.newInstance(toAppend.getClass(), 1); + appended[0] = toAppend; + } else { + appended = Arrays.copyOf(array, array.length + 1); + appended[array.length] = toAppend; + } + return appended; + } } diff --git a/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy b/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy index f279f8438bd..06a97350b19 100644 --- a/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy @@ -4,6 +4,7 @@ package datadog.trace.util import datadog.trace.test.util.DDSpecification import static datadog.environment.JavaVirtualMachine.isJavaVersionAtLeast +import static datadog.trace.util.CollectionUtils.append import static datadog.trace.util.CollectionUtils.tryMakeImmutableList import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap import static datadog.trace.util.CollectionUtils.tryMakeImmutableSet @@ -39,4 +40,25 @@ class CollectionUtilsTest extends DDSpecification { then: hopefullyImmutable.getClass().getName().contains(expectedClassName) } + + def "append grows a null array"() { + when: + String[] result = append(null, "a") + then: + result == ["a"] as String[] + } + + def "append grows an empty array"() { + when: + String[] result = append(new String[0], "a") + then: + result == ["a"] as String[] + } + + def "append grows a non-empty array"() { + when: + String[] result = append(["a", "b"] as String[], "c") + then: + result == ["a", "b", "c"] as String[] + } } From 174ed6be63a55cb1638d403efcc6f0bc1b24de28 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 11:58:13 +0100 Subject: [PATCH 07/14] Simplify PreserveLoadedStructure --- .../trace/agent/tooling/MatchRecorder.java | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java index 4b3c9186d88..627c3c1140b 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java @@ -4,6 +4,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed; import datadog.trace.agent.tooling.context.FieldBackedContextMatcher; +import java.util.Arrays; import java.util.BitSet; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -155,22 +156,14 @@ public void record( ClassLoader classLoader, Class classBeingRedefined, BitSet matches) { - // don't transform loaded classes unless they declare the marker (we must retransform those) - if (matches.get(id) && null != classBeingRedefined && !declaresMarker(classBeingRedefined)) { + // don't transform loaded classes unless they directly declare the marker + // - we must re-transform them to preserve the original structural change + if (matches.get(id) + && null != classBeingRedefined + && !Arrays.asList(classBeingRedefined.getInterfaces()).contains(structuralChangeMarker)) { matches.clear(id); } } - - // the marker is added to the exact type we structurally changed before it was loaded; only - // that type needs retransforming, not its sub-types, so this must not use isAssignableFrom - private boolean declaresMarker(Class classBeingRedefined) { - for (Class intf : classBeingRedefined.getInterfaces()) { - if (structuralChangeMarker.equals(intf)) { - return true; - } - } - return false; - } } /** Copies an already-computed match into another id, avoiding a redundant type match. */ From 57e838fb71762f4953b7e28da41b4e4fbad47d39 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 12:07:33 +0100 Subject: [PATCH 08/14] Extract array.contains utility --- .../trace/agent/tooling/MatchRecorder.java | 5 +++-- .../bytebuddy/iast/TaintableVisitor.java | 2 +- .../profiling/UnwrappingVisitor.java | 3 +-- .../RequestImplInstrumentation.java | 3 +-- .../datadog/trace/util/CollectionUtils.java | 22 ++++++++++++++----- .../trace/util/CollectionUtilsTest.groovy | 12 ++++++++++ 6 files changed, 35 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java index 627c3c1140b..afe11415552 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java @@ -4,7 +4,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed; import datadog.trace.agent.tooling.context.FieldBackedContextMatcher; -import java.util.Arrays; +import datadog.trace.util.CollectionUtils; import java.util.BitSet; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -160,7 +160,8 @@ public void record( // - we must re-transform them to preserve the original structural change if (matches.get(id) && null != classBeingRedefined - && !Arrays.asList(classBeingRedefined.getInterfaces()).contains(structuralChangeMarker)) { + && !CollectionUtils.contains( + classBeingRedefined.getInterfaces(), structuralChangeMarker)) { matches.clear(id); } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java index 72e34b610e5..179054a2a84 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java @@ -79,7 +79,7 @@ public void visit( final String superName, String[] interfaces) { owner = name; - if (interfaces == null || !Arrays.asList(interfaces).contains(TAINTABLE)) { + if (!CollectionUtils.contains(interfaces, TAINTABLE)) { interfaces = CollectionUtils.append(interfaces, TAINTABLE); if (signature != null) { signature += 'L' + TAINTABLE + ';'; diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java index 4f24df67bdf..8af589a443f 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java @@ -1,7 +1,6 @@ package datadog.trace.agent.tooling.bytebuddy.profiling; import datadog.trace.util.CollectionUtils; -import java.util.Arrays; import java.util.HashMap; import java.util.Map; import net.bytebuddy.asm.AsmVisitorWrapper; @@ -82,7 +81,7 @@ public void visit( String signature, String superName, String[] interfaces) { - if (interfaces == null || !Arrays.asList(interfaces).contains(TASK_WRAPPER)) { + if (!CollectionUtils.contains(interfaces, TASK_WRAPPER)) { interfaces = CollectionUtils.append(interfaces, TASK_WRAPPER); if (signature != null) { signature += 'L' + TASK_WRAPPER + ';'; diff --git a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java index 83e2397bf76..590a140525f 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java @@ -6,7 +6,6 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.util.CollectionUtils; -import java.util.Arrays; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.field.FieldDescription; import net.bytebuddy.description.field.FieldList; @@ -89,7 +88,7 @@ public void visit( String[] interfaces) { // Add the Cloneable interface, unless it's already there (e.g. reapplying this // change while retransforming a class that was already modified on initial load) - if (interfaces == null || !Arrays.asList(interfaces).contains(CLONEABLE)) { + if (!CollectionUtils.contains(interfaces, CLONEABLE)) { interfaces = CollectionUtils.append(interfaces, CLONEABLE); if (signature != null) { signature += 'L' + CLONEABLE + ';'; diff --git a/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java b/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java index ad032186350..93dd501563f 100644 --- a/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java +++ b/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java @@ -91,17 +91,29 @@ private static MethodHandle findCopyOf(Class clazz, Class arg) { return null; } - /** Appends an element to an array, growing it as needed. */ - public static T[] append(@Nullable T[] array, T toAppend) { + /** Appends value to an array, growing it as needed; treats null arrays as empty. */ + public static T[] append(@Nullable T[] array, T value) { T[] appended; if (array == null) { //noinspection unchecked - appended = (T[]) Array.newInstance(toAppend.getClass(), 1); - appended[0] = toAppend; + appended = (T[]) Array.newInstance(value.getClass(), 1); + appended[0] = value; } else { appended = Arrays.copyOf(array, array.length + 1); - appended[array.length] = toAppend; + appended[array.length] = value; } return appended; } + + /** Checks whether an array contains an element; treats null arrays as empty. */ + public static boolean contains(@Nullable T[] array, T value) { + if (array != null) { + for (T element : array) { + if (value.equals(element)) { + return true; + } + } + } + return false; + } } diff --git a/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy b/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy index 06a97350b19..c9139964b84 100644 --- a/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy @@ -5,6 +5,7 @@ import datadog.trace.test.util.DDSpecification import static datadog.environment.JavaVirtualMachine.isJavaVersionAtLeast import static datadog.trace.util.CollectionUtils.append +import static datadog.trace.util.CollectionUtils.contains import static datadog.trace.util.CollectionUtils.tryMakeImmutableList import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap import static datadog.trace.util.CollectionUtils.tryMakeImmutableSet @@ -61,4 +62,15 @@ class CollectionUtilsTest extends DDSpecification { then: result == ["a", "b", "c"] as String[] } + + def "contains tolerates a null array"() { + expect: + !contains(null, "a") + } + + def "contains reports whether the value is present"() { + expect: + contains(["a", "b"] as String[], "a") + !contains(["a", "b"] as String[], "c") + } } From e5b225fb7ee24dac05f6486cb184c71899811a16 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 12:35:28 +0100 Subject: [PATCH 09/14] Static import utility methods --- .../trace/agent/tooling/MatchRecorder.java | 7 +++-- .../bytebuddy/iast/TaintableVisitor.java | 8 +++--- .../profiling/UnwrappingVisitor.java | 8 +++--- .../RequestImplInstrumentation.java | 7 ++--- .../datadog/trace/util/CollectionUtils.java | 4 +-- .../trace/util/CollectionUtilsTest.groovy | 26 +++++++++---------- 6 files changed, 32 insertions(+), 28 deletions(-) diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java index afe11415552..f93fdb8eb5a 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/MatchRecorder.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.ANY_CLASS_LOADER; import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed; +import static datadog.trace.util.CollectionUtils.arrayContains; import datadog.trace.agent.tooling.context.FieldBackedContextMatcher; -import datadog.trace.util.CollectionUtils; import java.util.BitSet; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -157,11 +157,10 @@ public void record( Class classBeingRedefined, BitSet matches) { // don't transform loaded classes unless they directly declare the marker - // - we must re-transform them to preserve the original structural change + // - we must re-transform those to preserve the original structural change if (matches.get(id) && null != classBeingRedefined - && !CollectionUtils.contains( - classBeingRedefined.getInterfaces(), structuralChangeMarker)) { + && !arrayContains(classBeingRedefined.getInterfaces(), structuralChangeMarker)) { matches.clear(id); } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java index 179054a2a84..916907afdc6 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java @@ -1,7 +1,9 @@ package datadog.trace.agent.tooling.bytebuddy.iast; +import static datadog.trace.util.CollectionUtils.appendToArray; +import static datadog.trace.util.CollectionUtils.arrayContains; + import datadog.trace.api.iast.Taintable; -import datadog.trace.util.CollectionUtils; import java.util.Arrays; import java.util.HashSet; import java.util.Set; @@ -79,8 +81,8 @@ public void visit( final String superName, String[] interfaces) { owner = name; - if (!CollectionUtils.contains(interfaces, TAINTABLE)) { - interfaces = CollectionUtils.append(interfaces, TAINTABLE); + if (!arrayContains(interfaces, TAINTABLE)) { + interfaces = appendToArray(interfaces, TAINTABLE); if (signature != null) { signature += 'L' + TAINTABLE + ';'; } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java index 8af589a443f..54b03196444 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java @@ -1,6 +1,8 @@ package datadog.trace.agent.tooling.bytebuddy.profiling; -import datadog.trace.util.CollectionUtils; +import static datadog.trace.util.CollectionUtils.appendToArray; +import static datadog.trace.util.CollectionUtils.arrayContains; + import java.util.HashMap; import java.util.Map; import net.bytebuddy.asm.AsmVisitorWrapper; @@ -81,8 +83,8 @@ public void visit( String signature, String superName, String[] interfaces) { - if (!CollectionUtils.contains(interfaces, TASK_WRAPPER)) { - interfaces = CollectionUtils.append(interfaces, TASK_WRAPPER); + if (!arrayContains(interfaces, TASK_WRAPPER)) { + interfaces = appendToArray(interfaces, TASK_WRAPPER); if (signature != null) { signature += 'L' + TASK_WRAPPER + ';'; } diff --git a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java index 590a140525f..cf9fca51ee6 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-redis-client/vertx-redis-client-3.9/src/main/java/datadog/trace/instrumentation/vertx_redis_client/RequestImplInstrumentation.java @@ -1,11 +1,12 @@ package datadog.trace.instrumentation.vertx_redis_client; +import static datadog.trace.util.CollectionUtils.appendToArray; +import static datadog.trace.util.CollectionUtils.arrayContains; import static net.bytebuddy.matcher.ElementMatchers.none; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.util.CollectionUtils; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.field.FieldDescription; import net.bytebuddy.description.field.FieldList; @@ -88,8 +89,8 @@ public void visit( String[] interfaces) { // Add the Cloneable interface, unless it's already there (e.g. reapplying this // change while retransforming a class that was already modified on initial load) - if (!CollectionUtils.contains(interfaces, CLONEABLE)) { - interfaces = CollectionUtils.append(interfaces, CLONEABLE); + if (!arrayContains(interfaces, CLONEABLE)) { + interfaces = appendToArray(interfaces, CLONEABLE); if (signature != null) { signature += 'L' + CLONEABLE + ';'; } diff --git a/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java b/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java index 93dd501563f..58a7431d612 100644 --- a/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java +++ b/internal-api/src/main/java/datadog/trace/util/CollectionUtils.java @@ -92,7 +92,7 @@ private static MethodHandle findCopyOf(Class clazz, Class arg) { } /** Appends value to an array, growing it as needed; treats null arrays as empty. */ - public static T[] append(@Nullable T[] array, T value) { + public static T[] appendToArray(@Nullable T[] array, T value) { T[] appended; if (array == null) { //noinspection unchecked @@ -106,7 +106,7 @@ public static T[] append(@Nullable T[] array, T value) { } /** Checks whether an array contains an element; treats null arrays as empty. */ - public static boolean contains(@Nullable T[] array, T value) { + public static boolean arrayContains(@Nullable T[] array, T value) { if (array != null) { for (T element : array) { if (value.equals(element)) { diff --git a/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy b/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy index c9139964b84..e9929cc68e1 100644 --- a/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/util/CollectionUtilsTest.groovy @@ -4,8 +4,8 @@ package datadog.trace.util import datadog.trace.test.util.DDSpecification import static datadog.environment.JavaVirtualMachine.isJavaVersionAtLeast -import static datadog.trace.util.CollectionUtils.append -import static datadog.trace.util.CollectionUtils.contains +import static datadog.trace.util.CollectionUtils.appendToArray +import static datadog.trace.util.CollectionUtils.arrayContains import static datadog.trace.util.CollectionUtils.tryMakeImmutableList import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap import static datadog.trace.util.CollectionUtils.tryMakeImmutableSet @@ -42,35 +42,35 @@ class CollectionUtilsTest extends DDSpecification { hopefullyImmutable.getClass().getName().contains(expectedClassName) } - def "append grows a null array"() { + def "appendToArray grows a null array"() { when: - String[] result = append(null, "a") + String[] result = appendToArray(null, "a") then: result == ["a"] as String[] } - def "append grows an empty array"() { + def "appendToArray grows an empty array"() { when: - String[] result = append(new String[0], "a") + String[] result = appendToArray(new String[0], "a") then: result == ["a"] as String[] } - def "append grows a non-empty array"() { + def "appendToArray grows a non-empty array"() { when: - String[] result = append(["a", "b"] as String[], "c") + String[] result = appendToArray(["a", "b"] as String[], "c") then: result == ["a", "b", "c"] as String[] } - def "contains tolerates a null array"() { + def "arrayContains tolerates a null array"() { expect: - !contains(null, "a") + !arrayContains(null, "a") } - def "contains reports whether the value is present"() { + def "arrayContains reports whether the value is present"() { expect: - contains(["a", "b"] as String[], "a") - !contains(["a", "b"] as String[], "c") + arrayContains(["a", "b"] as String[], "a") + !arrayContains(["a", "b"] as String[], "c") } } From 6e51a81b6b98d1c3f894f78f61712fa849b5faaa Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 12:46:13 +0100 Subject: [PATCH 10/14] Remove ENABLE_ADVICE_TRANSFORMER test workaround - it's not required anymore --- .../okhttp2/IastHttpUrlInstrumentation.java | 11 +---------- .../test/groovy/IastOkHttp2InstrumentationTest.groovy | 3 --- .../okhttp3/IastHttpUrlInstrumentation.java | 11 +---------- .../test/groovy/IastOkHttp3InstrumentationTest.groovy | 3 --- 4 files changed, 2 insertions(+), 26 deletions(-) diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java index 7b8b0ee2879..fa9bb7f28c7 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java @@ -22,13 +22,6 @@ public class IastHttpUrlInstrumentation extends InstrumenterModule.TaintableIast Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { - /** - * Adding fields to an already loaded class is not possible, during testing - * com.squareup.okhttp.HttpUrl gets loaded before the instrumenter kicks in, so we must disable - * the advice transformer or none of the transformations will be applied - */ - protected static boolean ENABLE_ADVICE_TRANSFORMER = true; - private final String className = IastHttpUrlInstrumentation.class.getName(); public IastHttpUrlInstrumentation() { @@ -47,9 +40,7 @@ public String muzzleDirective() { @Override public void typeAdvice(TypeTransformer transformer) { - if (ENABLE_ADVICE_TRANSFORMER) { - transformer.applyAdvice(new TaintableVisitor(instrumentedType())); - } + transformer.applyAdvice(new TaintableVisitor(instrumentedType())); } @Override diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy index c836d978c20..ba303ac6469 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy @@ -5,7 +5,6 @@ import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.CodecModule import datadog.trace.api.iast.propagation.PropagationModule import datadog.trace.api.iast.sink.SsrfModule -import datadog.trace.instrumentation.okhttp2.IastHttpUrlInstrumentation import spock.lang.AutoCleanup import spock.lang.Shared @@ -15,8 +14,6 @@ class IastOkHttp2InstrumentationTest extends InstrumentationSpecification { @Override protected void configurePreAgent() { - // HttpUrl gets loaded early so we have to disable the advice transformer - IastHttpUrlInstrumentation.ENABLE_ADVICE_TRANSFORMER = false injectSysConfig('dd.iast.enabled', 'true') } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java index f56eaf349c3..2c7d52c8c09 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java @@ -22,13 +22,6 @@ public class IastHttpUrlInstrumentation extends InstrumenterModule.TaintableIast Instrumenter.HasTypeAdvice, Instrumenter.HasMethodAdvice { - /** - * Adding fields to a loaded class is not possible, during testing okhttp3.HttpUrl gets loaded - * before the instrumenter kicks in, so we must disable the advice transformer or none of the - * transformations will be applied happen - */ - protected static boolean ENABLE_ADVICE_TRANSFORMER = true; - private final String className = IastHttpUrlInstrumentation.class.getName(); public IastHttpUrlInstrumentation() { @@ -42,9 +35,7 @@ public String instrumentedType() { @Override public void typeAdvice(TypeTransformer transformer) { - if (ENABLE_ADVICE_TRANSFORMER) { - transformer.applyAdvice(new TaintableVisitor(instrumentedType())); - } + transformer.applyAdvice(new TaintableVisitor(instrumentedType())); } @Override diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/IastOkHttp3InstrumentationTest.groovy b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/IastOkHttp3InstrumentationTest.groovy index e11e2bf2e99..bdb195dc308 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/IastOkHttp3InstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/IastOkHttp3InstrumentationTest.groovy @@ -3,7 +3,6 @@ import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.CodecModule import datadog.trace.api.iast.propagation.PropagationModule import datadog.trace.api.iast.sink.SsrfModule -import datadog.trace.instrumentation.okhttp3.IastHttpUrlInstrumentation import okhttp3.OkHttpClient import okhttp3.Request import spock.lang.AutoCleanup @@ -15,8 +14,6 @@ class IastOkHttp3InstrumentationTest extends InstrumentationSpecification { @Override protected void configurePreAgent() { - // HttpUrl gets loaded early so we have to disable the advice transformer - IastHttpUrlInstrumentation.ENABLE_ADVICE_TRANSFORMER = false injectSysConfig('dd.iast.enabled', 'true') // disable tracer metrics because it uses OkHttp and class loading is // not isolated in tests From 058842ae68329e5ea4deaa342dd51a87d58e8aaf Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 13:21:14 +0100 Subject: [PATCH 11/14] Test structural changes survive retransformation, and an already loaded class is not structurally changed --- ...turalChangeRetransformationForkedTest.java | 197 ++++++++++++++++++ 1 file changed, 197 insertions(+) create mode 100644 dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/StructuralChangeRetransformationForkedTest.java diff --git a/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/StructuralChangeRetransformationForkedTest.java b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/StructuralChangeRetransformationForkedTest.java new file mode 100644 index 00000000000..9b2857a94c9 --- /dev/null +++ b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/StructuralChangeRetransformationForkedTest.java @@ -0,0 +1,197 @@ +package datadog.trace.agent.tooling; + +import static datadog.trace.util.CollectionUtils.appendToArray; +import static datadog.trace.util.CollectionUtils.arrayContains; +import static net.bytebuddy.matcher.ElementMatchers.none; +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 datadog.trace.agent.tooling.bytebuddy.SharedTypePools; +import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers; +import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade; +import datadog.trace.agent.tooling.muzzle.ReferenceMatcher; +import java.lang.instrument.ClassFileTransformer; +import java.lang.instrument.Instrumentation; +import java.util.EnumSet; +import java.util.concurrent.atomic.AtomicInteger; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.agent.ByteBuddyAgent; +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.field.FieldDescription; +import net.bytebuddy.description.field.FieldList; +import net.bytebuddy.description.method.MethodList; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.loading.ClassLoadingStrategy; +import net.bytebuddy.implementation.Implementation; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.Type; +import net.bytebuddy.pool.TypePool; +import net.bytebuddy.utility.OpenedClassReader; +import org.junit.jupiter.api.Test; + +class StructuralChangeRetransformationForkedTest { + + private static final String TARGET_CLASS_NAME = + "datadog.trace.agent.tooling.GeneratedStructuralChangeTarget"; + + @Test + void reappliesStructuralChangeWhenRetransformingClassModifiedOnInitialLoad() throws Exception { + Instrumentation instrumentation = ByteBuddyAgent.install(); + + StructuralChangeInstrumentation structuralChange = new StructuralChangeInstrumentation(); + ClassFileTransformer transformer = installTransformer(instrumentation, structuralChange); + try { + Class targetClass = loadTargetClass(); + + assertEquals(1, structuralChange.transformations.get()); + assertTrue(StructuralChangeMarker.class.isAssignableFrom(targetClass)); + + instrumentation.retransformClasses(targetClass); + + assertEquals(2, structuralChange.transformations.get()); + assertTrue(StructuralChangeMarker.class.isAssignableFrom(targetClass)); + } finally { + instrumentation.removeTransformer(transformer); + } + } + + @Test + void skipsStructuralChangeForClassLoadedBeforeTransformerInstallation() throws Exception { + Instrumentation instrumentation = ByteBuddyAgent.install(); + + Class targetClass = loadTargetClass(); + + StructuralChangeInstrumentation structuralChange = new StructuralChangeInstrumentation(); + ClassFileTransformer transformer = installTransformer(instrumentation, structuralChange); + try { + assertEquals(0, structuralChange.transformations.get()); + assertFalse(StructuralChangeMarker.class.isAssignableFrom(targetClass)); + + instrumentation.retransformClasses(targetClass); + + assertEquals(0, structuralChange.transformations.get()); + assertFalse(StructuralChangeMarker.class.isAssignableFrom(targetClass)); + } finally { + instrumentation.removeTransformer(transformer); + } + } + + private ClassFileTransformer installTransformer( + Instrumentation instrumentation, StructuralChangeInstrumentation structuralChange) { + TypePoolFacade.registerAsSupplier(); + MemoizedMatchers.registerAsSupplier(); + + InstrumenterIndex instrumenterIndex = InstrumenterIndex.readIndex(); + InstrumenterState.initialize(instrumenterIndex.instrumentationCount()); + CombiningTransformerBuilder transformerBuilder = + new CombiningTransformerBuilder( + new AgentBuilder.Default() + .disableClassFormatChanges() + .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) + .ignore(none()), + instrumenterIndex, + EnumSet.of(InstrumenterModule.TargetSystem.TRACING), + false); + transformerBuilder.applyInstrumentation(structuralChange); + InstrumenterState.resetDefaultState(); + + ClassFileTransformer transformer = transformerBuilder.installOn(instrumentation); + SharedTypePools.endInstall(); + return transformer; + } + + private Class loadTargetClass() { + return new ByteBuddy() + .subclass(Object.class) + .name(TARGET_CLASS_NAME) + .make() + .load(getClass().getClassLoader(), ClassLoadingStrategy.Default.WRAPPER) + .getLoaded(); + } + + public interface StructuralChangeMarker {} + + private static final class StructuralChangeInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, + Instrumenter.WithStructuralChange, + Instrumenter.HasMethodAdvice { + + private final AtomicInteger transformations = new AtomicInteger(); + + private StructuralChangeInstrumentation() { + super("structural-change-retransformation-test"); + } + + @Override + public String instrumentedType() { + return TARGET_CLASS_NAME; + } + + @Override + public void typeAdvice(TypeTransformer transformer) { + transformer.applyAdvice( + (builder, typeDescription, classLoader, module, protectionDomain) -> { + transformations.incrementAndGet(); + return builder.visit(new AddStructuralMarkerVisitor()); + }); + } + + @Override + public void methodAdvice(MethodTransformer transformer) {} + + @Override + public Class structuralChangeMarker() { + return StructuralChangeMarker.class; + } + + public static class Muzzle { + public static ReferenceMatcher create() { + return ReferenceMatcher.NO_REFERENCES; + } + } + } + + private static final class AddStructuralMarkerVisitor implements AsmVisitorWrapper { + + private static final String MARKER_NAME = Type.getInternalName(StructuralChangeMarker.class); + + @Override + public int mergeWriter(int flags) { + return flags; + } + + @Override + public int mergeReader(int flags) { + return flags; + } + + @Override + public ClassVisitor wrap( + TypeDescription instrumentedType, + ClassVisitor classVisitor, + Implementation.Context implementationContext, + TypePool typePool, + FieldList fields, + MethodList methods, + int writerFlags, + int readerFlags) { + return new ClassVisitor(OpenedClassReader.ASM_API, classVisitor) { + @Override + public void visit( + int version, + int access, + String name, + String signature, + String superName, + String[] interfaces) { + if (!arrayContains(interfaces, MARKER_NAME)) { + interfaces = appendToArray(interfaces, MARKER_NAME); + } + super.visit(version, access, name, signature, superName, interfaces); + } + }; + } + } +} From cf801e9e9d760ce146cb39f8655c461efee83b86 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 23 Sep 2026 16:52:04 +0100 Subject: [PATCH 12/14] Document new Instrumenter.WithStructuralChange feature --- docs/how_instrumentations_work.md | 50 +++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docs/how_instrumentations_work.md b/docs/how_instrumentations_work.md index 97e97fa062a..83b600b4c16 100644 --- a/docs/how_instrumentations_work.md +++ b/docs/how_instrumentations_work.md @@ -893,6 +893,56 @@ If reflection must be used the reflection usage should be added to See [GraalVM configuration docs](https://www.graalvm.org/jdk17/reference-manual/native-image/dynamic-features/Reflection/#manual-configuration). +## Structural Changes (Adding Fields, Methods, or Interfaces) + +Some instrumentations use `Instrumenter.HasTypeAdvice` to change the *structure* of the type itself — for example adding a marker +interface or a field via a custom `AsmVisitorWrapper`. Unlike method advice this will fail if the target type is already loaded, +causing the instrumentation to silently stop working. The JVM does not allow transformations to change the structure once a type +is loaded, only the method bodies can be changed. Conversely, if we structurally changed a type before it was loaded then we must +remember to reapply that same change if the type is ever retransformed. + +This is hard to get right, so we provide a feature to correctly handle both situations. + +Instrumentations whose `typeAdvice()` adds fields, methods, or interfaces should implement `Instrumenter.WithStructuralChange`: + +```java +public interface WithStructuralChange extends HasTypeAdvice { + /** The marker interface added by the structural change, used to detect already-loaded types. */ + Class structuralChangeMarker(); +} +``` + +`structuralChangeMarker()` returns the marker interface the type advice adds directly to the type. For example IAST's +[`TaintableIast`](../dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java) +adds `Taintable` to every type it instruments, so it returns `Taintable.class` as the structural marker. + +### How it works + +1. When building matchers for a `WithStructuralChange` instrumentation, `CombiningTransformerBuilder` adds a + `MatchRecorder.PreserveLoadedStructure` narrowing matcher. +2. On a *fresh* class load there is no `classBeingRedefined`, so the matcher has no effect and the structural change + is applied as usual. +3. On a *retransform*, the matcher only allows the structural change to re-apply if the loaded class already + directly declares the marker interface — i.e. the change was already applied on first load, so it must be + re-applied. Otherwise the match is dropped, skipping the change instead of failing `retransformClasses()`. + Since the type might already have the marker, the `AsmVisitorWrapper` must guard against adding it twice + (see the `arrayContains`/`appendToArray` check in `TaintableVisitor`). +4. If the instrumentation also implements `HasMethodAdvice`, the structural change is split into its own + transformation (see `buildTypeAdvice()` in `CombiningTransformerBuilder`) so the narrowing can't disable the + *method* advice when the structural change is skipped. + +### When to use it + +Implement `WithStructuralChange` whenever `typeAdvice()` adds a field, method, or interface to the instrumented type. +Pick (or add) a marker interface that's only ever added by that structural change, since it's used to detect that the +change already happened. + +**Examples in the codebase:** +- [`InstrumenterModule.TaintableIast`](../dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java) + — instrumentation that adds `Taintable` to IAST-instrumented types. +- [`TaintableVisitor`](../dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableVisitor.java) + — the `AsmVisitorWrapper` that actually adds the marker interface during type advice. + ## JPMS Module Opening Java 9 introduced the Java Platform Module System (JPMS), which restricts reflective access across module boundaries. From 2b6ccd5b7055f13d5028bc076a34a6d708e52176 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Thu, 24 Sep 2026 14:49:19 +0100 Subject: [PATCH 13/14] Add test to make sure the TaintedMap fix doesn't regress in the future --- .../datadog/iast/taint/TaintedMapTest.groovy | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedMapTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedMapTest.groovy index ad7f7499d9d..83b57ad5abb 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedMapTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedMapTest.groovy @@ -72,6 +72,26 @@ class TaintedMapTest extends DDSpecification { map.count() == 0 } + def 'put deduplicates the last element of a bucket'() { + given: + // capacity 1 forces every entry into the same bucket chain + final map = new TaintedMap.TaintedMapImpl(1) + final a = new Object() + final b = new Object() + final c = new Object() + map.put(new TaintedObject(a, [] as Range[])) + map.put(new TaintedObject(b, [] as Range[])) + final originalC = new TaintedObject(c, [] as Range[]) + map.put(originalC) + + when: + map.put(new TaintedObject(c, [] as Range[])) + + then: + map.count() == 3 + map.get(c) == originalC + } + def 'last put always exists'() { given: int capacity = 256 From e0fe52c08fe79e21e0d6676cda9ad64a33199da4 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Thu, 24 Sep 2026 17:41:52 +0100 Subject: [PATCH 14/14] Document that WithStructuralChange's typeAdvice must not rely on injected helper classes or context-stores --- .../java/datadog/trace/agent/tooling/Instrumenter.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java index cb3b7e16cdc..8f53ccd4775 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java @@ -81,7 +81,12 @@ interface WithTypeStructure { ElementMatcher structureMatcher(); } - /** Instrumentation whose type advice adds fields, methods, or interfaces to the original type. */ + /** + * Instrumentation whose type advice adds fields, methods, or interfaces to the original type. + * + *

Note: {@link #typeAdvice} must not rely on injected helper classes or context-stores; those + * features are limited to bytecode inserted by {@link HasMethodAdvice#methodAdvice}. + */ interface WithStructuralChange extends HasTypeAdvice { /** The marker interface added by the structural change, used to detect already-loaded types. */ Class structuralChangeMarker();