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..a5f8d0f883f 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 @@ -403,7 +403,7 @@ public void onStringTrim(@Nonnull final String self, @Nullable final String resu } @Override - public void onStringConstructor(@Nonnull String self, @Nonnull String result) { + public void onStringConstructor(@Nonnull CharSequence self, @Nonnull String result) { if (!canBeTainted(self)) { return; } 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..f880a31bb2f 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 @@ -810,6 +810,47 @@ class StringModuleTest extends IastModuleImplTestBase { "A==>\u00cc\u00cc\u00cc<==B" | "a==>ììì<==b" | "en" | 5 | 5 | [[1, 3]] } + void 'onStringConstructor (#input)'() { + given: + final taintedObjects = ctx.getTaintedObjects() + final self = addFromTaintFormat(taintedObjects, input) + final result = new String(self) + + when: + module.onStringConstructor(self, result) + def taintedObject = taintedObjects.get(result) + + then: + 1 * tracer.activeSpan() >> span + taintFormat(result, taintedObject.getRanges()) == expected + + where: + input | expected + "==>123<==" | "==>123<==" + sb("==>123<==") | "==>123<==" + sbf("==>123<==") | "==>123<==" + } + + void 'onStringConstructor empty (#input)'() { + given: + final taintedObjects = ctx.getTaintedObjects() + final self = addFromTaintFormat(taintedObjects, input) + final result = new String(self) + + when: + module.onStringConstructor(self, result) + + then: + null == taintedObjects.get(result) + result == expected + + where: + input | expected + "" | "" + sb("") | "" + sbf("") | "" + } + void 'test trim and make sure IastRequestContext is called'() { given: final taintedObjects = ctx.getTaintedObjects() @@ -1269,6 +1310,10 @@ class StringModuleTest extends IastModuleImplTestBase { return new StringBuilder(string) } + private static StringBuilder sbf() { + return sbf('') + } + private static StringBuffer sbf(final String string) { return new StringBuffer(string) } diff --git a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java index 9d845256077..75669f20c9a 100644 --- a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java +++ b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java @@ -201,13 +201,15 @@ public static String afterTrim( } @CallSite.After("void java.lang.String.(java.lang.String)") + @CallSite.After("void java.lang.String.(java.lang.StringBuffer)") + @CallSite.After("void java.lang.String.(java.lang.StringBuilder)") public static String afterStringConstructor( @CallSite.AllArguments @Nonnull final Object[] params, @CallSite.Return @Nonnull final String result) { final StringModule module = InstrumentationBridge.STRING; try { if (module != null) { - module.onStringConstructor((String) params[0], result); + module.onStringConstructor((CharSequence) params[0], result); } } catch (final Throwable e) { module.onUnexpectedException("afterStringConstructor threw", e); diff --git a/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy b/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy index 24c4f7b1b26..0bb3fe70a13 100644 --- a/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy @@ -176,12 +176,18 @@ class StringCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringSuite.stringConstructor("hello") + final result = TestStringSuite.stringConstructor(param) then: - result == 'hello' + result == expected 1 * module.onStringConstructor(_, _) 0 * _ + + where: + param | expected + 'hello' | 'hello' + new StringBuilder('hello') | 'hello' + new StringBuffer('hello') | 'hello' } void 'test string format'() { diff --git a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringSuite.java b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringSuite.java index 2e1c00f4f50..b25b6de9aae 100644 --- a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringSuite.java +++ b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringSuite.java @@ -92,6 +92,20 @@ public static String stringConstructor(String self) { return result; } + public static String stringConstructor(StringBuffer self) { + LOGGER.debug("Before string stringConstructor {}", self); + final String result = new String(self); + LOGGER.debug("After string stringConstructor {}", result); + return result; + } + + public static String stringConstructor(StringBuilder self) { + LOGGER.debug("Before string stringConstructor {}", self); + final String result = new String(self); + LOGGER.debug("After string stringConstructor {}", result); + return result; + } + public static String stringConstructor(final byte[] value) { LOGGER.debug("Before string stringConstructor {}", value); final String result = new String(value); 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..a42df53652c 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 @@ -36,7 +36,7 @@ void onStringJoin( void onStringRepeat(@Nonnull String self, int count, @Nonnull String result); - void onStringConstructor(@Nonnull String self, @Nonnull String result); + void onStringConstructor(@Nonnull CharSequence self, @Nonnull String result); void onStringFormat(@Nonnull String pattern, @Nonnull Object[] params, @Nonnull String result);