Try to match tainted objects with sources when checking vulnerabilities with unbounded objects

Author: manuel-alvarez-alvarez

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

@@ -2,6 +2,7 @@
 
 import static com.datadog.iast.util.ObjectVisitor.State.CONTINUE;
 import static com.datadog.iast.util.ObjectVisitor.State.EXIT;
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.HasDependencies;
 import com.datadog.iast.IastRequestContext;
@@ -58,7 +59,7 @@ public void registerDependencies(@Nonnull final Dependencies dependencies) {
     if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
       return null;
     }
-    final Evidence result = new Evidence(value.toString(), ranges);
+    final Evidence result = buildEvidence(value, ranges);
     report(span, type, result);
     return result;
   }
@@ -117,7 +118,7 @@ public void registerDependencies(@Nonnull final Dependencies dependencies) {
       return null;
     }
 
-    final Evidence result = new Evidence(evidence, notMarkedRanges);
+    final Evidence result = buildEvidence(evidence, notMarkedRanges);
     report(span, type, result);
     return result;
   }
@@ -159,7 +160,7 @@ public void registerDependencies(@Nonnull final Dependencies dependencies) {
     if (notMarkedRanges == null || notMarkedRanges.length == 0) {
       return null;
     }
-    final Evidence result = new Evidence(evidence.toString(), notMarkedRanges);
+    final Evidence result = buildEvidence(evidence, notMarkedRanges);
     report(span, type, result);
     return result;
   }
@@ -177,8 +178,21 @@ protected StackTraceElement getCurrentStackTrace() {
     return stackWalker.walk(SinkModuleBase::findValidPackageForVulnerability);
   }
 
-  protected String getServiceName(final AgentSpan span) {
-    return span != null ? span.getServiceName() : null;
+  protected Evidence buildEvidence(final Object value, final Range[] ranges) {
+    if (Ranges.isUnbound(ranges)) {
+      final Range range = ranges[0];
+      if (range != null && range.getSource() != null && range.getSource().getValue() != null) {
+        final String sourceValue = range.getSource().getValue();
+        final String evidenceValue = value.toString();
+        final int start = evidenceValue.indexOf(sourceValue);
+        if (start >= 0) {
+          return new Evidence(
+              evidenceValue,
+              new Range[] {new Range(start, sourceValue.length(), range.getSource(), NOT_MARKED)});
+        }
+      }
+    }
+    return new Evidence(value instanceof String ? (String) value : value.toString(), ranges);
   }
 
   static StackTraceElement findValidPackageForVulnerability(

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java

@@ -59,6 +59,14 @@ public static Range[] forObject(final @Nonnull Source source, final int mark) {
     return new Range[] {new Range(0, Integer.MAX_VALUE, source, mark)};
   }
 
+  public static boolean isUnbound(@Nonnull final Range[] ranges) {
+    if (ranges.length != 1) {
+      return false;
+    }
+    final Range range = ranges[0];
+    return range.getStart() == 0 && range.getLength() == Integer.MAX_VALUE;
+  }
+
   public static void copyShift(
       final @Nonnull Range[] src, final @Nonnull Range[] dst, final int dstPos, final int shift) {
     copyShift(src, dst, dstPos, shift, src.length);

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

@@ -1,13 +1,37 @@
 package com.datadog.iast.sink
 
 import com.datadog.iast.IastModuleImplTestBase
+import com.datadog.iast.IastRequestContext
+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.gateway.RequestContext
+import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 
-class AbstractSinkModuleTest  extends IastModuleImplTestBase {
+import static com.datadog.iast.model.VulnerabilityType.SSRF
+import static datadog.trace.api.iast.SourceTypes.REQUEST_PARAMETER_VALUE
+
+class AbstractSinkModuleTest extends IastModuleImplTestBase {
 
   final StackTraceElement ignoredPackageClassElement = element("org.springframework.Ignored")
   final StackTraceElement notIgnoredPackageClassElement = element("datadog.smoketest.NotIgnored")
   final StackTraceElement notInIastExclusionTrie = element("not.in.iast.exclusion.Class")
 
+  private IastRequestContext ctx
+  private AgentSpan span
+
+  void setup() {
+    ctx = new IastRequestContext()
+    final reqCtx = Mock(RequestContext) {
+      getData(RequestContextSlot.IAST) >> ctx
+    }
+    span = Mock(AgentSpan) {
+      getRequestContext() >> reqCtx
+    }
+    tracer.activeSpan() >> span
+  }
 
   void 'filter ignored package element from stack'() {
 
@@ -43,6 +67,31 @@ class AbstractSinkModuleTest  extends IastModuleImplTestBase {
     result == expected
   }
 
+  void 'test reporting evidence on objects'() {
+    given:
+    final sink = registerDependencies(new SinkModuleBase() {})
+    final propagation = new PropagationModuleImpl()
+    final input = new String(source.value)
+    ctx.getTaintedObjects().taint(source.value, Ranges.forCharSequence(input, source))
+
+    when:
+    propagation.taintIfTainted(toReport, input)
+    final evidence = sink.checkInjection(span, ctx, SSRF, toReport)
+
+    then:
+    1 * overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span) >> true
+    evidence.ranges.length == 1
+
+    final range = evidence.ranges[0]
+    final taintedEvidence = evidence.value.toString().substring(range.start, range.start + range.length)
+    taintedEvidence == input
+
+    where:
+    source                                                    | toReport
+    new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URL('https://datadog.com/index.html')
+    new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URI('https://datadog.com/index.html')
+  }
+
   private StackTraceElement element(final String declaringClass) {
     return new StackTraceElement(declaringClass, "method", "fileName", 1)
   }