Extract interface for tainted objects

DataDog · Dec 6, 2024 · 471f00b · 471f00b
1 parent 0db6312
commit 471f00b
Show file tree

Hide file tree

Showing 315 changed files with 2,809 additions and 3,540 deletions.
diff --git a/...nt-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java b/...nt-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java
@@ -2,8 +2,9 @@
 
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
-import datadog.trace.api.iast.Taintable.Source;
 import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.api.iast.taint.Source;
+import datadog.trace.api.iast.taint.TaintedObjects;
 import datadog.trace.bootstrap.ContextStore;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import javax.annotation.Nonnull;
@@ -30,7 +31,8 @@ public static <E> NamedContext getOrCreate(
     }
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null) {
-      final Source source = module.findSource(target);
+      final TaintedObjects to = IastContext.Provider.taintedObjects();
+      final Source source = module.findSource(to, target);
       if (source != null) {
         result = new NamedContextImpl(module, source);
       }
@@ -60,7 +62,7 @@ private static class NamedContextImpl extends NamedContext {
     @Nullable private String currentName;
 
     private boolean fetched;
-    @Nullable private IastContext context;
+    @Nullable private TaintedObjects to;
 
     public NamedContextImpl(@Nonnull final PropagationModule module, @Nonnull final Source source) {
       this.module = module;
@@ -69,7 +71,7 @@ public NamedContextImpl(@Nonnull final PropagationModule module, @Nonnull final
 
     @Override
     public void taintValue(@Nullable final String value) {
-      module.taintString(iastCtx(), value, source.getOrigin(), currentName, source.getValue());
+      module.taintObject(to(), value, source.getOrigin(), currentName, source.getValue());
     }
 
     @Override
@@ -79,7 +81,7 @@ public void taintName(@Nullable final String name) {
       // prevent tainting the same name more than once
       if (currentName != name) {
         currentName = name;
-        module.taintString(iastCtx(), name, source.getOrigin(), name, source.getValue());
+        module.taintObject(to(), name, source.getOrigin(), name, source.getValue());
       }
     }
 
@@ -88,12 +90,12 @@ public void setCurrentName(@Nullable final String name) {
       currentName = name;
     }
 
-    private IastContext iastCtx() {
+    private TaintedObjects to() {
       if (!fetched) {
         fetched = true;
-        context = IastContext.Provider.get();
+        to = IastContext.Provider.taintedObjects();
       }
-      return context;
+      return to;
     }
   }
 }
diff --git a/...trap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy b/...trap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy
@@ -3,8 +3,8 @@ package datadog.trace.bootstrap.instrumentation.iast
 
 import datadog.trace.api.iast.InstrumentationBridge
 import datadog.trace.api.iast.SourceTypes
-import datadog.trace.api.iast.Taintable.Source
 import datadog.trace.api.iast.propagation.PropagationModule
+import datadog.trace.api.iast.taint.Source
 import datadog.trace.bootstrap.ContextStore
 import datadog.trace.test.util.DDSpecification
 
@@ -31,14 +31,14 @@ class NamedContextTest extends DDSpecification {
 
     then:
     1 * store.get(target) >> null
-    1 * module.findSource(target) >> source
+    1 * module.findSource(_, target) >> source
     1 * store.put(target, _)
 
     when:
     context.taintName(name)
 
     then:
-    1 * module.taintString(_, name, source.origin, name, source.value)
+    1 * module.taintObject(_, name, source.origin, name, source.value)
 
     when:
     context.taintName(name)
@@ -50,7 +50,7 @@ class NamedContextTest extends DDSpecification {
     context.taintValue(value)
 
     then:
-    1 * module.taintString(_, value, source.origin, name, source.value)
+    1 * module.taintObject(_, value, source.origin, name, source.value)
     0 * _
   }
 
@@ -62,7 +62,7 @@ class NamedContextTest extends DDSpecification {
     final ctx = NamedContext.getOrCreate(store, target)
 
     then:
-    1 * module.findSource(target) >> null
+    1 * module.findSource(_, target) >> null
     1 * store.put(target, _)
 
     when:
@@ -82,5 +82,20 @@ class NamedContextTest extends DDSpecification {
     byte origin
     String name
     String value
+
+    @Override
+    Source attachValue(Object newValue) {
+      return new SourceImpl(origin: origin, name: name, value: newValue as String)
+    }
+
+    @Override
+    boolean isReference() {
+      return false
+    }
+
+    @Override
+    Object getRawValue() {
+      return value
+    }
   }
 }
diff --git a/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/AbstractBenchmark.java b/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/AbstractBenchmark.java
@@ -3,14 +3,15 @@
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
 import com.datadog.iast.IastSystem;
-import com.datadog.iast.model.Range;
-import com.datadog.iast.model.Source;
-import com.datadog.iast.taint.TaintedObjects;
+import com.datadog.iast.model.SourceImpl;
 import datadog.trace.api.Config;
 import datadog.trace.api.ProductActivation;
 import datadog.trace.api.gateway.InstrumentationGateway;
 import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.taint.Range;
+import datadog.trace.api.iast.taint.Source;
+import datadog.trace.api.iast.taint.TaintedObjects;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -95,7 +96,7 @@ protected <E> E notTainted(final E value) {
   }
 
   protected Source source() {
-    return new Source((byte) 0, "key", "value");
+    return new SourceImpl((byte) 0, "key", "value");
   }
 
   private static long computeHash(final Object value) {

diff --git a/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderAppendBenchmark.java b/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderAppendBenchmark.java
@@ -3,7 +3,7 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -17,14 +17,14 @@ protected Context initializeContext() {
     final IastRequestContext context = new IastRequestContext();
     final String notTainted = notTainted("I am not a tainted string");
     final String tainted =
-        tainted(context, "I am a tainted string", new Range(5, 6, source(), NOT_MARKED));
+        tainted(context, "I am a tainted string", new RangeImpl(5, 6, source(), NOT_MARKED));
     final StringBuilder notTaintedBuilder =
         notTainted(new StringBuilder("I am not a tainted string builder"));
     final StringBuilder taintedBuilder =
         tainted(
             context,
             new StringBuilder("I am a tainted string builder"),
-            new Range(5, 6, source(), NOT_MARKED));
+            new RangeImpl(5, 6, source(), NOT_MARKED));
     return new Context(context, notTainted, tainted, notTaintedBuilder, taintedBuilder);
   }
 

diff --git a/...ent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderBatchBenchmark.java b/...ent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderBatchBenchmark.java
@@ -4,7 +4,7 @@
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
 import java.util.ArrayList;
@@ -36,7 +36,8 @@ protected StringBuilderBatchBenchmark.Context initializeContext() {
       final String value;
       if (current < limit) {
         value =
-            tainted(context, UUID.randomUUID().toString(), new Range(3, 6, source(), NOT_MARKED));
+            tainted(
+                context, UUID.randomUUID().toString(), new RangeImpl(3, 6, source(), NOT_MARKED));
       } else {
         value = notTainted(UUID.randomUUID().toString());
       }

diff --git a/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderInitBenchmark.java b/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderInitBenchmark.java
@@ -3,7 +3,7 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -17,7 +17,7 @@ protected Context initializeContext() {
     final IastRequestContext context = new IastRequestContext();
     final String notTainted = notTainted("I am not a tainted string");
     final String tainted =
-        tainted(context, "I am a tainted string", new Range(3, 6, source(), NOT_MARKED));
+        tainted(context, "I am a tainted string", new RangeImpl(3, 6, source(), NOT_MARKED));
     return new Context(context, notTainted, tainted);
   }
 

diff --git a/.../agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderToStringBenchmark.java b/.../agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderToStringBenchmark.java
@@ -3,7 +3,7 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -21,7 +21,7 @@ protected Context initializeContext() {
         tainted(
             context,
             new StringBuilder("I am a tainted string builder"),
-            new Range(5, 7, source(), NOT_MARKED));
+            new RangeImpl(5, 7, source(), NOT_MARKED));
     return new Context(context, notTaintedBuilder, taintedBuilder);
   }
 

diff --git a/...ava-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatBenchmark.java b/...ava-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatBenchmark.java
@@ -3,7 +3,7 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.instrumentation.java.lang.StringCallSite;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -16,7 +16,7 @@ protected StringConcatBenchmark.Context initializeContext() {
     final IastRequestContext context = new IastRequestContext();
     final String notTainted = notTainted("I am not a tainted string");
     final String tainted =
-        tainted(context, "I am a tainted string", new Range(3, 5, source(), NOT_MARKED));
+        tainted(context, "I am a tainted string", new RangeImpl(3, 5, source(), NOT_MARKED));
     return new StringConcatBenchmark.Context(context, notTainted, tainted);
   }
 

diff --git a/...ent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBatchBenchmark.java b/...ent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBatchBenchmark.java
@@ -4,7 +4,7 @@
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
 import java.lang.invoke.MethodHandle;
@@ -56,7 +56,7 @@ protected StringConcatFactoryBatchBenchmark.Context initializeContext() {
       double current = i / (double) stringCount;
       final String value;
       if (current < limit) {
-        value = tainted(context, "Yep, tainted", new Range(3, 5, source(), NOT_MARKED));
+        value = tainted(context, "Yep, tainted", new RangeImpl(3, 5, source(), NOT_MARKED));
       } else {
         value = notTainted("Nop, tainted");
       }

diff --git a/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBenchmark.java b/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBenchmark.java
@@ -3,7 +3,7 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
+import com.datadog.iast.model.RangeImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -16,7 +16,8 @@ public class StringConcatFactoryBenchmark
   protected StringConcatFactoryBenchmark.Context initializeContext() {
     final IastContext context = new IastRequestContext();
     final String notTainted = notTainted("Nop, tainted");
-    final String tainted = tainted(context, "Yep, tainted", new Range(3, 5, source(), NOT_MARKED));
+    final String tainted =
+        tainted(context, "Yep, tainted", new RangeImpl(3, 5, source(), NOT_MARKED));
     return new StringConcatFactoryBenchmark.Context(context, notTainted, tainted);
   }
 

diff --git a/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringJoinBenchmark.java b/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringJoinBenchmark.java
@@ -3,10 +3,11 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
-import com.datadog.iast.model.Source;
+import com.datadog.iast.model.RangeImpl;
+import com.datadog.iast.model.SourceImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.taint.Range;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.Fork;
 
@@ -23,7 +24,8 @@ protected StringJoinBenchmark.Context initializeContext() {
         .taint(
             tainted,
             new Range[] {
-              new Range(0, tainted.length(), new Source((byte) 0, "key", "value"), NOT_MARKED)
+              new RangeImpl(
+                  0, tainted.length(), new SourceImpl((byte) 0, "key", "value"), NOT_MARKED)
             });
 
     final String taintedDelimiter = new String("-");
@@ -32,8 +34,11 @@ protected StringJoinBenchmark.Context initializeContext() {
         .taint(
             taintedDelimiter,
             new Range[] {
-              new Range(
-                  0, taintedDelimiter.length(), new Source((byte) 1, "key", "value"), NOT_MARKED)
+              new RangeImpl(
+                  0,
+                  taintedDelimiter.length(),
+                  new SourceImpl((byte) 1, "key", "value"),
+                  NOT_MARKED)
             });
 
     return new StringJoinBenchmark.Context(

diff --git a/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringSubsequenceBenchmark.java b/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringSubsequenceBenchmark.java
@@ -3,10 +3,11 @@
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
-import com.datadog.iast.model.Range;
-import com.datadog.iast.model.Source;
+import com.datadog.iast.model.RangeImpl;
+import com.datadog.iast.model.SourceImpl;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.taint.Range;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.Fork;
 
@@ -29,7 +30,7 @@ protected StringSubsequenceBenchmark.Context initializeContext() {
         .taint(
             taintedLoseRange,
             new Range[] {
-              new Range(0, RANGE_SIZE, new Source((byte) 0, "key", "value"), NOT_MARKED)
+              new RangeImpl(0, RANGE_SIZE, new SourceImpl((byte) 0, "key", "value"), NOT_MARKED)
             });
 
     final String taintedModifyRange = new String(DEFAULT_STRING);
@@ -38,7 +39,7 @@ protected StringSubsequenceBenchmark.Context initializeContext() {
         .taint(
             taintedModifyRange,
             new Range[] {
-              new Range(1, RANGE_SIZE, new Source((byte) 1, "key", "value"), NOT_MARKED)
+              new RangeImpl(1, RANGE_SIZE, new SourceImpl((byte) 1, "key", "value"), NOT_MARKED)
             });
 
     return new StringSubsequenceBenchmark.Context(

diff --git a/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/taint/TaintedMapGetsBenchmark.java b/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/taint/TaintedMapGetsBenchmark.java
@@ -3,7 +3,7 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
-import com.datadog.iast.model.Range;
+import datadog.trace.api.iast.taint.Range;
 import java.util.ArrayList;
 import java.util.List;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -45,12 +45,12 @@ public void setup(BenchmarkParams params) {
     for (int i = 0; i < INITIAL_OP_COUNT; i++) {
       final Object k = new Object();
       initialObjectList.add(k);
-      map.put(new TaintedObject(k, new Range[0]));
+      map.put(new TaintedObjectEntry(k, new Range[0]));
     }
     for (int i = 0; i < OP_COUNT; i++) {
       final Object k = new Object();
       objectList.add(k);
-      map.put(new TaintedObject(k, new Range[0]));
+      map.put(new TaintedObjectEntry(k, new Range[0]));
     }
   }