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 4ef50cb6b1e..7d8124b9118 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 @@ -128,6 +128,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 99b17559b4b..06ed21ce670 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 @@ -131,6 +131,80 @@ 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' | 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)'() { 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 23bcdd860a5..d449f1a8425 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(