DataDog · mcculls · Sep 23, 2026 · Sep 23, 2026 · Sep 23, 2026 · Sep 23, 2026
@@ -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;
         }

@@ -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

@@ -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<InstrumenterModule.TargetSystem> 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);

@@ -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
   }
 

@@ -2,6 +2,7 @@
 
 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 java.util.BitSet;
@@ -140,6 +141,56 @@ 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 directly declare the marker
+      // - we must re-transform those to preserve the original structural change
+      if (matches.get(id)
+          && null != classBeingRedefined
+          && !arrayContains(classBeingRedefined.getInterfaces(), structuralChangeMarker)) {
+        matches.clear(id);
+      }
+    }
+  }
+
+  /** 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<ClassLoader> matcher;