Increase IAST propagation to StringBuilder subSequence (#8026)

DataDog · Nov 29, 2024 · 9f01834 · 9f01834
1 parent b1bfb13
commit 9f01834
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 34 deletions.
diff --git a/...ava-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java b/...ava-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java
@@ -139,4 +139,21 @@ public static String afterSubstring(
     }
     return result;
   }
+
+  @CallSite.After("java.lang.CharSequence java.lang.StringBuilder.subSequence(int, int)")
+  public static CharSequence afterSubSequence(
+      @CallSite.This final CharSequence self,
+      @CallSite.Argument final int beginIndex,
+      @CallSite.Argument final int endIndex,
+      @CallSite.Return final CharSequence result) {
+    final StringModule module = InstrumentationBridge.STRING;
+    if (module != null) {
+      try {
+        module.onStringSubSequence(self, beginIndex, endIndex, result);
+      } catch (final Throwable e) {
+        module.onUnexpectedException("afterSubSequence threw", e);
+      }
+    }
+    return result;
+  }
 }
diff --git a/.../src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy b/.../src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy
@@ -174,76 +174,60 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
     ex.stackTrace.find { it.className == StringBuilderCallSite.name } == null
   }
 
-  def 'test string builder substring call site'() {
+  def 'test string #type substring call site'() {
     setup:
     final iastModule = Mock(StringModule)
     InstrumentationBridge.registerIastModule(iastModule)
 
     when:
-    final result = TestStringBuilderSuite.substring(param, beginIndex)
+    final result = suite.substring(param, beginIndex)
 
     then:
     result == expected
     1 * iastModule.onStringSubSequence(param, beginIndex, param.length(), expected)
     0 * _
 
     where:
-    param         | beginIndex | expected
-    sb('012345')  | 1          | '12345'
+    type      | suite                        | param         | beginIndex | expected
+    "builder" | new TestStringBuilderSuite() | sb('012345')  | 1          | '12345'
+    "buffer"  | new TestStringBufferSuite()  | sbf('012345') | 1          | '12345'
   }
 
-  def 'test string buffer substring call site'() {
+  def 'test string #type substring with endIndex call site'() {
     setup:
     final iastModule = Mock(StringModule)
     InstrumentationBridge.registerIastModule(iastModule)
 
     when:
-    final result = TestStringBufferSuite.substring(param, beginIndex)
-
-    then:
-    result == expected
-    1 * iastModule.onStringSubSequence(param, beginIndex, param.length(), expected)
-    0 * _
-
-    where:
-    param         | beginIndex | expected
-    sbf('012345') | 1          | '12345'
-  }
-
-  def 'test string builder substring with endIndex call site'() {
-    setup:
-    final iastModule = Mock(StringModule)
-    InstrumentationBridge.registerIastModule(iastModule)
-
-    when:
-    final result = TestStringBuilderSuite.substring(param, beginIndex, endIndex)
+    final result = suite.substring(param, beginIndex, endIndex)
 
     then:
     result == expected
     1 * iastModule.onStringSubSequence(param, beginIndex, endIndex, expected)
     0 * _
 
     where:
-    param         | beginIndex | endIndex | expected
-    sb('012345')  | 1          | 5        | '1234'
+    type      | suite                        | param         | beginIndex | endIndex | expected
+    "builder" | new TestStringBuilderSuite() | sb('012345')  | 1          | 5        | '1234'
+    "buffer"  | new TestStringBufferSuite()  | sbf('012345') | 1          | 5        | '1234'
   }
 
-  def 'test string buffer substring with endIndex call site'() {
+  def 'test string #type subSequence with endIndex call site'() {
     setup:
     final iastModule = Mock(StringModule)
     InstrumentationBridge.registerIastModule(iastModule)
 
     when:
-    final result = TestStringBufferSuite.substring(param, beginIndex, endIndex)
+    final result = suite.subSequence(param, beginIndex, endIndex)
 
     then:
     result == expected
     1 * iastModule.onStringSubSequence(param, beginIndex, endIndex, expected)
     0 * _
 
     where:
-    param         | beginIndex | endIndex | expected
-    sbf('012345') | 1          | 5        | '1234'
+    type      | suite                        | param         | beginIndex | endIndex | expected
+    "builder" | new TestStringBuilderSuite() | sb('012345')  | 1          | 5        | '1234'
   }
 
   private static class BrokenToString {

diff --git a/...agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java b/...agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java
@@ -12,5 +12,11 @@ public interface TestAbstractStringBuilderSuite<E> {
 
   void append(final E target, final Object param);
 
+  String substring(final E self, final int beginIndex, final int endIndex);
+
+  String substring(final E self, final int beginIndex);
+
+  CharSequence subSequence(final E self, final int beginIndex, final int endIndex);
+
   String toString(final E target);
 }
diff --git a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBufferSuite.java b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBufferSuite.java
@@ -52,17 +52,28 @@ public String toString(final StringBuffer buffer) {
     return result;
   }
 
-  public static String substring(StringBuffer self, int beginIndex, int endIndex) {
+  @Override
+  public String substring(final StringBuffer self, final int beginIndex, final int endIndex) {
     LOGGER.debug("Before string buffer substring {} from {} to {}", self, beginIndex, endIndex);
     final String result = self.substring(beginIndex, endIndex);
     LOGGER.debug("After string buffer substring {}", result);
     return result;
   }
 
-  public static String substring(StringBuffer self, int beginIndex) {
+  @Override
+  public String substring(final StringBuffer self, final int beginIndex) {
     LOGGER.debug("Before string buffer substring {} from {}", self, beginIndex);
     final String result = self.substring(beginIndex);
     LOGGER.debug("After string buffer substring {}", result);
     return result;
   }
+
+  @Override
+  public CharSequence subSequence(
+      final StringBuffer self, final int beginIndex, final int endIndex) {
+    LOGGER.debug("Before string builder subSequence {} from {} to {}", self, beginIndex, endIndex);
+    final CharSequence result = self.subSequence(beginIndex, endIndex);
+    LOGGER.debug("After string builder subSequence {}", result);
+    return result;
+  }
 }
diff --git a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBuilderSuite.java b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBuilderSuite.java
@@ -77,17 +77,28 @@ public String plus(final Object... items) {
     return result;
   }
 
-  public static String substring(StringBuilder self, int beginIndex, int endIndex) {
+  @Override
+  public String substring(final StringBuilder self, final int beginIndex, final int endIndex) {
     LOGGER.debug("Before string builder substring {} from {} to {}", self, beginIndex, endIndex);
     final String result = self.substring(beginIndex, endIndex);
     LOGGER.debug("After string builder substring {}", result);
     return result;
   }
 
-  public static String substring(StringBuilder self, int beginIndex) {
+  @Override
+  public String substring(final StringBuilder self, final int beginIndex) {
     LOGGER.debug("Before string builder substring {} from {}", self, beginIndex);
     final String result = self.substring(beginIndex);
     LOGGER.debug("After string builder substring {}", result);
     return result;
   }
+
+  @Override
+  public CharSequence subSequence(
+      final StringBuilder self, final int beginIndex, final int endIndex) {
+    LOGGER.debug("Before string builder subSequence {} from {} to {}", self, beginIndex, endIndex);
+    final CharSequence result = self.subSequence(beginIndex, endIndex);
+    LOGGER.debug("After string builder subSequence {}", result);
+    return result;
+  }
 }