Make new tests + change to after

DataDog · Dec 20, 2024 · 407adcb · 407adcb
1 parent 4d3d99c
commit 407adcb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/StringModuleImpl.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/StringModuleImpl.java
@@ -833,7 +833,7 @@ public void onStringValueOf(Object param, @Nonnull String result) {
 
   @Override
   public void onStringBuilderSetLength(@Nonnull CharSequence self, int length) {
-    if (self.length() <= length) {
+    if (self.length() != length) {
       return;
     }
     final IastContext ctx = IastContext.Provider.get();

diff --git a/...ava-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy b/...ava-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy
@@ -1436,26 +1436,32 @@ class StringModuleTest extends IastModuleImplTestBase {
     taintFormat(result, taintedObject.getRanges()) == "==>my_input<=="
   }
 
-  void 'onStringBuilderSetLength empty or string not changed after setLength (#self, #length)'() {
+  void 'onStringBuilderSetLength is empty or different lengths (#self, #length)'() {
     given:
-    self?.setLength(length)
+    self?.setLength(self.length())
 
     when:
     module.onStringBuilderSetLength(self, length)
 
     then:
+    mockCalls * tracer.activeSpan() >> null
     0 * _
 
     where:
-    self              | length
-    sb()              | 0
-    sb("not_changed") | 10
+    self      | length | mockCalls
+    sb("123") | 2      | 0
+    sb()      | 0      | 1
   }
 
   void 'onStringBuilderSetLength (#input, #length)'() {
     final taintedObjects = ctx.getTaintedObjects()
     def self = addFromTaintFormat(taintedObjects, input)
-    final result = self.toString().substring(0, length)
+    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)