Allow any types in invokedynamic advice method signatures

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceSignatureEraser.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+import io.opentelemetry.javaagent.extension.instrumentation.internal.AsmApi;
+import java.util.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.tree.ClassNode;
+
+/**
+ * This transformer is a workaround to make the life of instrumentation developers easier.
+ *
+ * <p>When inserting the instrumentation and calling the advice enter/exit methods, we insert an
+ * invokedynamic instruction to perform this method call. This instruction has a {@link
+ * java.lang.invoke.MethodType} associated, which normally corresponds to the method type of the
+ * called advice. This can however lead to problems, when the advice method signature contains types
+ * which are not visible to the instrumented class.
+ *
+ * <p>To prevent this, we instead associate the invokedynamic instruction with a type where all
+ * class reference are replaced with references to {@link Object}. To lookup the correct advice
+ * method, we pass the original type as string argument to the invokedynamic bootstrapping method.
+ *
+ * <p>Because bytebuddy currently doesn't support this customization, we perform the type erasure on
+ * the .class via ASM before bytebuddy parses the advice. In addition, we insert an annotation to
+ * preserve the original descriptor of the method.
+ */
+// TODO: replace this workaround when buddy natively supports altering the MethodType for
+// invokedynamic advices
+class AdviceSignatureEraser {
+
+  private static final Pattern TYPE_REFERENCE_PATTERN = Pattern.compile("L[^;]+;");
+
+  public static final String ORIGNINAL_DESCRIPTOR_ANNOTATION_TYPE =
+      "L" + OriginalDescriptor.class.getName().replace('.', '/') + ";";
+
+  private AdviceSignatureEraser() {}
+
+  static byte[] transform(byte[] bytes) {
+    ClassReader cr = new ClassReader(bytes);
+    ClassWriter cw = new ClassWriter(cr, 0);
+    ClassNode classNode = new ClassNode();
+    cr.accept(classNode, 0);
+
+    Set<String> methodsToTransform = listAdviceMethods(classNode);
+    if (methodsToTransform.isEmpty()) {
+      return bytes;
+    }
+
+    ClassVisitor cv =
+        new ClassVisitor(AsmApi.VERSION, cw) {
+          @Override
+          public MethodVisitor visitMethod(
+              int access, String name, String descriptor, String signature, String[] exceptions) {
+            if (methodsToTransform.contains(name + descriptor)) {
+              String erased = eraseTypes(descriptor);
+              MethodVisitor visitor = super.visitMethod(access, name, erased, null, exceptions);
+              AnnotationVisitor av =
+                  visitor.visitAnnotation(ORIGNINAL_DESCRIPTOR_ANNOTATION_TYPE, false);
+              av.visit("value", descriptor);
+              av.visitEnd();
+              return visitor;
+            } else {
+              return super.visitMethod(access, name, descriptor, signature, exceptions);
+            }
+          }
+        };
+    classNode.accept(cv);
+    return cw.toByteArray();
+  }
+
+  private static String eraseTypes(String descriptor) {
+    Matcher matcher = TYPE_REFERENCE_PATTERN.matcher(descriptor);
+    StringBuffer result = new StringBuffer();
+    while (matcher.find()) {
+      String reference = matcher.group();
+      if (reference.startsWith("Ljava/")) {
+        // do not erase java.* references
+        matcher.appendReplacement(result, reference);
+      } else {
+        matcher.appendReplacement(result, "Ljava/lang/Object;");
+      }
+    }
+    matcher.appendTail(result);
+    return result.toString();
+  }
+
+  private static Set<String> listAdviceMethods(ClassNode classNode) {
+    return classNode.methods.stream()
+        .filter(
+            mn ->
+                AdviceTransformer.hasAnnotation(mn, AdviceTransformer.ADVICE_ON_METHOD_ENTER)
+                    || AdviceTransformer.hasAnnotation(mn, AdviceTransformer.ADVICE_ON_METHOD_EXIT))
+        .map(mn -> mn.name + mn.desc)
+        .collect(Collectors.toSet());
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceTransformer.java

@@ -225,13 +225,13 @@ private static List<AdviceLocal> getLocals(MethodNode source) {
     return result;
   }
 
-  private static final Type ADVICE_ON_METHOD_ENTER = Type.getType(Advice.OnMethodEnter.class);
+  static final Type ADVICE_ON_METHOD_ENTER = Type.getType(Advice.OnMethodEnter.class);
 
   private static boolean isEnterAdvice(MethodNode source) {
     return hasAnnotation(source, ADVICE_ON_METHOD_ENTER);
   }
 
-  private static final Type ADVICE_ON_METHOD_EXIT = Type.getType(Advice.OnMethodExit.class);
+  static final Type ADVICE_ON_METHOD_EXIT = Type.getType(Advice.OnMethodExit.class);
 
   private static boolean isExitAdvice(MethodNode source) {
     return hasAnnotation(source, ADVICE_ON_METHOD_EXIT);
@@ -250,7 +250,7 @@ private static AnnotationNode getAnnotationNode(MethodNode source, Type type) {
     return null;
   }
 
-  private static boolean hasAnnotation(MethodNode source, Type type) {
+  static boolean hasAnnotation(MethodNode source, Type type) {
     return getAnnotationNode(source, type) != null;
   }
 

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyBootstrap.java

@@ -20,6 +20,8 @@
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.annotation.AnnotationDescription;
+import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.utility.JavaConstant;
 
 /**
@@ -131,7 +133,12 @@ private static ConstantCallSite internalBootstrap(
         case BOOTSTRAP_KIND_ADVICE:
           // See the getAdviceBootstrapArguments method for the argument definitions
           return bootstrapAdvice(
-              lookup, adviceMethodName, adviceMethodType, (String) args[1], (String) args[2]);
+              lookup,
+              adviceMethodName,
+              adviceMethodType,
+              (String) args[1],
+              (String) args[2],
+              (String) args[3]);
         case BOOTSTRAP_KIND_PROXY:
           // See getProxyFactory for the argument definitions
           return bootstrapProxyMethod(
@@ -153,8 +160,9 @@ private static ConstantCallSite internalBootstrap(
   private static ConstantCallSite bootstrapAdvice(
       MethodHandles.Lookup lookup,
       String adviceMethodName,
-      MethodType adviceMethodType,
+      MethodType invokedynamicMethodType,
       String moduleClassName,
+      String adviceMethodDescriptor,
       String adviceClassName)
       throws NoSuchMethodException, IllegalAccessException, ClassNotFoundException {
     CallDepth callDepth = CallDepth.forClass(IndyBootstrap.class);
@@ -177,10 +185,14 @@ private static ConstantCallSite bootstrapAdvice(
       // Advices are not inlined. They are loaded as normal classes by the
       // InstrumentationModuleClassloader and invoked via a method call from the instrumented method
       Class<?> adviceClass = instrumentationClassloader.loadClass(adviceClassName);
+      MethodType actualAdviceMethodType =
+          MethodType.fromMethodDescriptorString(adviceMethodDescriptor, instrumentationClassloader);
+
       MethodHandle methodHandle =
           instrumentationClassloader
               .getLookup()
-              .findStatic(adviceClass, adviceMethodName, adviceMethodType);
+              .findStatic(adviceClass, adviceMethodName, actualAdviceMethodType)
+              .asType(invokedynamicMethodType);
       return new ConstantCallSite(methodHandle);
     } finally {
       callDepth.decrementAndGet();
@@ -195,9 +207,19 @@ static Advice.BootstrapArgumentResolver.Factory getAdviceBootstrapArguments(
             Arrays.asList(
                 JavaConstant.Simple.ofLoaded(BOOTSTRAP_KIND_ADVICE),
                 JavaConstant.Simple.ofLoaded(moduleName),
+                JavaConstant.Simple.ofLoaded(getOriginalSignature(adviceMethod)),
                 JavaConstant.Simple.ofLoaded(adviceMethod.getDeclaringType().getName()));
   }
 
+  private static String getOriginalSignature(MethodDescription.InDefinedShape adviceMethod) {
+    for (AnnotationDescription an : adviceMethod.getDeclaredAnnotations()) {
+      if (OriginalDescriptor.class.getName().equals(an.getAnnotationType().getName())) {
+        return (String) an.getValue("value").resolve();
+      }
+    }
+    throw new IllegalStateException("OriginalSignature annotation is not present!");
+  }
+
   private static ConstantCallSite bootstrapProxyMethod(
       MethodHandles.Lookup lookup,
       String proxyMethodName,

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyTypeTransformerImpl.java

@@ -101,9 +101,11 @@ public byte[] resolve() {
       if (result != null) {
         dump(name, result);
         InstrumentationModuleClassLoader.bytecodeOverride.put(name.replace('/', '.'), result);
-        return result;
+      } else {
+        result = bytes;
       }
-      return bytes;
+      result = AdviceSignatureEraser.transform(result);
+      return result;
     }
   }
 

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/OriginalDescriptor.java

@@ -0,0 +1,10 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+public @interface OriginalDescriptor {
+  String value();
+}

testing-common/integration-tests/src/main/java/io/opentelemetry/javaagent/IndyInstrumentationTestModule.java

@@ -84,13 +84,15 @@ public void transform(TypeTransformer transformer) {
       transformer.applyAdviceToMethod(named("exceptionPlease"), prefix + "$ThrowExceptionAdvice");
       transformer.applyAdviceToMethod(
           named("noExceptionPlease"), prefix + "$SuppressExceptionAdvice");
+      transformer.applyAdviceToMethod(
+          named("instrumentWithErasedTypes"), prefix + "$SignatureErasureAdvice");
     }
 
     @SuppressWarnings({"unused"})
     public static class AssignFieldViaReturnAdvice {
 
       @Advice.OnMethodEnter(inline = false)
-      @Advice.AssignReturned.ToFields(@ToField(value = "privateField"))
+      @Advice.AssignReturned.ToFields(@ToField(value = "privateField", typing = DYNAMIC))
       public static String onEnter(@Advice.Argument(0) String toAssign) {
         return toAssign;
       }
@@ -110,7 +112,7 @@ public static Object[] onEnter(@Advice.Argument(0) String toAssign) {
     public static class AssignArgumentViaReturnAdvice {
 
       @Advice.OnMethodEnter(inline = false)
-      @Advice.AssignReturned.ToArguments(@ToArgument(0))
+      @Advice.AssignReturned.ToArguments(@ToArgument(value = 0, typing = DYNAMIC))
       public static String onEnter(@Advice.Argument(1) String toAssign) {
         return toAssign;
       }
@@ -130,7 +132,7 @@ public static Object[] onEnter(@Advice.Argument(1) String toAssign) {
     public static class AssignReturnViaReturnAdvice {
 
       @Advice.OnMethodExit(inline = false)
-      @Advice.AssignReturned.ToReturned
+      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
       public static String onExit(@Advice.Argument(0) String toAssign) {
         return toAssign;
       }
@@ -150,7 +152,7 @@ public static Object[] onExit(@Advice.Argument(0) String toAssign) {
     public static class GetHelperClassAdvice {
 
       @Advice.OnMethodExit(inline = false)
-      @Advice.AssignReturned.ToReturned
+      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
       public static Class<?> onExit(@Advice.Argument(0) boolean localHelper) {
         if (localHelper) {
           return LocalHelper.class;
@@ -175,7 +177,7 @@ public static void onMethodEnter() {
         throw new RuntimeException("This exception should be suppressed");
       }
 
-      @Advice.AssignReturned.ToReturned
+      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
       @Advice.OnMethodExit(
           suppress = Throwable.class,
           onThrowable = Throwable.class,
@@ -184,6 +186,23 @@ public static void onMethodExit(@Advice.Thrown Throwable throwable) {
         throw new RuntimeException("This exception should be suppressed");
       }
     }
+
+    @SuppressWarnings({"unused", "ThrowSpecificExceptions"})
+    public static class SignatureErasureAdvice {
+      @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+      public static LocalHelper onMethodEnter() {
+        return new LocalHelper();
+      }
+
+      @Advice.AssignReturned.ToReturned(typing = DYNAMIC)
+      @Advice.OnMethodExit(
+          suppress = Throwable.class,
+          onThrowable = Throwable.class,
+          inline = false)
+      public static LocalHelper onMethodExit(@Advice.Enter LocalHelper enterVal) {
+        return enterVal;
+      }
+    }
   }
 
   public static class GlobalHelper {}

testing-common/integration-tests/src/test/java/indy/IndyInstrumentationTest.java

@@ -60,6 +60,10 @@ private Class<?> getHelperClass(boolean local) {
     return null;
   }
 
+  private Object instrumentWithErasedTypes() {
+    return "replace_me";
+  }
+
   @AfterEach
   public void reset() {
     privateField = null;
@@ -112,6 +116,12 @@ void testThrowExceptionIntoUserCode() {
     assertThatThrownBy(this::exceptionPlease).isInstanceOf(RuntimeException.class);
   }
 
+  @Test
+  void testAdviceSignatureReferenceInternalHelper() {
+    Object result = instrumentWithErasedTypes();
+    assertThat(result.getClass().getName()).contains("LocalHelper");
+  }
+
   @Test
   void testHelperClassLoading() {
     Class<?> localHelper = getHelperClass(true);