From c79490bdea40a7528a582f11f7d653dbe597731f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Thu, 19 Dec 2024 15:47:22 +0100 Subject: [PATCH 1/3] Add support for setLength in the StringBuilder --- .../iast/propagation/StringModuleImpl.java | 24 +++++++++++++ .../iast/propagation/StringModuleTest.groovy | 36 +++++++++++++++++++ .../java/lang/StringBuilderCallSite.java | 14 ++++++++ .../lang/StringBuilderCallSiteTest.groovy | 18 ++++++++++ .../bar/TestAbstractStringBuilderSuite.java | 2 ++ .../java/foo/bar/TestStringBufferSuite.java | 11 ++++-- .../java/foo/bar/TestStringBuilderSuite.java | 7 ++++ .../api/iast/propagation/StringModule.java | 2 ++ 8 files changed, 112 insertions(+), 2 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 index 7d8124b9118..6758f74c6f1 100644 --- 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 @@ -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. diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy index 06ed21ce670..1d84976db2e 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy @@ -1436,6 +1436,42 @@ class StringModuleTest extends IastModuleImplTestBase { taintFormat(result, taintedObject.getRanges()) == "==>my_input<==" } + void 'onStringBuilderSetLength empty or string not changed after setLength (#self, #length)'() { + given: + self?.setLength(length) + + when: + module.onStringBuilderSetLength(self, length) + + then: + 0 * _ + + where: + self | length + sb() | 0 + sb("not_changed") | 10 + } + + void 'onStringBuilderSetLength (#input, #length)'() { + final taintedObjects = ctx.getTaintedObjects() + def self = addFromTaintFormat(taintedObjects, input) + final result = self.toString().substring(0, length) + + 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) } diff --git a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java index 730564445d2..39708f027fd 100644 --- a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java +++ b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java @@ -181,4 +181,18 @@ public static CharSequence afterSubSequence( } return result; } + + // TODO - CHANGE TO AFTER + @CallSite.Before("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); + } + } + } } diff --git a/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy b/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy index da45bf012e2..dd8b81ca764 100644 --- a/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy +++ b/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() { diff --git a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java index f65e99e97fa..660a0d7bfad 100644 --- a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java +++ b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestAbstractStringBuilderSuite.java @@ -20,5 +20,7 @@ public interface TestAbstractStringBuilderSuite { CharSequence subSequence(final E self, final int beginIndex, final int endIndex); + void setLength(final E self, final int length); + 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 index 736e1387bd1..71ef6707a42 100644 --- 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 @@ -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); + } } 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 index 188c7b0c09c..36eb5dcb900 100644 --- 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 @@ -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); + } } diff --git a/internal-api/src/main/java/datadog/trace/api/iast/propagation/StringModule.java b/internal-api/src/main/java/datadog/trace/api/iast/propagation/StringModule.java index d449f1a8425..ebdd0fd7e70 100644 --- a/internal-api/src/main/java/datadog/trace/api/iast/propagation/StringModule.java +++ b/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, From 4ecb27087cd14df692aef039553e3fb051164340 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 20 Dec 2024 12:42:27 +0100 Subject: [PATCH 2/3] Change from before to after method --- .../trace/instrumentation/java/lang/StringBuilderCallSite.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java index 39708f027fd..146f62f9627 100644 --- a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java +++ b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java @@ -182,8 +182,7 @@ public static CharSequence afterSubSequence( return result; } - // TODO - CHANGE TO AFTER - @CallSite.Before("void java.lang.StringBuilder.setLength(int)") + @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; From 407adcb3cb01bb92ca592cda89a337dce681afd9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 20 Dec 2024 13:55:51 +0100 Subject: [PATCH 3/3] Make new tests + change to after --- .../iast/propagation/StringModuleImpl.java | 2 +- .../iast/propagation/StringModuleTest.groovy | 18 ++++++++++++------ 2 files changed, 13 insertions(+), 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 index 6758f74c6f1..2dcbc9b3c79 100644 --- 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/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy index 1d84976db2e..e631cdc3270 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy +++ b/dd-java-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)