Fix issue with call sites on constructors without DUP bytecode (#9698)

Author: manuel-alvarez-alvarez

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java

@@ -23,6 +23,7 @@
 import net.bytebuddy.implementation.Implementation;
 import net.bytebuddy.jar.asm.ClassVisitor;
 import net.bytebuddy.jar.asm.Handle;
+import net.bytebuddy.jar.asm.Label;
 import net.bytebuddy.jar.asm.MethodVisitor;
 import net.bytebuddy.jar.asm.Opcodes;
 import net.bytebuddy.jar.asm.Type;
@@ -128,23 +129,43 @@ public MethodVisitor visitMethod(
   private static class CallSiteMethodVisitor extends MethodVisitor
       implements CallSiteAdvice.MethodHandler {
     protected final Advices advices;
+    protected int lastOpcode;
+    protected boolean newFollowedByDup;
 
     private CallSiteMethodVisitor(
         @Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) {
       super(ASM_API, delegated);
       this.advices = advices;
     }
 
+    @Override
+    public void visitTypeInsn(int opcode, String type) {
+      lastOpcode = opcode;
+      if (opcode == Opcodes.NEW) {
+        newFollowedByDup = false;
+      }
+      super.visitTypeInsn(opcode, type);
+    }
+
+    @Override
+    public void visitInsn(final int opcode) {
+      if (opcode == Opcodes.DUP) {
+        newFollowedByDup = lastOpcode == Opcodes.NEW;
+      }
+      lastOpcode = opcode;
+      super.visitInsn(opcode);
+    }
+
     @Override
     public void visitMethodInsn(
         final int opcode,
         final String owner,
         final String name,
         final String descriptor,
         final boolean isInterface) {
-
+      lastOpcode = opcode;
       CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor);
-      if (advice instanceof InvokeAdvice) {
+      if (applyInvokeAdvice(advice, name, descriptor)) {
         invokeAdvice((InvokeAdvice) advice, opcode, owner, name, descriptor, isInterface);
       } else {
         mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
@@ -167,6 +188,7 @@ public void visitInvokeDynamicInsn(
         final String descriptor,
         final Handle bootstrapMethodHandle,
         final Object... bootstrapMethodArguments) {
+      lastOpcode = Opcodes.INVOKEDYNAMIC;
       CallSiteAdvice advice = advices.findAdvice(bootstrapMethodHandle);
       if (advice instanceof InvokeDynamicAdvice) {
         invokeDynamicAdvice(
@@ -303,6 +325,75 @@ private Type[] methodParamTypesWithThis(String owner, String methodDescriptor) {
           methodParameterTypesWithThis.length - 1);
       return methodParameterTypesWithThis;
     }
+
+    protected boolean applyInvokeAdvice(
+        final CallSiteAdvice advice, final String methodName, final String methodDescriptor) {
+      if (!(advice instanceof InvokeAdvice)) {
+        return false;
+      }
+      if (!"<init>".equals(methodName)) {
+        return true;
+      }
+      // TODO: do not ignore ctors where there is no DUP after NEW
+      return newFollowedByDup;
+    }
+
+    // Keep track of the latest opcode to match a DUP with its previous NEW
+    @Override
+    public void visitIntInsn(final int opcode, final int operand) {
+      lastOpcode = opcode;
+      super.visitIntInsn(opcode, operand);
+    }
+
+    @Override
+    public void visitVarInsn(final int opcode, final int var) {
+      lastOpcode = opcode;
+      super.visitVarInsn(opcode, var);
+    }
+
+    @Override
+    public void visitFieldInsn(
+        final int opcode, final String owner, final String name, final String descriptor) {
+      lastOpcode = opcode;
+      super.visitFieldInsn(opcode, owner, name, descriptor);
+    }
+
+    @Override
+    public void visitJumpInsn(final int opcode, final Label label) {
+      lastOpcode = opcode;
+      super.visitJumpInsn(opcode, label);
+    }
+
+    @Override
+    public void visitLdcInsn(final Object value) {
+      lastOpcode = Opcodes.LDC;
+      super.visitLdcInsn(value);
+    }
+
+    @Override
+    public void visitIincInsn(final int var, final int increment) {
+      lastOpcode = Opcodes.IINC;
+      super.visitIincInsn(var, increment);
+    }
+
+    @Override
+    public void visitTableSwitchInsn(
+        final int min, final int max, final Label dflt, final Label... labels) {
+      lastOpcode = Opcodes.TABLESWITCH;
+      super.visitTableSwitchInsn(min, max, dflt, labels);
+    }
+
+    @Override
+    public void visitLookupSwitchInsn(final Label dflt, final int[] keys, final Label[] labels) {
+      lastOpcode = Opcodes.LOOKUPSWITCH;
+      super.visitLookupSwitchInsn(dflt, keys, labels);
+    }
+
+    @Override
+    public void visitMultiANewArrayInsn(final String descriptor, final int numDimensions) {
+      lastOpcode = Opcodes.MULTIANEWARRAY;
+      super.visitMultiANewArrayInsn(descriptor, numDimensions);
+    }
   }
 
   private static class CallSiteCtorMethodVisitor extends CallSiteMethodVisitor {
@@ -390,5 +481,14 @@ protected void invokeAdvice(
         super.invokeAdvice(advice, opcode, owner, name, descriptor, isInterface);
       }
     }
+
+    @Override
+    protected boolean applyInvokeAdvice(
+        final CallSiteAdvice advice, final String methodName, final String descriptor) {
+      if (isSuperCall) {
+        return advice instanceof InvokeAdvice && "<init>".equals(methodName);
+      }
+      return super.applyInvokeAdvice(advice, methodName, descriptor);
+    }
   }
 }

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy

@@ -193,6 +193,14 @@ class BaseCallSiteTest extends DDSpecification {
   final CallSiteTransformer transformer,
   final ClassLoader loader = Thread.currentThread().contextClassLoader) {
     final classContent = loader.getResourceAsStream("${source.getInternalName()}.class").bytes
+    return transformType(source, classContent, target, transformer, loader)
+  }
+
+  protected static byte[] transformType(final Type source,
+  final byte[] classContent,
+  final Type target,
+  final CallSiteTransformer transformer,
+  final ClassLoader loader = Thread.currentThread().contextClassLoader) {
     final classFileTransformer = new AgentBuilder.Default()
     .type(named(source.className))
     .transform(new AgentBuilder.Transformer() {

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy

@@ -1,9 +1,22 @@
 package datadog.trace.agent.tooling.csi
 
 import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteTransformer
+import net.bytebuddy.ByteBuddy
 import net.bytebuddy.asm.AsmVisitorWrapper
+import net.bytebuddy.description.method.MethodDescription
+import net.bytebuddy.description.modifier.Ownership
+import net.bytebuddy.description.modifier.Visibility
 import net.bytebuddy.description.type.TypeDescription
 import net.bytebuddy.dynamic.DynamicType
+import net.bytebuddy.dynamic.loading.ClassLoadingStrategy
+import net.bytebuddy.implementation.Implementation
+import net.bytebuddy.implementation.bytecode.ByteCodeAppender
+import net.bytebuddy.implementation.bytecode.StackManipulation
+import net.bytebuddy.implementation.bytecode.TypeCreation
+import net.bytebuddy.implementation.bytecode.member.MethodInvocation
+import net.bytebuddy.implementation.bytecode.member.MethodReturn
+import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess
+import net.bytebuddy.jar.asm.MethodVisitor
 import net.bytebuddy.jar.asm.Type
 
 import java.util.concurrent.atomic.AtomicInteger
@@ -103,6 +116,56 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest {
     new StringBuilder("test") | "Inconsistent stackmap frames"
   }
 
+  void 'test new invocation without dup'() {
+    setup:
+    def (clazz, bytes) = buildConsumerWithNewAndNoDup()
+    final source = Type.getType(clazz)
+    final target = renameType(source, 'Test')
+    final pointcut = stringReaderPointcut()
+    final advice = Mock(InvokeAdvice)
+    final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(AFTER, advice, pointcut)]))
+
+    when:
+    final transformedClass = transformType(source, bytes, target, callSiteTransformer)
+    final transformed = loadClass(target, transformedClass)
+    final instance = transformed.newInstance()
+    instance.accept('test')
+
+    then:
+    0 * advice._
+  }
+
+
+  private static Tuple2<Class<?>, byte[]> buildConsumerWithNewAndNoDup() {
+    final newType = new ByteBuddy()
+      .subclass(Cloneable)
+      .name('foo.CtorIssue')
+      .modifiers(Visibility.PUBLIC)
+      .defineMethod("accept", void, Visibility.PUBLIC, Ownership.MEMBER)
+      .withParameters(String)
+      .intercept(
+      new Implementation.Simple(new ByteCodeAppender() {
+        @Override
+        ByteCodeAppender.Size apply(MethodVisitor mv,
+          Implementation.Context ctx,
+          MethodDescription method) {
+          StackManipulation compound = new StackManipulation.Compound(
+            TypeCreation.of(new TypeDescription.ForLoadedType(StringReader)),
+            // ignore DUP opcode
+            MethodVariableAccess.REFERENCE.loadFrom(1),
+            MethodInvocation.invoke(new MethodDescription.ForLoadedConstructor(StringReader.getConstructor(String))),
+            MethodReturn.VOID
+            )
+          final size = compound.apply(mv, ctx)
+          return new ByteCodeAppender.Size(size.getMaximalSize(), method.getStackSize())
+        }
+      })
+      )
+      .make()
+      .load(Thread.currentThread().contextClassLoader, ClassLoadingStrategy.Default.INJECTION)
+    return Tuple.tuple(newType.loaded, newType.bytes)
+  }
+
   static class StringCallSites implements CallSites, TestCallSites {
 
     @Override