Skip to content

Commit d3151b6

Browse files
MariovidoMariovido
authored
Add experimental taint propagation to the String replace, replaceFirst, replaceAll methods (#7741)
1 parent 8b31030 commit d3151b6

File tree

18 files changed

+890
-28
lines changed

18 files changed

+890
-28
lines changed

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

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313
import com.datadog.iast.taint.TaintedObjects;
1414
import com.datadog.iast.util.RangeBuilder;
1515
import com.datadog.iast.util.Ranged;
16+
import com.datadog.iast.util.StringUtils;
1617
import datadog.trace.api.iast.IastContext;
1718
import datadog.trace.api.iast.propagation.StringModule;
19+
import de.thetaphi.forbiddenapis.SuppressForbidden;
1820
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
1921
import java.util.Arrays;
2022
import java.util.Deque;
@@ -631,6 +633,115 @@ public void onIndent(@Nonnull String self, int indentation, @Nonnull String resu
631633
}
632634
}
633635

636+
@Override
637+
@SuppressFBWarnings("ES_COMPARING_PARAMETER_STRING_WITH_EQ")
638+
public void onStringReplace(
639+
@Nonnull String self, char oldChar, char newChar, @Nonnull String result) {
640+
if (self == result || !canBeTainted(result)) {
641+
return;
642+
}
643+
final IastContext ctx = IastContext.Provider.get();
644+
if (ctx == null) {
645+
return;
646+
}
647+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
648+
final TaintedObject taintedSelf = taintedObjects.get(self);
649+
if (taintedSelf == null) {
650+
return;
651+
}
652+
653+
final Range[] rangesSelf = taintedSelf.getRanges();
654+
if (rangesSelf.length == 0) {
655+
return;
656+
}
657+
658+
taintedObjects.taint(result, rangesSelf);
659+
}
660+
661+
/** This method is used to make an {@code CallSite.Around} of the {@code String.replace} method */
662+
@Override
663+
@SuppressFBWarnings("ES_COMPARING_PARAMETER_STRING_WITH_EQ")
664+
public String onStringReplace(
665+
@Nonnull String self, CharSequence oldCharSeq, CharSequence newCharSeq) {
666+
final IastContext ctx = IastContext.Provider.get();
667+
if (ctx == null) {
668+
return self.replace(oldCharSeq, newCharSeq);
669+
}
670+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
671+
final TaintedObject taintedSelf = taintedObjects.get(self);
672+
Range[] rangesSelf = new Range[0];
673+
if (taintedSelf != null) {
674+
rangesSelf = taintedSelf.getRanges();
675+
}
676+
677+
final TaintedObject taintedInput = taintedObjects.get(newCharSeq);
678+
Range[] rangesInput = null;
679+
if (taintedInput != null) {
680+
rangesInput = taintedInput.getRanges();
681+
}
682+
683+
if (rangesSelf.length == 0 && rangesInput == null) {
684+
return self.replace(oldCharSeq, newCharSeq);
685+
}
686+
687+
return StringUtils.replaceAndTaint(
688+
taintedObjects,
689+
self,
690+
Pattern.compile((String) oldCharSeq),
691+
(String) newCharSeq,
692+
rangesSelf,
693+
rangesInput,
694+
Integer.MAX_VALUE);
695+
}
696+
697+
/**
698+
* This method is used to make an {@code CallSite.Around} of the {@code String.replaceFirst} and
699+
* {@code String.replaceAll} methods
700+
*/
701+
@Override
702+
@SuppressForbidden
703+
public String onStringReplace(
704+
@Nonnull String self, String regex, String replacement, int numReplacements) {
705+
final IastContext ctx = IastContext.Provider.get();
706+
if (ctx == null) {
707+
if (numReplacements > 1) {
708+
return self.replaceAll(regex, replacement);
709+
} else {
710+
return self.replaceFirst(regex, replacement);
711+
}
712+
}
713+
714+
final TaintedObjects taintedObjects = ctx.getTaintedObjects();
715+
final TaintedObject taintedSelf = taintedObjects.get(self);
716+
Range[] rangesSelf = new Range[0];
717+
if (taintedSelf != null) {
718+
rangesSelf = taintedSelf.getRanges();
719+
}
720+
721+
final TaintedObject taintedInput = taintedObjects.get(replacement);
722+
Range[] rangesInput = null;
723+
if (taintedInput != null) {
724+
rangesInput = taintedInput.getRanges();
725+
}
726+
727+
if (rangesSelf.length == 0 && rangesInput == null) {
728+
if (numReplacements > 1) {
729+
return self.replaceAll(regex, replacement);
730+
} else {
731+
return self.replaceFirst(regex, replacement);
732+
}
733+
}
734+
735+
return StringUtils.replaceAndTaint(
736+
taintedObjects,
737+
self,
738+
Pattern.compile(regex),
739+
replacement,
740+
rangesSelf,
741+
rangesInput,
742+
numReplacements);
743+
}
744+
634745
/**
635746
* Adds the tainted ranges belonging to the current parameter added via placeholder taking care of
636747
* an optional tainted placeholder.

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

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -373,12 +373,11 @@ private static int updateRangesWithIndentation(
373373
if (range.getStart() + range.getLength() > end) {
374374
newLength -= lineOffset;
375375
}
376-
newRanges[i] = new Range(newStart, newLength, range.getSource(), range.getMarks());
376+
newRanges[i] = copyWithPosition(range, newStart, newLength);
377377
} else if (range.getStart() + range.getLength() >= start) {
378378
final Range newRange = newRanges[i];
379379
final int newLength = newRange.getLength() + indentation;
380-
newRanges[i] =
381-
new Range(newRange.getStart(), newLength, newRange.getSource(), newRange.getMarks());
380+
newRanges[i] = copyWithPosition(newRange, newRange.getStart(), newLength);
382381
}
383382

384383
if (range.getStart() + range.getLength() - 1 <= end) {
@@ -390,4 +389,41 @@ private static int updateRangesWithIndentation(
390389

391390
return rangeStart;
392391
}
392+
393+
/**
394+
* Split the range in two taking into account the new length of the characters.
395+
*
396+
* <p>In case start and end are out of the range, it will return the range without splitting but
397+
* taking into account the offset. In the case that the new length is less than or equal to 0, it
398+
* will return an empty array.
399+
*
400+
* @param start is the start of the character sequence
401+
* @param end is the end of the character sequence
402+
* @param newLength is the new length of the character sequence
403+
* @param range is the range to split
404+
* @param offset is the offset to apply to the range
405+
* @param diffLength is the difference between the new length and the old length
406+
*/
407+
public static Range[] splitRanges(
408+
int start, int end, int newLength, Range range, int offset, int diffLength) {
409+
start += offset;
410+
end += offset;
411+
int rangeStart = range.getStart() + offset;
412+
int rangeEnd = rangeStart + range.getLength() + diffLength;
413+
414+
int firstLength = start - rangeStart;
415+
int secondLength = range.getLength() - firstLength - newLength + diffLength;
416+
if (rangeStart > end || rangeEnd <= start) {
417+
if (firstLength <= 0) {
418+
return Ranges.EMPTY;
419+
}
420+
return new Range[] {copyWithPosition(range, rangeStart, firstLength)};
421+
}
422+
423+
Range[] splittedRanges = new Range[2];
424+
splittedRanges[0] = copyWithPosition(range, rangeStart, firstLength);
425+
splittedRanges[1] = copyWithPosition(range, rangeEnd - secondLength, secondLength);
426+
427+
return splittedRanges;
428+
}
393429
}

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,15 @@ public RangeBuilder(final int maxSize, final int arrayChunkSize) {
4747
}
4848

4949
public boolean add(final Range range) {
50+
return add(range, 0);
51+
}
52+
53+
public boolean add(final Range range, final int offset) {
5054
if (size >= maxSize) {
5155
return false;
5256
}
5357
if (head == null) {
54-
addNewEntry(new SingleEntry(range));
58+
addNewEntry(new SingleEntry(range.shift(offset)));
5559
} else {
5660
final ArrayEntry entry;
5761
if (tail instanceof ArrayEntry && tail.size() < arrayChunkSize) {
@@ -60,7 +64,7 @@ public boolean add(final Range range) {
6064
entry = new ArrayEntry(arrayChunkSize);
6165
addNewEntry(entry);
6266
}
63-
entry.add(range);
67+
entry.add(range.shift(offset));
6468
}
6569
size += 1;
6670
return true;
@@ -77,6 +81,9 @@ public boolean add(final Range[] ranges, final int offset) {
7781
if (ranges.length == 0) {
7882
return true;
7983
}
84+
if (ranges.length == 1) {
85+
return add(ranges[0], offset);
86+
}
8087

8188
if (tail instanceof ArrayEntry && ranges.length <= (arrayChunkSize - tail.size())) {
8289
// compact intermediate ranges

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

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
package com.datadog.iast.util;
22

3+
import com.datadog.iast.model.Range;
4+
import com.datadog.iast.taint.Ranges;
5+
import com.datadog.iast.taint.TaintedObjects;
6+
import java.util.regex.Matcher;
7+
import java.util.regex.Pattern;
38
import javax.annotation.Nonnull;
9+
import javax.annotation.Nullable;
410

511
public abstract class StringUtils {
612

@@ -48,4 +54,142 @@ public static int leadingWhitespaces(@Nonnull final String value, int start, int
4854
}
4955
return whitespaces;
5056
}
57+
58+
/**
59+
* Returns the string replaced with the regex and the values tainted (if needed)
60+
*
61+
* @param taintedObjects the ctx object to save the range of the tainted values
62+
* @param target the string to be replaced
63+
* @param pattern the pattern to be replaced
64+
* @param replacement the replacement string
65+
* @param ranges the ranges of the string to be replaced
66+
* @param rangesInput the ranges of the input string
67+
* @param numOfReplacements the number of replacements to be made
68+
*/
69+
@Nonnull
70+
public static String replaceAndTaint(
71+
@Nonnull TaintedObjects taintedObjects,
72+
@Nonnull final String target,
73+
Pattern pattern,
74+
String replacement,
75+
Range[] ranges,
76+
@Nullable Range[] rangesInput,
77+
int numOfReplacements) {
78+
if (numOfReplacements <= 0) {
79+
return target;
80+
}
81+
Matcher matcher = pattern.matcher(target);
82+
boolean result = matcher.find();
83+
if (result) {
84+
int offset = 0;
85+
RangeBuilder newRanges = new RangeBuilder();
86+
87+
int firstRange = 0;
88+
int newLength = replacement.length();
89+
90+
boolean canAddRange = true;
91+
StringBuffer sb = new StringBuffer();
92+
do {
93+
int start = matcher.start();
94+
int end = matcher.end();
95+
int diffLength = newLength - (end - start);
96+
97+
boolean rangesAdded = false;
98+
while (firstRange < ranges.length && canAddRange) {
99+
Range range = ranges[firstRange];
100+
int rangeStart = range.getStart();
101+
int rangeEnd = rangeStart + range.getLength();
102+
// If the replaced value is between one range
103+
if (rangeStart <= start && rangeEnd >= end) {
104+
Range[] splittedRanges =
105+
Ranges.splitRanges(start, end, newLength, range, offset, diffLength);
106+
107+
if (splittedRanges.length > 0 && splittedRanges[0].getLength() > 0) {
108+
canAddRange = newRanges.add(splittedRanges[0]);
109+
}
110+
111+
if (rangesInput != null) {
112+
canAddRange = newRanges.add(rangesInput, start + offset);
113+
rangesAdded = true;
114+
}
115+
116+
if (splittedRanges.length > 1 && splittedRanges[1].getLength() > 0) {
117+
canAddRange = newRanges.add(splittedRanges[1]);
118+
}
119+
120+
firstRange++;
121+
break;
122+
// If the replaced value starts in the range and not end there
123+
} else if (rangeStart <= start && rangeEnd > start) {
124+
Range[] splittedRanges =
125+
Ranges.splitRanges(start, end, newLength, range, offset, diffLength);
126+
127+
if (splittedRanges.length > 0 && splittedRanges[0].getLength() > 0) {
128+
canAddRange = newRanges.add(splittedRanges[0]);
129+
}
130+
131+
if (rangesInput != null && !rangesAdded) {
132+
canAddRange = newRanges.add(rangesInput, start + offset);
133+
rangesAdded = true;
134+
}
135+
136+
// If the replaced value ends in the range
137+
} else if (rangeEnd >= end) {
138+
Range[] splittedRanges =
139+
Ranges.splitRanges(start, end, newLength, range, offset, diffLength);
140+
141+
if (rangesInput != null && !rangesAdded) {
142+
canAddRange = newRanges.add(rangesInput, start + offset);
143+
rangesAdded = true;
144+
}
145+
146+
if (splittedRanges.length > 1 && splittedRanges[1].getLength() > 0) {
147+
canAddRange = newRanges.add(splittedRanges[1]);
148+
}
149+
150+
firstRange++;
151+
break;
152+
// Middle ranges
153+
} else if (rangeStart >= start) {
154+
firstRange++;
155+
continue;
156+
} else {
157+
canAddRange = newRanges.add(range, rangeStart + offset);
158+
}
159+
160+
firstRange++;
161+
}
162+
163+
// In case there are no ranges
164+
if (rangesInput != null && !rangesAdded && canAddRange) {
165+
canAddRange = newRanges.add(rangesInput, start + offset);
166+
}
167+
168+
matcher.appendReplacement(sb, replacement);
169+
170+
offset = diffLength;
171+
numOfReplacements--;
172+
if (numOfReplacements > 0) {
173+
result = matcher.find();
174+
}
175+
} while (result && numOfReplacements > 0);
176+
177+
// In the case there is no tainted object
178+
if (firstRange < ranges.length && canAddRange) {
179+
for (int i = firstRange; i < ranges.length && canAddRange; i++) {
180+
canAddRange = newRanges.add(ranges[i], offset);
181+
}
182+
}
183+
184+
matcher.appendTail(sb);
185+
String finalString = sb.toString();
186+
Range[] finalRanges = newRanges.toArray();
187+
if (finalRanges.length > 0) {
188+
taintedObjects.taint(finalString, finalRanges);
189+
}
190+
return finalString;
191+
}
192+
193+
return target;
194+
}
51195
}

0 commit comments

Comments
 (0)