Ensure vulnerabilities are reported with taintable values (#7801)

Author: manuel-alvarez-alvarez

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Source.java

@@ -52,6 +52,18 @@ public String getValue() {
     return asString(value);
   }
 
+  /** This will expose the internal reference so be careful */
+  @Nullable
+  public Object getRawValue() {
+    if (value == null) {
+      return null;
+    }
+    if (value instanceof Reference<?>) {
+      return ((Reference<?>) value).get();
+    }
+    return value;
+  }
+
   @SuppressFBWarnings(
       value = "DM_STRING_CTOR",
       justification = "New string instance requires constructor")

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java

@@ -21,6 +21,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.Pair;
 import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.Taintable;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.instrumentation.iastinstrumenter.IastExclusionTrie;
@@ -108,17 +109,35 @@ protected final Evidence checkInjection(
   protected final Evidence checkInjection(
       final IastContext ctx,
       final VulnerabilityType type,
-      final Object value,
+      Object value,
       @Nullable final EvidenceBuilder evidenceBuilder,
       @Nullable final LocationSupplier locationSupplier) {
 
     final TaintedObjects to = ctx.getTaintedObjects();
-    final TaintedObject tainted = to.get(value);
-    if (tainted == null) {
-      return null;
+    final Range[] valueRanges;
+    if (value instanceof Taintable) {
+      final Taintable taintable = (Taintable) value;
+      if (!taintable.$DD$isTainted()) {
+        return null;
+      }
+      final Source source = (Source) taintable.$$DD$getSource();
+      final Object origin = source.getRawValue();
+      final TaintedObject tainted = origin == null ? null : to.get(origin);
+      if (origin != null && tainted != null) {
+        valueRanges = Ranges.getNotMarkedRanges(tainted.getRanges(), type.mark());
+        value = origin;
+      } else {
+        valueRanges = Ranges.forObject((Source) taintable.$$DD$getSource(), type.mark());
+        value = String.format("Tainted reference detected in " + value.getClass());
+      }
+    } else {
+      final TaintedObject tainted = to.get(value);
+      if (tainted == null) {
+        return null;
+      }
+      valueRanges = Ranges.getNotMarkedRanges(tainted.getRanges(), type.mark());
     }
 
-    final Range[] valueRanges = Ranges.getNotMarkedRanges(tainted.getRanges(), type.mark());
     if (valueRanges == null || valueRanges.length == 0) {
       return null;
     }

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/sink/AbstractSinkModuleTest.groovy

@@ -5,6 +5,9 @@ import com.datadog.iast.model.Source
 import com.datadog.iast.overhead.Operations
 import com.datadog.iast.propagation.PropagationModuleImpl
 import com.datadog.iast.taint.Ranges
+import datadog.trace.api.iast.Taintable
+
+import java.lang.ref.WeakReference
 
 import static com.datadog.iast.model.VulnerabilityType.SSRF
 import static datadog.trace.api.iast.SourceTypes.REQUEST_PARAMETER_VALUE
@@ -80,7 +83,63 @@ class AbstractSinkModuleTest extends IastModuleImplTestBase {
     new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URI('https://dAtAdOg.com/index.html') | false
   }
 
+  void 'test reporting with taintables'() {
+    setup:
+    final sink = new SinkModuleBase(dependencies) {}
+    final value = 'datadog.com'
+    final valueRef = new WeakReference<>(value)
+    final source = new Source(REQUEST_PARAMETER_VALUE, 'url', valueRef)
+
+    and:
+    final taintable = new MockTaintable(source: source)
+
+    when: 'original source value is not tainted'
+    def evidence = sink.checkInjection(SSRF, taintable)
+
+    then: 'a fallback evidence is provided'
+    evidence.value == "Tainted reference detected in " + taintable.class
+    evidence.ranges.length == 1
+    evidence.ranges[0].start == 0
+    evidence.ranges[0].length == evidence.value.length()
+
+    when: 'original source value is tainted'
+    ctx.getTaintedObjects().taint(value, Ranges.forCharSequence(value, source))
+    evidence = sink.checkInjection(SSRF, taintable)
+
+    then: 'the proper value is set in the evidence'
+    evidence.value == value
+    evidence.ranges.length == 1
+    evidence.ranges[0].start == 0
+    evidence.ranges[0].length == value.length()
+
+    when: 'original source cleared by the GC'
+    valueRef.clear()
+    evidence = sink.checkInjection(SSRF, taintable)
+
+    then: 'a fallback evidence is provided'
+    evidence.value == "Tainted reference detected in " + taintable.class
+    evidence.ranges.length == 1
+    evidence.ranges[0].start == 0
+    evidence.ranges[0].length == evidence.value.length()
+  }
+
   private StackTraceElement element(final String declaringClass) {
     return new StackTraceElement(declaringClass, "method", "fileName", 1)
   }
+
+  private static class MockTaintable implements Taintable {
+    private Source source
+
+    @SuppressWarnings('CodeNarc')
+    @Override
+    Source $$DD$getSource() {
+      return source
+    }
+
+    @SuppressWarnings('CodeNarc')
+    @Override
+    void $$DD$setSource(Source source) {
+      this.source = source
+    }
+  }
 }