Replace changing source for CallDepthThreadLocalMap approach

Author: Mariovido

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

@@ -1,7 +1,6 @@
 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;
@@ -669,39 +668,6 @@ 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) {
@@ -818,32 +784,6 @@ 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/taint/Ranges.java

@@ -144,20 +144,6 @@ 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}

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

@@ -19,6 +19,7 @@
 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;
@@ -78,13 +79,21 @@ public static void onExit(@Advice.This final ResultSet resultSet) {
 
   @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()) {
@@ -99,11 +108,7 @@ public static void onExit(
       }
       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);
-        }
+        module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument);
       } else {
         module.taintString(ctx, value, SourceTypes.SQL_TABLE);
       }

internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java

@@ -347,11 +347,4 @@ int taintObjectDeeply(
   Source findSource(@Nullable IastContext ctx, @Nullable Object target);
 
   void markIfTainted(@Nullable Object target, int mark);
-
-  void changeSource(@Nullable Object target, byte origin);
-
-  void changeSource(@Nullable Object target, byte origin, @Nullable CharSequence name);
-
-  void changeSource(
-      @Nullable IastContext ctx, @Nullable Object target, byte origin, @Nullable CharSequence name);
 }