From ad752ecd40573f0da91f570875cc115c07ebab32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Mon, 25 Nov 2024 11:13:32 +0100 Subject: [PATCH 1/2] Add support for StringBuilder.append with start and end and append StrinbBuilder.append with buffer --- .../iast/propagation/StringModuleImpl.java | 39 +++++++++++ .../iast/propagation/StringModuleTest.groovy | 70 +++++++++++++++++++ .../java/lang/StringBuilderCallSite.java | 21 ++++++ .../lang/StringBuilderCallSiteTest.groovy | 18 +++++ .../bar/TestAbstractStringBuilderSuite.java | 2 + .../java/foo/bar/TestStringBufferSuite.java | 7 ++ .../java/foo/bar/TestStringBuilderSuite.java | 7 ++ .../api/iast/propagation/StringModule.java | 3 + 8 files changed, 167 insertions(+) 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 b31d7ca15fb..c3c0e8b40fa 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 @@ -126,6 +126,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) { 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 6a96cdd1603..5139d63a4a4 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 @@ -126,6 +126,76 @@ 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<==' | 0 | 3 | '123==>456<==' + 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' + } + void 'onStringBuilderInit null or empty (#builder, #param)'() { given: final result = builder?.append(param) 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 80feb171f2a..e73f9bbc22f 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 @@ -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 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 a9680db938f..7a0512cb3fd 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 @@ -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) 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 9fbf35809d4..069f48115bc 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 @@ -10,6 +10,8 @@ public interface TestAbstractStringBuilderSuite { 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 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 94b7f04e732..cfe3f02417a 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 @@ -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); 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 1edd8b1500b..26157f151ec 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 @@ -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); 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 37a28f2a0a3..3a3e7b7c49e 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 @@ -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( From 4eb61ee82492a51ebfad3eb078003f13e08c1a36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Thu, 28 Nov 2024 10:15:35 +0100 Subject: [PATCH 2/2] Add more test cases for the stringbuilder append --- .../com/datadog/iast/propagation/StringModuleTest.groovy | 4 ++++ 1 file changed, 4 insertions(+) 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 5139d63a4a4..18cf3fb6873 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 @@ -190,10 +190,14 @@ class StringModuleTest extends IastModuleImplTestBase { 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)'() {