From 4925a5023343e7abf918bef4738b708ef61f20c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Mon, 25 Nov 2024 00:13:11 +0100 Subject: [PATCH] Fix issue with call sites in super calls to constructor (#7991) Fix issue with call sites in super calls to constructor --- .../plugin/csi/impl/AdviceGeneratorImpl.java | 6 +- .../csi/impl/AdviceGeneratorTest.groovy | 20 ++--- .../bytebuddy/csi/CallSiteTransformer.java | 85 ++++++++++++++++++- .../tooling/bytebuddy/csi/CallSiteUtils.java | 13 ++- .../agent/tooling/csi/CallSiteAdvice.java | 10 ++- .../agent/tooling/csi/BaseCallSiteTest.groovy | 17 ++++ .../csi/CallSiteInstrumentationTest.groovy | 46 ++++++++++ .../tooling/csi/CallSiteUtilsTest.groovy | 38 +++++++++ .../agent/tooling/csi/SuperInCtorExample.java | 10 +++ .../java/io/StringReaderCallSiteTest.groovy | 28 +++++- .../java/foo/bar/TestCustomStringReader.java | 14 +++ .../controller/IastWebController.java | 17 ++++ .../AbstractIastSpringBootTest.groovy | 11 +++ 13 files changed, 293 insertions(+), 22 deletions(-) create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/SuperInCtorExample.java create mode 100644 dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomStringReader.java diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java index 9bab598ffdc..04c7a9be704 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java +++ b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java @@ -344,12 +344,10 @@ private static void writeAdviceMethodCall( final MethodCallExpr invokeStatic = new MethodCallExpr() .setScope(new NameExpr("handler")) - .setName("method") - .addArgument(opCode("INVOKESTATIC")) + .setName("advice") .addArgument(new StringLiteralExpr(method.getOwner().getInternalName())) .addArgument(new StringLiteralExpr(method.getMethodName())) - .addArgument(new StringLiteralExpr(method.getMethodType().getDescriptor())) - .addArgument(new BooleanLiteralExpr(false)); + .addArgument(new StringLiteralExpr(method.getMethodType().getDescriptor())); body.addStatement(invokeStatic); } if (requiresCast(advice)) { diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy index dcc5dc8559a..a0b3e15e2cc 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy +++ b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy @@ -47,7 +47,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { pointcut('java/security/MessageDigest', 'getInstance', '(Ljava/lang/String;)Ljava/security/MessageDigest;') statements( 'handler.dupParameters(descriptor, StackDupMode.COPY);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$BeforeAdvice", "before", "(Ljava/lang/String;)V", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$BeforeAdvice", "before", "(Ljava/lang/String;)V");', 'handler.method(opcode, owner, name, descriptor, isInterface);' ) } @@ -78,7 +78,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { advices(0) { pointcut('java/lang/String', 'replaceAll', '(Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;') statements( - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;", false);' + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");' ) } } @@ -110,7 +110,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);', 'handler.method(opcode, owner, name, descriptor, isInterface);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdvice", "after", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdvice", "after", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");', ) } } @@ -142,7 +142,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);', 'handler.method(opcode, owner, name, descriptor, isInterface);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceCtor", "after", "([Ljava/lang/Object;Ljava/net/URL;)Ljava/net/URL;", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceCtor", "after", "([Ljava/lang/Object;Ljava/net/URL;)Ljava/net/URL;");', ) } } @@ -208,7 +208,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY);', 'handler.invokeDynamic(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicAfterAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;)Ljava/lang/String;", false);' + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicAfterAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;)Ljava/lang/String;");' ) } } @@ -297,7 +297,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { 'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY);', 'handler.invokeDynamic(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);', 'handler.loadConstantArray(bootstrapMethodArguments);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicWithConstantsAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;", false);' + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicWithConstantsAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;");' ) } } @@ -393,7 +393,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'int[] parameterIndices = new int[] { 0 };', 'handler.dupParameters(descriptor, parameterIndices, owner);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;)V", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;)V");', 'handler.method(opcode, owner, name, descriptor, isInterface);', ) } @@ -402,7 +402,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'int[] parameterIndices = new int[] { 1 };', 'handler.dupParameters(descriptor, parameterIndices, null);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "([Ljava/lang/Object;)V", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "([Ljava/lang/Object;)V");', 'handler.method(opcode, owner, name, descriptor, isInterface);', ) } @@ -411,7 +411,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'int[] parameterIndices = new int[] { 0 };', 'handler.dupInvoke(owner, descriptor, parameterIndices);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;I)V", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;I)V");', 'handler.method(opcode, owner, name, descriptor, isInterface);', ) } @@ -443,7 +443,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { statements( 'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);', 'handler.method(opcode, owner, name, descriptor, isInterface);', - 'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$SuperTypeReturnAdvice", "after", "([Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;", false);', + 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$SuperTypeReturnAdvice", "after", "([Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");', 'handler.instruction(Opcodes.CHECKCAST, "java/lang/StringBuilder");' ) } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java index a56a4954e3f..991523069c1 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java @@ -1,5 +1,6 @@ package datadog.trace.agent.tooling.bytebuddy.csi; +import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY; import static net.bytebuddy.jar.asm.ClassWriter.COMPUTE_MAXS; import datadog.trace.agent.tooling.HelperInjector; @@ -9,6 +10,8 @@ import datadog.trace.agent.tooling.csi.InvokeAdvice; import datadog.trace.agent.tooling.csi.InvokeDynamicAdvice; import java.security.ProtectionDomain; +import java.util.Deque; +import java.util.LinkedList; import javax.annotation.Nonnull; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.field.FieldDescription; @@ -24,9 +27,13 @@ import net.bytebuddy.jar.asm.Type; import net.bytebuddy.pool.TypePool; import net.bytebuddy.utility.JavaModule; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class CallSiteTransformer implements Instrumenter.TransformingAdvice { + private static Logger LOGGER = LoggerFactory.getLogger(CallSiteTransformer.class); + private static final Instrumenter.TransformingAdvice NO_OP = (builder, typeDescription, classLoader, module, pd) -> builder; @@ -111,7 +118,9 @@ public MethodVisitor visitMethod( final String[] exceptions) { final MethodVisitor delegated = super.visitMethod(access, name, descriptor, signature, exceptions); - return new CallSiteMethodVisitor(advices, delegated); + return "".equals(name) + ? new CallSiteCtorMethodVisitor(advices, delegated) + : new CallSiteMethodVisitor(advices, delegated); } } @@ -132,6 +141,7 @@ public void visitMethodInsn( final String name, final String descriptor, final boolean isInterface) { + CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor); if (advice instanceof InvokeAdvice) { ((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface); @@ -197,6 +207,11 @@ public void method( mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface); } + @Override + public void advice(String owner, String name, String descriptor) { + mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false); + } + @Override public void invokeDynamic( final String name, @@ -265,4 +280,72 @@ private Type[] methodParamTypesWithThis(String owner, String methodDescriptor) { return methodParameterTypesWithThis; } } + + private static class CallSiteCtorMethodVisitor extends CallSiteMethodVisitor { + + private final Deque newInvocations = new LinkedList<>(); + private boolean isSuperCall = false; + + private CallSiteCtorMethodVisitor( + @Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) { + super(advices, delegated); + } + + @Override + public void visitEnd() { + super.visitEnd(); + if (!newInvocations.isEmpty()) { + LOGGER.debug( + SEND_TELEMETRY, + "There is an issue handling NEW bytecodes, remaining types {}", + newInvocations); + } + } + + @Override + public void visitTypeInsn(final int opcode, final String type) { + if (opcode == Opcodes.NEW) { + newInvocations.addLast(type); + } + super.visitTypeInsn(opcode, type); + } + + @Override + public void visitMethodInsn( + final int opcode, + final String owner, + final String name, + final String descriptor, + final boolean isInterface) { + try { + if (opcode == Opcodes.INVOKESPECIAL && "".equals(name)) { + if (owner.equals(newInvocations.peekLast())) { + newInvocations.removeLast(); + isSuperCall = false; + } else { + // no new before call to + isSuperCall = true; + } + } + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + } finally { + isSuperCall = false; + } + } + + @Override + public void advice(final String owner, final String name, final String descriptor) { + if (isSuperCall) { + // append this to the stack after super call + mv.visitIntInsn(Opcodes.ALOAD, 0); + } + mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false); + } + + @Override + public void dupParameters(final String methodDescriptor, final StackDupMode mode) { + super.dupParameters( + methodDescriptor, isSuperCall ? StackDupMode.PREPEND_ARRAY_SUPER_CTOR : mode); + } + } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java index d2b6a28cb74..39511a60a76 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java @@ -151,6 +151,7 @@ public static void dup(final MethodVisitor mv, final Type[] parameters, final St break; case PREPEND_ARRAY: case PREPEND_ARRAY_CTOR: + case PREPEND_ARRAY_SUPER_CTOR: case APPEND_ARRAY: dupN(mv, parameters, mode); break; @@ -280,11 +281,21 @@ private static void dupN( mv.visitInsn(POP); break; case PREPEND_ARRAY_CTOR: - // move the array before the NEW and DUP opcodes + // move the array before the uninitialized entry created by NEW and DUP + // stack start = [uninitialized, uninitialized, arg_0, ..., arg_n] + // stack end = [array, uninitialized, uninitialized, arg_0, ..., arg_n] mv.visitInsn(DUP_X2); loadArray(mv, arraySize, parameters); mv.visitInsn(POP); break; + case PREPEND_ARRAY_SUPER_CTOR: + // move the array before the uninitialized entry + // stack start = [uninitialized, arg_0, ..., arg_n] + // stack end = [array, uninitialized, arg_0, ..., arg_n] + mv.visitInsn(DUP_X1); + loadArray(mv, arraySize, parameters); + mv.visitInsn(POP); + break; case APPEND_ARRAY: loadArray(mv, arraySize, parameters); break; diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java index 28c41fb820d..0e43f3602e8 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java @@ -27,6 +27,9 @@ interface MethodHandler { /** Performs a method invocation (static, special, virtual, interface...) */ void method(int opcode, String owner, String name, String descriptor, boolean isInterface); + /** Performs an advice invocation (always static) */ + void advice(String owner, String name, String descriptor); + /** Performs a dynamic method invocation */ void invokeDynamic( String name, @@ -62,11 +65,10 @@ enum StackDupMode { COPY, /** Copies the parameters in an array and prepends it */ PREPEND_ARRAY, - /** - * Copies the parameters in an array, prepends it and swaps the array with the uninitialized - * instance in a ctor - */ + /** Copies the parameters in an array and adds it between NEW and DUP opcodes */ PREPEND_ARRAY_CTOR, + /** Copies the parameters in an array and adds it before the uninitialized instance in a ctor */ + PREPEND_ARRAY_SUPER_CTOR, /** Copies the parameters in an array and appends it */ APPEND_ARRAY } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy index 8c69fc0bd21..2ee1e555758 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy @@ -15,6 +15,8 @@ import net.bytebuddy.jar.asm.Type import net.bytebuddy.matcher.ElementMatcher import net.bytebuddy.utility.JavaModule import net.bytebuddy.utility.nullability.MaybeNull + +import java.lang.reflect.Constructor import java.security.MessageDigest @@ -81,6 +83,10 @@ class BaseCallSiteTest extends DDSpecification { return buildPointcut(String.getDeclaredMethod('concat', String)) } + protected static Pointcut stringReaderPointcut() { + return buildPointcut(StringReader.getDeclaredConstructor(String)) + } + protected static Pointcut messageDigestGetInstancePointcut() { return buildPointcut(MessageDigest.getDeclaredMethod('getInstance', String)) } @@ -100,6 +106,10 @@ class BaseCallSiteTest extends DDSpecification { return buildPointcut(Type.getType(executable.getDeclaringClass()).internalName, executable.name, Type.getType(executable).descriptor) } + protected static Pointcut buildPointcut(final Constructor executable) { + return buildPointcut(Type.getType(executable.getDeclaringClass()).internalName, "", Type.getType(executable).descriptor) + } + protected static Pointcut buildPointcut(final String type, final String method, final String descriptor) { return new Pointcut(type: type, method: method, descriptor: descriptor) } @@ -157,6 +167,13 @@ class BaseCallSiteTest extends DDSpecification { return clazz.getConstructor().newInstance() } + protected static Class loadClass(final Type type, + final byte[] data, + final ClassLoader loader = Thread.currentThread().contextClassLoader) { + final classLoader = new ByteArrayClassLoader(loader, [(type.className): data]) + return classLoader.loadClass(type.className) + } + protected static byte[] transformType(final Type source, final Type target, final CallSiteTransformer transformer, diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy index 0823d9b419e..b56b94779ad 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy @@ -1,8 +1,14 @@ package datadog.trace.agent.tooling.csi +import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteTransformer import net.bytebuddy.asm.AsmVisitorWrapper import net.bytebuddy.description.type.TypeDescription import net.bytebuddy.dynamic.DynamicType +import net.bytebuddy.jar.asm.Type + +import java.util.concurrent.atomic.AtomicInteger + +import static datadog.trace.agent.tooling.csi.CallSiteAdvice.StackDupMode.PREPEND_ARRAY_CTOR class CallSiteInstrumentationTest extends BaseCallSiteTest { @@ -60,6 +66,36 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest { 0 * builder.visit(_ as AsmVisitorWrapper) >> builder } + void 'test call site transformer with super call in ctor'() { + setup: + SuperInCtorExampleAdvice.CALLS.set(0) + final source = Type.getType(SuperInCtorExample) + final target = renameType(source, 'Test') + final pointcut = stringReaderPointcut() + final InvokeAdvice advice = new InvokeAdvice() { + @Override + void apply(CallSiteAdvice.MethodHandler handler, int opcode, String owner, String name, String descriptor, boolean isInterface) { + handler.dupParameters(descriptor, PREPEND_ARRAY_CTOR) + handler.method(opcode, owner, name, descriptor, isInterface) + handler.advice( + Type.getType(SuperInCtorExampleAdvice).internalName, + 'onInvoke', + Type.getMethodType(Type.getType(StringReader), Type.getType(Object[]), Type.getType(StringReader)).getDescriptor(), + ) + } + } + final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(advice, pointcut)])) + + when: + final transformedClass = transformType(source, target, callSiteTransformer) + final transformed = loadClass(target, transformedClass) + final reader = transformed.newInstance("test") + + then: + reader != null + SuperInCtorExampleAdvice.CALLS.get() > 0 + } + static class StringCallSites implements CallSites, TestCallSites { @Override @@ -82,4 +118,14 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest { handler.method(opcode, owner, name, descriptor, isInterface) } } + + static class SuperInCtorExampleAdvice { + + private static final AtomicInteger CALLS = new AtomicInteger(0) + + static StringReader onInvoke(Object[] args, StringReader result) { + CALLS.incrementAndGet() + return result + } + } } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy index 6a345b4b868..e3703981457 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy @@ -240,6 +240,44 @@ class CallSiteUtilsTest extends DDSpecification { ] | items.subList(2, items.size()) } + void 'test stack dup with array before super ctor of #items'() { + setup: + final stack = buildStack(items) + final visitor = mockMethodVisitor(stack) + + when: + CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_SUPER_CTOR) + + then: 'the first element of the stack should be an array with the parameters' + final arrayFromStack = stack.remove(0) + arrayFromStack.type.descriptor == '[Ljava/lang/Object;' + final array = (arrayFromStack.value as Object[]).toList() as List + [array, expected].transpose().each { arrayItem, expectedItem -> + assert arrayItem.value == expectedItem.value // some of the items might be boxed so be careful + } + + then: 'the rest of the stack should contain the original values' + final result = fromStack(stack) + result == items + + where: + items | expected + [forObject('uninitialized'), forInt(1)] | items.subList(1, items.size()) + [forObject('uninitialized'), forInt(1), forInt(2)] | items.subList(1, items.size()) + [forObject('uninitialized'), forLong(1L)] | items.subList(1, items.size()) + [forObject('uninitialized'), forLong(1L), forLong(2L)] | items.subList(1, items.size()) + [forObject('uninitialized'), forInt(1), forLong(2L)] | items.subList(1, items.size()) + [forObject('uninitialized'), forInt(1), forInt(2), forLong(3L)] | items.subList(1, items.size()) + [forObject('uninitialized'), forInt(1), forInt(2), forInt(3), forLong(4L)] | items.subList(1, items.size()) + [ + forObject('uninitialized'), + forObject('PI = '), + forDouble(3.14D), + forChar((char) '?'), + forBoolean(true) + ] | items.subList(1, items.size()) + } + private MethodVisitor mockMethodVisitor(final List stack) { return Mock(MethodVisitor) { visitInsn(Opcodes.DUP) >> { handleDUP(stack) } diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/SuperInCtorExample.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/SuperInCtorExample.java new file mode 100644 index 00000000000..b3d5356ecf4 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/SuperInCtorExample.java @@ -0,0 +1,10 @@ +package datadog.trace.agent.tooling.csi; + +import java.io.StringReader; + +public class SuperInCtorExample extends StringReader { + + public SuperInCtorExample(String s) { + super(s + new StringReader(s + "Test" + new StringBuilder("another test"))); + } +} diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy index 914801845a0..e106adfb98f 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy @@ -2,11 +2,12 @@ package datadog.trace.instrumentation.java.io import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule +import foo.bar.TestCustomStringReader import foo.bar.TestStringReaderSuite -class StringReaderCallSiteTest extends BaseIoCallSiteTest{ +class StringReaderCallSiteTest extends BaseIoCallSiteTest { - void 'test StringReader.'(){ + void 'test StringReader.'() { given: PropagationModule iastModule = Mock(PropagationModule) InstrumentationBridge.registerIastModule(iastModule) @@ -18,4 +19,27 @@ class StringReaderCallSiteTest extends BaseIoCallSiteTest{ then: 1 * iastModule.taintObjectIfTainted(_ as StringReader, input) } + + void 'test super call to StringReader.'() { + given: + PropagationModule iastModule = Mock(PropagationModule) + InstrumentationBridge.registerIastModule(iastModule) + final input = 'Test input' + + when: + new TestCustomStringReader(input) + + then: + // new StringReader + 3 * iastModule.taintObjectIfTainted( + { it -> !(it instanceof TestCustomStringReader) }, + { String it -> + it.startsWith("New") } + ) + + // super(...) + 1 * iastModule.taintObjectIfTainted( + { it instanceof TestCustomStringReader }, + { String it -> it.startsWith("Super") }) + } } diff --git a/dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomStringReader.java b/dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomStringReader.java new file mode 100644 index 00000000000..34082f66ad8 --- /dev/null +++ b/dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomStringReader.java @@ -0,0 +1,14 @@ +package foo.bar; + +import java.io.StringReader; + +public class TestCustomStringReader extends StringReader { + + public TestCustomStringReader(String s) { + super( + "Super " + + s + + (new StringReader( + "New_1" + new StringReader("New_2" + new StringReader("New_3" + s))))); + } +} diff --git a/dd-smoke-tests/iast-util/src/main/java/datadog/smoketest/springboot/controller/IastWebController.java b/dd-smoke-tests/iast-util/src/main/java/datadog/smoketest/springboot/controller/IastWebController.java index f92723bb2b1..3cfd1ca77f0 100644 --- a/dd-smoke-tests/iast-util/src/main/java/datadog/smoketest/springboot/controller/IastWebController.java +++ b/dd-smoke-tests/iast-util/src/main/java/datadog/smoketest/springboot/controller/IastWebController.java @@ -2,6 +2,7 @@ import com.google.gson.Gson; import com.google.gson.JsonParser; +import datadog.communication.util.IOUtils; import datadog.smoketest.springboot.TestBean; import ddtest.client.sources.Hasher; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; @@ -413,6 +414,11 @@ public String untrustedDeserializationSnakeYaml(@RequestParam("yaml") String par return "OK"; } + @GetMapping("/test_custom_string_reader") + public String testCustomStringReader(@RequestParam("param") String param) throws IOException { + return String.join("", IOUtils.readLines(new CustomStringReader(param))); + } + private void withProcess(final Operation op) { Process process = null; try { @@ -429,4 +435,15 @@ private void withProcess(final Operation op) { private interface Operation { E run() throws Throwable; } + + public static class CustomStringReader extends StringReader { + + public CustomStringReader(String s) { + super( + "Super " + + s + + (new StringReader( + "New_1" + new StringReader("New_2" + new StringReader("New_3" + s))))); + } + } } diff --git a/dd-smoke-tests/iast-util/src/testFixtures/groovy/datadog/smoketest/AbstractIastSpringBootTest.groovy b/dd-smoke-tests/iast-util/src/testFixtures/groovy/datadog/smoketest/AbstractIastSpringBootTest.groovy index e5bd0a58e84..96b5dcaabe3 100644 --- a/dd-smoke-tests/iast-util/src/testFixtures/groovy/datadog/smoketest/AbstractIastSpringBootTest.groovy +++ b/dd-smoke-tests/iast-util/src/testFixtures/groovy/datadog/smoketest/AbstractIastSpringBootTest.groovy @@ -1184,5 +1184,16 @@ abstract class AbstractIastSpringBootTest extends AbstractIastServerSmokeTest { hasVulnerability { vul -> vul.type == 'UNTRUSTED_DESERIALIZATION' } } + void 'test custom string reader'() { + setup: + final url = "http://localhost:${httpPort}/test_custom_string_reader?param=Test" + final request = new Request.Builder().url(url).get().build() + + when: + final response = client.newCall(request).execute() + + then: + response.body().string().contains("Test") + } }