Add IAST taint tracking for DB values

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

@@ -21,13 +21,16 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.Pair;
 import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.SourceTypes;
 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;
 import datadog.trace.instrumentation.iastinstrumenter.SourceMapperImpl;
 import datadog.trace.util.stacktrace.StackWalker;
+import java.util.Arrays;
 import java.util.Iterator;
+import java.util.List;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
@@ -39,6 +42,8 @@
 public abstract class SinkModuleBase {
 
   private static final int MAX_EVIDENCE_LENGTH = Config.get().getIastTruncationMaxValueLength();
+  private static final List<VulnerabilityType> DB_INJECTION_TYPES =
+      Arrays.asList(VulnerabilityType.SQL_INJECTION, VulnerabilityType.XSS);
 
   protected final OverheadController overheadController;
   protected final Reporter reporter;
@@ -143,9 +148,21 @@ protected final Evidence checkInjection(
       return null;
     }
 
+    // filter out ranges that are not SQL_TABLE for types that are not DB_INJECTION_TYPES
+    final Range[] filteredRanges;
+    if (!DB_INJECTION_TYPES.contains(type)) {
+      filteredRanges = Ranges.excludeRangesBySource(valueRanges, SourceTypes.SQL_TABLE);
+    } else {
+      filteredRanges = valueRanges;
+    }
+
+    if (filteredRanges == null || filteredRanges.length == 0) {
+      return null;
+    }
+
     final StringBuilder evidence = new StringBuilder();
     final RangeBuilder ranges = new RangeBuilder();
-    addToEvidence(type, evidence, ranges, value, valueRanges, evidenceBuilder);
+    addToEvidence(type, evidence, ranges, value, filteredRanges, evidenceBuilder);
 
     // check if finally we have an injection
     if (ranges.isEmpty()) {

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

@@ -426,4 +426,22 @@ public static Range[] splitRanges(
 
     return splittedRanges;
   }
+
+  /**
+   * Remove the ranges that have the same origin as the input source.
+   *
+   * @param ranges the ranges to filter
+   * @param source the byte value of the source to exclude (see {@link SourceTypes})
+   */
+  public static Range[] excludeRangesBySource(Range[] ranges, byte source) {
+    RangeBuilder newRanges = new RangeBuilder(ranges.length);
+
+    for (Range range : ranges) {
+      if (range.getSource().getOrigin() != source) {
+        newRanges.add(range);
+      }
+    }
+
+    return newRanges.toArray();
+  }
 }

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy

@@ -12,6 +12,8 @@ import static com.datadog.iast.util.HttpHeader.LOCATION
 import static com.datadog.iast.util.HttpHeader.REFERER
 import static datadog.trace.api.iast.SourceTypes.GRPC_BODY
 import static datadog.trace.api.iast.SourceTypes.REQUEST_HEADER_VALUE
+import static datadog.trace.api.iast.SourceTypes.REQUEST_QUERY
+import static datadog.trace.api.iast.SourceTypes.SQL_TABLE
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED
 import static com.datadog.iast.taint.Ranges.mergeRanges
 import static datadog.trace.api.iast.SourceTypes.REQUEST_HEADER_NAME
@@ -378,6 +380,22 @@ class RangesTest extends DDSpecification {
     1     | 3   | 2         | range(8, 8) | 0      | 0          | []
   }
 
+  void 'test excludeRangesBySource method'() {
+    when:
+    final result = Ranges.excludeRangesBySource(ranges as Range[], source)
+
+    then:
+    final expectedArray = expected as Range[]
+    result == expectedArray
+
+    where:
+    ranges                                          | source              | expected
+    [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | SQL_TABLE           | [range(5, 3)]
+    [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_HEADER_NAME | [rangeWithSource(0, 5, SQL_TABLE)]
+    [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_QUERY       | [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)]
+    []                                              | SQL_TABLE           | []
+  }
+
   Range[] rangesFromSpec(List<List<Object>> spec) {
     def ranges = new Range[spec.size()]
     int j = 0

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java

@@ -0,0 +1,117 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Source;
+import datadog.trace.api.iast.SourceTypes;
+import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.sql.ResultSet;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class IastResultSetInstrumentation extends InstrumenterModule.Iast
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public IastResultSetInstrumentation() {
+    super("jdbc", "jdbc-resultset", "iast-resultset");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "java.sql.ResultSet";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named("java.sql.ResultSet"));
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("next").and(takesArguments(0))),
+        IastResultSetInstrumentation.class.getName() + "$NextAdvice");
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("getString").or(named("getNString")))
+            .and(takesArguments(int.class).or(takesArguments(String.class))),
+        IastResultSetInstrumentation.class.getName() + "$GetParameterAdvice");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("java.sql.ResultSet", Integer.class.getName());
+  }
+
+  public static class NextAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.This final ResultSet resultSet) {
+      ContextStore<ResultSet, Integer> contextStore =
+          InstrumentationContext.get(ResultSet.class, Integer.class);
+      if (contextStore.get(resultSet) != null) {
+        contextStore.put(resultSet, contextStore.get(resultSet) + 1);
+      } else {
+        // first time
+        contextStore.put(resultSet, 1);
+      }
+    }
+  }
+
+  @RequiresRequestContext(RequestContextSlot.IAST)
+  public static class GetParameterAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      CallDepthThreadLocalMap.incrementCallDepth(ResultSet.class);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Source(SourceTypes.SQL_TABLE)
+    public static void onExit(
+        @Advice.Argument(0) Object argument,
+        @Advice.Return final String value,
+        @Advice.This final ResultSet resultSet,
+        @ActiveRequestContext RequestContext reqCtx) {
+      if (CallDepthThreadLocalMap.decrementCallDepth(ResultSet.class) > 0) {
+        return;
+      }
+      ContextStore<ResultSet, Integer> contextStore =
+          InstrumentationContext.get(ResultSet.class, Integer.class);
+      if (contextStore.get(resultSet) > Config.get().getIastDbRowsToTaint()) {
+        return;
+      }
+      if (value == null) {
+        return;
+      }
+      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      if (module == null) {
+        return;
+      }
+      IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
+      if (argument instanceof String) {
+        module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument);
+      } else {
+        module.taintString(ctx, value, SourceTypes.SQL_TABLE);
+      }
+    }
+  }
+}