Add IAST taint tracking for DB values

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java

@@ -1,6 +1,7 @@
 package com.datadog.iast.propagation;
 
 import static com.datadog.iast.model.Source.PROPAGATION_PLACEHOLDER;
+import static com.datadog.iast.taint.Ranges.changeHighestPriorityRange;
 import static com.datadog.iast.taint.Ranges.highestPriorityRange;
 import static com.datadog.iast.util.ObjectVisitor.State.CONTINUE;
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
@@ -668,6 +669,39 @@ public void markIfTainted(@Nullable Object target, int mark) {
     }
   }
 
+  @Override
+  public void changeSource(@Nullable Object target, byte origin) {
+    if (target == null) {
+      return;
+    }
+    changeSource(target, origin, null);
+  }
+
+  @Override
+  public void changeSource(@Nullable Object target, byte origin, @Nullable CharSequence name) {
+    if (target == null) {
+      return;
+    }
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    changeSource(ctx, target, origin, name);
+  }
+
+  @Override
+  public void changeSource(
+      @Nullable final IastContext ctx,
+      @Nullable Object target,
+      byte origin,
+      @Nullable CharSequence name) {
+    if (ctx == null || target == null) {
+      return;
+    }
+    final TaintedObjects to = ctx.getTaintedObjects();
+    changeHighestPrioritySource(to, target, origin, name);
+  }
+
   @Override
   public boolean isTainted(@Nullable final Object target) {
     if (target == null) {
@@ -784,6 +818,32 @@ private static Source highestPrioritySource(
     }
   }
 
+  private static void changeHighestPrioritySource(
+      final @Nonnull TaintedObjects to,
+      final @Nonnull Object object,
+      final byte origin,
+      @Nullable final CharSequence name) {
+    Source previousValue = highestPrioritySource(to, object);
+    if (previousValue == null) {
+      return;
+    }
+    Source newSource = newSource(object, origin, name, previousValue.getValue());
+    if (object instanceof Taintable) {
+      ((Taintable) object).$$DD$setSource(newSource);
+    } else {
+      TaintedObject taintedObject = to.get(object);
+      if (taintedObject == null) {
+        return;
+      }
+      final Range[] ranges = getRanges(to, object);
+      if (ranges == null || ranges.length == 0) {
+        return;
+      }
+      changeHighestPriorityRange(ranges, newSource);
+      taintedObject.setRanges(ranges);
+    }
+  }
+
   private static void internalTaint(
       @Nonnull final TaintedObjects to,
       @Nonnull final Object value,

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

@@ -144,6 +144,20 @@ public static Range highestPriorityRange(@Nonnull final Range[] ranges) {
     return ranges[0];
   }
 
+  public static void changeHighestPriorityRange(
+      @Nonnull final Range[] ranges, @Nonnull final Source source) {
+    for (int i = 0; i < ranges.length; i++) {
+      if (ranges[i].getMarks() == NOT_MARKED) {
+        Range newRange =
+            new Range(ranges[i].getStart(), ranges[i].getLength(), source, ranges[i].getMarks());
+        ranges[i] = newRange;
+      }
+    }
+    Range newRange =
+        new Range(ranges[0].getStart(), ranges[0].getLength(), source, ranges[0].getMarks());
+    ranges[0] = newRange;
+  }
+
   /**
    * Checks if all ranges are coming from any header, in case no ranges are provided it will return
    * {@code true}
@@ -426,4 +440,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,112 @@
+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.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 {
+
+  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.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) {
+      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) {
+        if (module.isTainted(value)) {
+          module.changeSource(ctx, value, SourceTypes.SQL_TABLE, (String) argument);
+        } else {
+          module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument);
+        }
+      } else {
+        module.taintString(ctx, value, SourceTypes.SQL_TABLE);
+      }
+    }
+  }
+}