Increase IAST propagation to StringBuilder append (#8010)

Author: Mariovido

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

@@ -128,6 +128,45 @@ public void onStringBuilderAppend(
     }
   }
 
+  @Override
+  public void onStringBuilderAppend(
+      @Nonnull CharSequence builder, @Nullable CharSequence param, int start, int end) {
+    if (!canBeTainted(builder) || !canBeTainted(param)) {
+      return;
+    }
+    if (start < 0 || end > param.length() || start > end) {
+      return;
+    }
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    final TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    final TaintedObject paramTainted = taintedObjects.get(param);
+    if (paramTainted == null) {
+      return;
+    }
+    // We get the ranges for the new substring that will be appended to the builder
+    final Range[] paramRanges = paramTainted.getRanges();
+    final Range[] newParamRanges = Ranges.forSubstring(start, end - start, paramRanges);
+    if (newParamRanges == null) {
+      return;
+    }
+    final TaintedObject builderTainted = taintedObjects.get(builder);
+    // If the builder is not tainted we must shift the ranges of the parameter
+    // Else we shift the ranges of the parameter and merge with the ranges of the builder
+    final int shift = builder.length() - (end - start);
+    if (builderTainted == null) {
+      final Range[] ranges = new Range[newParamRanges.length];
+      Ranges.copyShift(newParamRanges, ranges, 0, shift);
+      taintedObjects.taint(builder, ranges);
+    } else {
+      final Range[] builderRanges = builderTainted.getRanges();
+      final Range[] ranges = mergeRanges(shift, builderRanges, newParamRanges);
+      builderTainted.setRanges(ranges);
+    }
+  }
+
   @Override
   public void onStringBuilderToString(
       @Nonnull final CharSequence builder, @Nonnull final String result) {

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy

@@ -131,6 +131,80 @@ class StringModuleTest extends IastModuleImplTestBase {
     sb('1==>234<==5==>678<==9') | 'a==>bcd<==e==>fgh<==i' | 1         | '1==>234<==5==>678<==9a==>bcd<==e==>fgh<==i'
   }
 
+  void 'onStringBuilderAppend null or empty (#builder, #param, #start, #end)'() {
+    given:
+    final result = builder?.append(param, start, end)
+
+    when:
+    module.onStringBuilderAppend(result, param, start, end)
+
+    then:
+    0 * _
+
+    where:
+    builder | param | start | end
+    sb('')  | null  | 0     | 0
+    sb('')  | ''    | 0     | 0
+  }
+
+  void 'onStringBuilderAppend without span (#builder, #param, #start, #end)'() {
+    given:
+    final result = builder?.append(param, start, end)
+
+    when:
+    module.onStringBuilderAppend(result, param)
+
+    then:
+    mockCalls * tracer.activeSpan() >> null
+    0 * _
+
+    where:
+    builder | param | start | end | mockCalls
+    sb('1') | null  | 0     | 0   | 0
+    sb('3') | '4'   | 0     | 0   | 1
+  }
+
+  void 'onStringBuilderAppend (#builder, #param, #start, #end)'() {
+    given:
+    final taintedObjects = ctx.getTaintedObjects()
+    def builderTainted = addFromTaintFormat(taintedObjects, builder)
+    objectHolder.add(builderTainted)
+    def paramTainted = addFromTaintFormat(taintedObjects, param)
+    objectHolder.add(paramTainted)
+    builderTainted?.append(paramTainted, start, end)
+
+    and:
+    final result = getStringFromTaintFormat(expected)
+    objectHolder.add(result)
+    final shouldBeTainted = fromTaintFormat(expected) != null
+
+    when:
+    module.onStringBuilderAppend(builderTainted, paramTainted, start, end)
+    def taintedObject = taintedObjects.get(builderTainted)
+
+    then:
+    if (shouldBeTainted) {
+      assert taintedObject != null
+      assert taintedObject.get() as String == result
+      assert taintFormat(taintedObject.get() as String, taintedObject.getRanges()) == expected
+    } else {
+      assert taintedObject == null
+    }
+
+    where:
+    builder                     | param                   | start | end | expected
+    sb('123')                   | '456'                   | 0     | 3   | '123456'
+    sb('==>123<==')             | '456'                   | 0     | 3   | '==>123<==456'
+    sb('==>123<==')             | '456'                   | 1     | 3   | '==>123<==56'
+    sb('123')                   | '==>456<=='             | 0     | 3   | '123==>456<=='
+    sb('123')                   | '==>456<=='             | 1     | 2   | '123==>5<=='
+    sb('==>123<==')             | '==>456<=='             | 0     | 3   | '==>123<====>456<=='
+    sb('1==>234<==5==>678<==9') | 'a==>bcd<==e'           | 0     | 5   | '1==>234<==5==>678<==9a==>bcd<==e'
+    sb('1==>234<==5==>678<==9') | 'a==>bcd<==e==>fgh<==i' | 0     | 9   | '1==>234<==5==>678<==9a==>bcd<==e==>fgh<==i'
+    sb('1==>234<==5==>678<==9') | 'a==>bcd<==e==>fgh<==i' | 5     | 9   | '1==>234<==5==>678<==9==>fgh<==i'
+    sb('1==>234<==5==>678<==9') | 'a==>bcd<==e==>fgh<==i' | 5     | 8   | '1==>234<==5==>678<==9==>fgh<=='
+  }
+
   void 'onStringBuilderInit null or empty (#builder, #param)'() {
     given:
     final result = builder?.append(param)

dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java

@@ -38,6 +38,7 @@ public static CharSequence afterInit(
 
   @CallSite.After("java.lang.StringBuilder java.lang.StringBuilder.append(java.lang.String)")
   @CallSite.After("java.lang.StringBuilder java.lang.StringBuilder.append(java.lang.CharSequence)")
+  @CallSite.After("java.lang.StringBuilder java.lang.StringBuilder.append(java.lang.StringBuffer)")
   @CallSite.After("java.lang.StringBuffer java.lang.StringBuffer.append(java.lang.String)")
   @CallSite.After("java.lang.StringBuffer java.lang.StringBuffer.append(java.lang.CharSequence)")
   @Nonnull
@@ -56,6 +57,26 @@ public static CharSequence afterAppend(
     return result;
   }
 
+  @CallSite.After(
+      "java.lang.StringBuilder java.lang.StringBuilder.append(java.lang.CharSequence, int, int)")
+  @Nonnull
+  public static CharSequence afterAppendWithSubstring(
+      @CallSite.This @Nonnull final CharSequence self,
+      @CallSite.Argument(0) @Nullable final CharSequence param,
+      @CallSite.Argument(1) final int start,
+      @CallSite.Argument(2) final int end,
+      @CallSite.Return @Nonnull final CharSequence result) {
+    final StringModule module = InstrumentationBridge.STRING;
+    if (module != null) {
+      try {
+        module.onStringBuilderAppend(self, param, start, end);
+      } catch (final Throwable e) {
+        module.onUnexpectedException("afterAppendWithSubstring threw", e);
+      }
+    }
+    return result;
+  }
+
   @CallSite.Around("java.lang.StringBuilder java.lang.StringBuilder.append(java.lang.Object)")
   @CallSite.Around("java.lang.StringBuffer java.lang.StringBuffer.append(java.lang.Object)")
   @Nonnull

dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy

@@ -90,6 +90,24 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
     new TestStringBufferSuite()  | new StringBuffer('Hello ')
   }
 
+  void 'test string builder append call site with start and end'() {
+    setup:
+    final iastModule = Mock(StringModule)
+    InstrumentationBridge.registerIastModule(iastModule)
+
+    when:
+    suite.append(target, param, start, end)
+
+    then:
+    target.toString() == expected
+    1 * iastModule.onStringBuilderAppend(target, param, start, end)
+    0 * _
+
+    where:
+    suite                        | target                      | param    | start | end | expected
+    new TestStringBuilderSuite() | new StringBuilder('Hello ') | 'World!' | 0     | 5   | 'Hello World'
+  }
+
   void 'test string builder toString call site'() {
     setup:
     final iastModule = Mock(StringModule)

dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java

@@ -10,6 +10,8 @@ public interface TestAbstractStringBuilderSuite<E> {
 
   void append(final E target, final CharSequence param);
 
+  void append(final E target, final CharSequence param, int start, int end);
+
   void append(final E target, final Object param);
 
   String substring(final E self, final int beginIndex, final int endIndex);

dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBufferSuite.java

@@ -37,6 +37,13 @@ public void append(final StringBuffer buffer, final CharSequence param) {
     LOGGER.debug("After string buffer append {}", result);
   }
 
+  @Override
+  public void append(final StringBuffer builder, final CharSequence param, int start, int end) {
+    LOGGER.debug("Before string buffer append {} with start {} and end {}", param, start, end);
+    final StringBuffer result = builder.append(param, start, end);
+    LOGGER.debug("After string buffer append {}", result);
+  }
+
   @Override
   public void append(final StringBuffer buffer, final Object param) {
     LOGGER.debug("Before string buffer append {}", param);

dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBuilderSuite.java

@@ -38,6 +38,13 @@ public void append(final StringBuilder builder, final CharSequence param) {
     LOGGER.debug("After string builder append {}", result);
   }
 
+  @Override
+  public void append(final StringBuilder builder, final CharSequence param, int start, int end) {
+    LOGGER.debug("Before string builder append {} with start {} and end {}", param, start, end);
+    final StringBuilder result = builder.append(param, start, end);
+    LOGGER.debug("After string builder append {}", result);
+  }
+
   @Override
   public void append(final StringBuilder builder, final Object param) {
     LOGGER.debug("Before string builder append {}", param);

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

@@ -13,6 +13,9 @@ public interface StringModule extends IastModule {
 
   void onStringBuilderAppend(@Nonnull CharSequence builder, @Nullable CharSequence param);
 
+  void onStringBuilderAppend(
+      @Nonnull CharSequence builder, @Nullable CharSequence param, int start, int end);
+
   void onStringBuilderToString(@Nonnull CharSequence builder, @Nonnull String result);
 
   void onStringConcatFactory(