Add experimental taint propagation to the String replace, replaceFirst, replaceAll methods

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

@@ -13,8 +13,10 @@
 import com.datadog.iast.taint.TaintedObjects;
 import com.datadog.iast.util.RangeBuilder;
 import com.datadog.iast.util.Ranged;
+import com.datadog.iast.util.StringUtils;
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.propagation.StringModule;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.Arrays;
 import java.util.Deque;
@@ -631,6 +633,110 @@ public void onIndent(@Nonnull String self, int indentation, @Nonnull String resu
     }
   }
 
+  @Override
+  @SuppressFBWarnings("ES_COMPARING_PARAMETER_STRING_WITH_EQ")
+  public void onStringReplace(
+      @Nonnull String self, char oldChar, char newChar, @Nonnull String result) {
+    if (self == result || !canBeTainted(result)) {
+      return;
+    }
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    final TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    final TaintedObject taintedSelf = taintedObjects.get(self);
+    if (taintedSelf == null) {
+      return;
+    }
+
+    final Range[] rangesSelf = taintedSelf.getRanges();
+    if (rangesSelf.length == 0) {
+      return;
+    }
+
+    taintedObjects.taint(result, rangesSelf);
+  }
+
+  @Override
+  @SuppressFBWarnings("ES_COMPARING_PARAMETER_STRING_WITH_EQ")
+  public String onStringReplace(
+      @Nonnull String self, CharSequence oldCharSeq, CharSequence newCharSeq) {
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return self.replace(oldCharSeq, newCharSeq);
+    }
+    final TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    final TaintedObject taintedSelf = taintedObjects.get(self);
+    Range[] rangesSelf = new Range[0];
+    if (taintedSelf != null) {
+      rangesSelf = taintedSelf.getRanges();
+    }
+
+    final TaintedObject taintedInput = taintedObjects.get(newCharSeq);
+    Range[] rangesInput = null;
+    if (taintedInput != null) {
+      rangesInput = taintedInput.getRanges();
+    }
+
+    if (rangesSelf.length == 0 && rangesInput == null) {
+      return self.replace(oldCharSeq, newCharSeq);
+    }
+
+    return StringUtils.replaceAndTaint(
+        taintedObjects,
+        self,
+        Pattern.compile((String) oldCharSeq),
+        (String) newCharSeq,
+        rangesSelf,
+        rangesInput,
+        Integer.MAX_VALUE);
+  }
+
+  @Override
+  @SuppressForbidden
+  public String onStringReplace(
+      @Nonnull String self, String regex, String replacement, int numReplacements) {
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      if (numReplacements > 1) {
+        return self.replaceAll(regex, replacement);
+      } else {
+        return self.replaceFirst(regex, replacement);
+      }
+    }
+
+    final TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    final TaintedObject taintedSelf = taintedObjects.get(self);
+    Range[] rangesSelf = new Range[0];
+    if (taintedSelf != null) {
+      rangesSelf = taintedSelf.getRanges();
+    }
+
+    final TaintedObject taintedInput = taintedObjects.get(replacement);
+    Range[] rangesInput = null;
+    if (taintedInput != null) {
+      rangesInput = taintedInput.getRanges();
+    }
+
+    if (rangesSelf.length == 0 && rangesInput == null) {
+      if (numReplacements > 1) {
+        return self.replaceAll(regex, replacement);
+      } else {
+        return self.replaceFirst(regex, replacement);
+      }
+    }
+
+    return StringUtils.replaceAndTaint(
+        taintedObjects,
+        self,
+        Pattern.compile(regex),
+        replacement,
+        rangesSelf,
+        rangesInput,
+        numReplacements);
+  }
+
   /**
    * Adds the tainted ranges belonging to the current parameter added via placeholder taking care of
    * an optional tainted placeholder.

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

@@ -409,12 +409,11 @@ private static int updateRangesWithIndentation(
         if (range.getStart() + range.getLength() > end) {
           newLength -= lineOffset;
         }
-        newRanges[i] = new Range(newStart, newLength, range.getSource(), range.getMarks());
+        newRanges[i] = copyWithPosition(range, newStart, newLength);
       } else if (range.getStart() + range.getLength() >= start) {
         final Range newRange = newRanges[i];
         final int newLength = newRange.getLength() + indentation;
-        newRanges[i] =
-            new Range(newRange.getStart(), newLength, newRange.getSource(), newRange.getMarks());
+        newRanges[i] = copyWithPosition(newRange, newRange.getStart(), newLength);
       }
 
       if (range.getStart() + range.getLength() - 1 <= end) {
@@ -426,4 +425,41 @@ private static int updateRangesWithIndentation(
 
     return rangeStart;
   }
+
+  /**
+   * Split the range in two taking into account the new length of the characters.
+   *
+   * <p>In case start and end are out of the range, it will return the range without splitting but
+   * taking into account the offset. In the case that the new length is less than or equal to 0, it
+   * will return an empty array.
+   *
+   * @param start is the start of the character sequence
+   * @param end is the end of the character sequence
+   * @param newLength is the new length of the character sequence
+   * @param range is the range to split
+   * @param offset is the offset to apply to the range
+   * @param diffLength is the difference between the new length and the old length
+   */
+  public static Range[] splitRanges(
+      int start, int end, int newLength, Range range, int offset, int diffLength) {
+    start += offset;
+    end += offset;
+    int rangeStart = range.getStart() + offset;
+    int rangeEnd = rangeStart + range.getLength() + diffLength;
+
+    int firstLength = start - rangeStart;
+    int secondLength = range.getLength() - firstLength - newLength + diffLength;
+    if (rangeStart > end || rangeEnd <= start) {
+      if (firstLength <= 0) {
+        return new Range[0];
+      }
+      return new Range[] {copyWithPosition(range, rangeStart, firstLength)};
+    }
+
+    Range[] splittedRanges = new Range[2];
+    splittedRanges[0] = copyWithPosition(range, rangeStart, firstLength);
+    splittedRanges[1] = copyWithPosition(range, rangeEnd - secondLength, secondLength);
+
+    return splittedRanges;
+  }
 }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/StringUtils.java

@@ -1,6 +1,12 @@
 package com.datadog.iast.util;
 
+import com.datadog.iast.model.Range;
+import com.datadog.iast.taint.Ranges;
+import com.datadog.iast.taint.TaintedObjects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 public abstract class StringUtils {
 
@@ -48,4 +54,171 @@ public static int leadingWhitespaces(@Nonnull final String value, int start, int
     }
     return whitespaces;
   }
+
+  /**
+   * Returns the string replaced with the regex and the values tainted (if needed)
+   *
+   * @param taintedObjects the ctx object to save the range of the tainted values
+   * @param target the string to be replaced
+   * @param pattern the pattern to be replaced
+   * @param replacement the replacement string
+   * @param ranges the ranges of the string to be replaced
+   * @param rangesInput the ranges of the input string
+   * @param numOfReplacements the number of replacements to be made
+   */
+  @Nonnull
+  public static String replaceAndTaint(
+      @Nonnull TaintedObjects taintedObjects,
+      @Nonnull final String target,
+      Pattern pattern,
+      String replacement,
+      Range[] ranges,
+      @Nullable Range[] rangesInput,
+      int numOfReplacements) {
+    if (numOfReplacements <= 0) {
+      return target;
+    }
+    Matcher matcher = pattern.matcher(target);
+    boolean result = matcher.find();
+    if (result) {
+      int offset = 0;
+
+      int newRangesSize = ranges.length * 2;
+      if (rangesInput != null) {
+        newRangesSize += rangesInput.length;
+      }
+      RangeBuilder newRanges = new RangeBuilder(newRangesSize);
+
+      int firstRange = 0;
+      int newLength = replacement.length();
+
+      StringBuffer sb = new StringBuffer();
+      do {
+        int start = matcher.start();
+        int end = matcher.end();
+        int diffLength = newLength - (end - start);
+
+        boolean rangesAdded = false;
+        while (firstRange < ranges.length) {
+          Range range = ranges[firstRange];
+          int rangeStart = range.getStart();
+          int rangeEnd = rangeStart + range.getLength();
+          // If the replaced value is between one range
+          if (rangeStart <= start && rangeEnd >= end) {
+            Range[] splittedRanges =
+                Ranges.splitRanges(start, end, newLength, range, offset, diffLength);
+
+            if (splittedRanges.length > 0 && splittedRanges[0].getLength() > 0) {
+              newRanges.add(splittedRanges[0]);
+            }
+
+            if (rangesInput != null) {
+              for (Range rangeInput : rangesInput) {
+                newRanges.add(
+                    Ranges.copyWithPosition(
+                        rangeInput,
+                        start + rangeInput.getStart() + offset,
+                        rangeInput.getLength()));
+              }
+              rangesAdded = true;
+            }
+
+            if (splittedRanges.length > 1 && splittedRanges[1].getLength() > 0) {
+              newRanges.add(splittedRanges[1]);
+            }
+
+            firstRange++;
+            break;
+            // If the replaced value starts in the range and not end there
+          } else if (rangeStart <= start && rangeEnd > start) {
+            Range[] splittedRanges =
+                Ranges.splitRanges(start, end, newLength, range, offset, diffLength);
+
+            if (splittedRanges.length > 0 && splittedRanges[0].getLength() > 0) {
+              newRanges.add(splittedRanges[0]);
+            }
+
+            if (rangesInput != null && !rangesAdded) {
+              for (Range rangeInput : rangesInput) {
+                newRanges.add(
+                    Ranges.copyWithPosition(
+                        rangeInput,
+                        start + rangeInput.getStart() + offset,
+                        rangeInput.getLength()));
+              }
+              rangesAdded = true;
+            }
+
+            // If the replaced value ends in the range
+          } else if (rangeEnd >= end) {
+            Range[] splittedRanges =
+                Ranges.splitRanges(start, end, newLength, range, offset, diffLength);
+
+            if (rangesInput != null && !rangesAdded) {
+              for (Range rangeInput : rangesInput) {
+                newRanges.add(
+                    Ranges.copyWithPosition(
+                        rangeInput,
+                        start + rangeInput.getStart() + offset,
+                        rangeInput.getLength()));
+              }
+              rangesAdded = true;
+            }
+
+            if (splittedRanges.length > 1 && splittedRanges[1].getLength() > 0) {
+              newRanges.add(splittedRanges[1]);
+            }
+
+            firstRange++;
+            break;
+            // Middle ranges
+          } else if (rangeStart >= start) {
+            firstRange++;
+            continue;
+          } else {
+            newRanges.add(
+                Ranges.copyWithPosition(range, range.getStart() + offset, range.getLength()));
+          }
+
+          firstRange++;
+        }
+
+        // In case there are no ranges
+        if (rangesInput != null && !rangesAdded) {
+          for (Range rangeInput : rangesInput) {
+            newRanges.add(
+                Ranges.copyWithPosition(
+                    rangeInput, start + rangeInput.getStart() + offset, rangeInput.getLength()));
+          }
+        }
+
+        matcher.appendReplacement(sb, replacement);
+
+        offset = diffLength;
+        numOfReplacements--;
+        if (numOfReplacements > 0) {
+          result = matcher.find();
+        }
+      } while (result && numOfReplacements > 0);
+
+      // In the case there is no tainted object
+      if (firstRange < ranges.length) {
+        for (int i = firstRange; i < ranges.length; i++) {
+          newRanges.add(
+              Ranges.copyWithPosition(
+                  ranges[i], ranges[i].getStart() + offset, ranges[i].getLength()));
+        }
+      }
+
+      matcher.appendTail(sb);
+      String finalString = sb.toString();
+      Range[] finalRanges = newRanges.toArray();
+      if (finalRanges.length > 0) {
+        taintedObjects.taint(finalString, finalRanges);
+      }
+      return finalString;
+    }
+
+    return target;
+  }
 }