Fix issue with call sites in super calls to constructor (#7991)

Fix issue with call sites in super calls to constructor
DataDog · Nov 24, 2024 · 4925a50 · 4925a50
1 parent c8030bd
commit 4925a50
Show file tree

Hide file tree

Showing 13 changed files with 293 additions and 22 deletions.
diff --git a/...strumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java b/...strumentation-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/...mentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy b/...mentation-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/...-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java b/...-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 "<init>".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<String> 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 && "<init>".equals(name)) {
+          if (owner.equals(newInvocations.peekLast())) {
+            newInvocations.removeLast();
+            isSuperCall = false;
+          } else {
+            // no new before call to <init>
+            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/.../agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java b/.../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/...ava-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java b/...ava-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/...ent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy b/...ent/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, "<init>", 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,