Increase IAST propagation to StringBuilder setLength (#8119)

Author: Mariovido

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

@@ -831,6 +831,30 @@ public void onStringValueOf(Object param, @Nonnull String result) {
     }
   }
 
+  @Override
+  public void onStringBuilderSetLength(@Nonnull CharSequence self, int length) {
+    if (self.length() != length) {
+      return;
+    }
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    final TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    final TaintedObject selfTainted = taintedObjects.get(self);
+    if (selfTainted == null) {
+      return;
+    }
+    final Range[] rangesSelf = selfTainted.getRanges();
+    if (rangesSelf.length == 0) {
+      return;
+    }
+    Range[] newRanges = Ranges.forSubstring(0, length, rangesSelf);
+    if (newRanges != null && newRanges.length > 0) {
+      selfTainted.setRanges(newRanges);
+    }
+  }
+
   /**
    * 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/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy

@@ -1436,6 +1436,48 @@ class StringModuleTest extends IastModuleImplTestBase {
     taintFormat(result, taintedObject.getRanges()) == "==>my_input<=="
   }
 
+  void 'onStringBuilderSetLength is empty or different lengths (#self, #length)'() {
+    given:
+    self?.setLength(self.length())
+
+    when:
+    module.onStringBuilderSetLength(self, length)
+
+    then:
+    mockCalls * tracer.activeSpan() >> null
+    0 * _
+
+    where:
+    self      | length | mockCalls
+    sb("123") | 2      | 0
+    sb()      | 0      | 1
+  }
+
+  void 'onStringBuilderSetLength (#input, #length)'() {
+    final taintedObjects = ctx.getTaintedObjects()
+    def self = addFromTaintFormat(taintedObjects, input)
+    if (self instanceof StringBuilder) {
+      ((StringBuilder) self).setLength(length)
+    } else if (self instanceof StringBuffer) {
+      ((StringBuffer) self).setLength(length)
+    }
+    final result = self.toString()
+
+    when:
+    module.onStringBuilderSetLength(self, length)
+    def taintedObject = taintedObjects.get(self)
+
+    then:
+    1 * tracer.activeSpan() >> span
+    taintFormat(result, taintedObject.getRanges()) == expected
+
+    where:
+    input                         | length | expected
+    sb("==>0123<==")              | 3      | "==>012<=="
+    sb("0123==>456<==78")         | 5      | "0123==>4<=="
+    sb("01==>234<==5==>678<==90") | 8      | "01==>234<==5==>67<=="
+  }
+
   private static Date date(final String pattern, final String value) {
     return new SimpleDateFormat(pattern).parse(value)
   }

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

@@ -181,4 +181,17 @@ public static CharSequence afterSubSequence(
     }
     return result;
   }
+
+  @CallSite.After("void java.lang.StringBuilder.setLength(int)")
+  public static void afterSetLength(
+      @CallSite.This final CharSequence self, @CallSite.Argument final int length) {
+    final StringModule module = InstrumentationBridge.STRING;
+    if (module != null) {
+      try {
+        module.onStringBuilderSetLength(self, length);
+      } catch (final Throwable e) {
+        module.onUnexpectedException("afterSetLength threw", e);
+      }
+    }
+  }
 }

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

@@ -250,6 +250,24 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
     "buffer"  | new TestStringBufferSuite()  | sbf('012345') | 1          | 5        | '1234'
   }
 
+  def 'test string #type setLength with length: #length call site'() {
+    setup:
+    final iastModule = Mock(StringModule)
+    InstrumentationBridge.registerIastModule(iastModule)
+
+    when:
+    suite.setLength(param, length)
+
+    then:
+    param.toString() == expected
+    1 * iastModule.onStringBuilderSetLength(param, length)
+    0 * _
+
+    where:
+    type      | suite                        | param         | length | expected
+    "builder" | new TestStringBuilderSuite() | sb('012345')  | 5      | '01234'
+  }
+
   private static class BrokenToString {
     @Override
     String toString() {

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

@@ -20,5 +20,7 @@ public interface TestAbstractStringBuilderSuite<E> {
 
   CharSequence subSequence(final E self, final int beginIndex, final int endIndex);
 
+  void setLength(final E self, final int length);
+
   String toString(final E target);
 }

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

@@ -78,9 +78,16 @@ public String substring(final StringBuffer self, final int beginIndex) {
   @Override
   public CharSequence subSequence(
       final StringBuffer self, final int beginIndex, final int endIndex) {
-    LOGGER.debug("Before string builder subSequence {} from {} to {}", self, beginIndex, endIndex);
+    LOGGER.debug("Before string buffer subSequence {} from {} to {}", self, beginIndex, endIndex);
     final CharSequence result = self.subSequence(beginIndex, endIndex);
-    LOGGER.debug("After string builder subSequence {}", result);
+    LOGGER.debug("After string buffer subSequence {}", result);
     return result;
   }
+
+  @Override
+  public void setLength(final StringBuffer self, final int length) {
+    LOGGER.debug("Before string buffer setLength {} with length {}", self, length);
+    self.setLength(length);
+    LOGGER.debug("After string buffer setLength {}", self);
+  }
 }

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

@@ -108,4 +108,11 @@ public CharSequence subSequence(
     LOGGER.debug("After string builder subSequence {}", result);
     return result;
   }
+
+  @Override
+  public void setLength(final StringBuilder self, final int length) {
+    LOGGER.debug("Before string builder setLength {} with length {}", self, length);
+    self.setLength(length);
+    LOGGER.debug("After string builder setLength {}", self);
+  }
 }

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

@@ -18,6 +18,8 @@ void onStringBuilderAppend(
 
   void onStringBuilderToString(@Nonnull CharSequence builder, @Nonnull String result);
 
+  void onStringBuilderSetLength(@Nonnull CharSequence self, int length);
+
   void onStringConcatFactory(
       @Nullable String result,
       @Nullable String[] args,