Add muzzle reference detection for invokedynamic calls. (DataDog/dd-trace-java#1712)

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/ReferenceCreator.java

@@ -30,6 +30,7 @@ import java.util.Set;
 import net.bytebuddy.jar.asm.ClassReader;
 import net.bytebuddy.jar.asm.ClassVisitor;
 import net.bytebuddy.jar.asm.FieldVisitor;
+import net.bytebuddy.jar.asm.Handle;
 import net.bytebuddy.jar.asm.Label;
 import net.bytebuddy.jar.asm.MethodVisitor;
 import net.bytebuddy.jar.asm.Opcodes;
@@ -374,6 +375,49 @@ public class ReferenceCreator extends ClassVisitor {
       super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
     }
 
+    @Override
+    public void visitTypeInsn(int opcode, String type) {
+      Type typeObj = underlyingType(Type.getObjectType(type));
+      if (typeObj.getSort() == Type.OBJECT) {
+        addReference(
+            new Reference.Builder(typeObj.getInternalName())
+                .withSource(refSourceClassName, currentLineNumber)
+                .withFlag(computeMinimumClassAccess(refSourceType, typeObj))
+                .build());
+      }
+      super.visitTypeInsn(opcode, type);
+    }
+
+    @Override
+    public void visitInvokeDynamicInsn(
+        String name,
+        String descriptor,
+        Handle bootstrapMethodHandle,
+        Object... bootstrapMethodArguments) {
+      // This part might be unnecessary...
+      addReference(
+          new Reference.Builder(bootstrapMethodHandle.getOwner())
+              .withSource(refSourceClassName, currentLineNumber)
+              .withFlag(
+                  computeMinimumClassAccess(
+                      refSourceType, Type.getObjectType(bootstrapMethodHandle.getOwner())))
+              .build());
+      for (Object arg : bootstrapMethodArguments) {
+        if (arg instanceof Handle) {
+          Handle handle = (Handle) arg;
+          addReference(
+              new Reference.Builder(handle.getOwner())
+                  .withSource(refSourceClassName, currentLineNumber)
+                  .withFlag(
+                      computeMinimumClassAccess(
+                          refSourceType, Type.getObjectType(handle.getOwner())))
+                  .build());
+        }
+      }
+      super.visitInvokeDynamicInsn(
+          name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
+    }
+
     @Override
     public void visitLdcInsn(Object value) {
       if (value instanceof Type) {

testing-common/src/test/groovy/muzzle/ReferenceCreatorTest.groovy

@@ -22,6 +22,7 @@ import static muzzle.TestClasses.MethodBodyAdvice
 import io.opentelemetry.auto.test.AgentTestRunner
 import io.opentelemetry.javaagent.tooling.muzzle.Reference
 import io.opentelemetry.javaagent.tooling.muzzle.ReferenceCreator
+import spock.lang.Ignore
 
 class ReferenceCreatorTest extends AgentTestRunner {
   def "method body creates references"() {
@@ -81,6 +82,25 @@ class ReferenceCreatorTest extends AgentTestRunner {
     references.get('muzzle.TestClasses$MethodBodyAdvice$A') != null
   }
 
+  def "instanceof creates references"() {
+    setup:
+    Map<String, Reference> references = ReferenceCreator.createReferencesFrom(TestClasses.InstanceofAdvice.getName(), this.getClass().getClassLoader())
+
+    expect:
+    references.get('muzzle.TestClasses$MethodBodyAdvice$A') != null
+  }
+
+  // TODO: remove ignore when we drop java 7 support.
+  @Ignore
+  def "invokedynamic creates references"() {
+    setup:
+    Map<String, Reference> references = ReferenceCreator.createReferencesFrom(TestClasses.InDyAdvice.getName(), this.getClass().getClassLoader())
+
+    expect:
+    references.get('muzzle.TestClasses$MethodBodyAdvice$SomeImplementation') != null
+    references.get('muzzle.TestClasses$MethodBodyAdvice$B') != null
+  }
+
   private static Reference.Method findMethod(Set<Reference.Method> methods, String methodName, String methodDesc) {
     for (Reference.Method method : methods) {
       if (method == new Reference.Method(methodName, methodDesc)) {

testing-common/src/test/java/muzzle/TestClasses.java

@@ -36,7 +36,7 @@ public class TestClasses {
     public static class A {
       public B b = new B();
       protected Object protectedField = null;
-      private Object privateField = null;
+      private final Object privateField = null;
       public static B staticB = new B();
     }
 
@@ -90,4 +90,19 @@ public class TestClasses {
       MethodBodyAdvice.A.class.getName();
     }
   }
+
+  public static class InstanceofAdvice {
+    public static boolean instanceofMethod(Object a) {
+      return a instanceof MethodBodyAdvice.A;
+    }
+  }
+
+  // TODO Can't test this until java 7 is dropped.
+  public static class InDyAdvice {
+    //    public static MethodBodyAdvice.SomeInterface indyMethod(
+    //        final MethodBodyAdvice.SomeImplementation a) {
+    //      Runnable aStaticMethod = MethodBodyAdvice.B::aStaticMethod;
+    //      return a::someMethod;
+    //    }
+  }
 }