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,13 +2,15 @@
 
 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.Dependencies;
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.Reporter;
 import com.datadog.iast.model.Evidence;
 import com.datadog.iast.model.Location;
 import com.datadog.iast.model.Range;
+import com.datadog.iast.model.Source;
 import com.datadog.iast.model.Vulnerability;
 import com.datadog.iast.model.VulnerabilityType;
 import com.datadog.iast.model.VulnerabilityType.InjectionType;
@@ -57,7 +59,7 @@ protected SinkModuleBase(@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;
   }
@@ -119,7 +121,7 @@ protected SinkModuleBase(@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;
   }
@@ -161,7 +163,7 @@ protected SinkModuleBase(@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;
   }
@@ -179,6 +181,24 @@ protected StackTraceElement getCurrentStackTrace() {
     return stackWalker.walk(SinkModuleBase::findValidPackageForVulnerability);
   }
 
+  protected Evidence buildEvidence(final Object value, final Range[] ranges) {
+    final Range unbound = Ranges.findUnbound(ranges);
+    if (unbound != null) {
+      final Source source = unbound.getSource();
+      if (source != null && source.getValue() != null) {
+        final String sourceValue = source.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(), source, NOT_MARKED)});
+        }
+      }
+    }
+    return new Evidence(value instanceof String ? (String) value : value.toString(), ranges);
+  }
+
   static StackTraceElement findValidPackageForVulnerability(
       @Nonnull final Stream<StackTraceElement> stream) {
     final StackTraceElement[] first = new StackTraceElement[1];

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

@@ -59,6 +59,15 @@ public static Range[] forObject(final @Nonnull Source source, final int mark) {
     return new Range[] {new Range(0, Integer.MAX_VALUE, source, mark)};
   }
 
+  @Nullable
+  public static Range findUnbound(@Nonnull final Range[] ranges) {
+    if (ranges.length != 1) {
+      return null;
+    }
+    final Range range = ranges[0];
+    return range.getStart() == 0 && range.getLength() == Integer.MAX_VALUE ? range : null;
+  }
+
   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,37 @@ class AbstractSinkModuleTest  extends IastModuleImplTestBase {
     result == expected
   }
 
+  void 'test reporting evidence on objects'() {
+    given:
+    overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span) >> true
+    final sink = new SinkModuleBase(dependencies) {}
+    final propagation = new PropagationModuleImpl()
+    final input = new String(source.value)
+    ctx.getTaintedObjects().taint(input, Ranges.forCharSequence(input, source))
+
+    when:
+    propagation.taintIfTainted(toReport, input)
+    final evidence = sink.checkInjection(span, ctx, SSRF, toReport)
+
+    then:
+    evidence.ranges.length == 1
+    final range = evidence.ranges[0]
+    if (matches) {
+      final taintedEvidence = evidence.value.substring(range.start, range.start + range.length)
+      taintedEvidence == input
+    } else {
+      final taintedEvidence = evidence.value
+      taintedEvidence != input
+    }
+
+    where:
+    source                                                    | toReport                                  | matches
+    new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URL('https://datadog.com/index.html') | true
+    new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URI('https://datadog.com/index.html') | true
+    new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URI('https://dAtAdOg.com/index.html') | false
+    new Source(REQUEST_PARAMETER_VALUE, 'url', 'datadog.com') | new URI('https://dAtAdOg.com/index.html') | false
+  }
+
   private StackTraceElement element(final String declaringClass) {
     return new StackTraceElement(declaringClass, "method", "fileName", 1)
   }