Verify that there are no VirtualField.find calls in indy ready advice (open-telemetry#13689)

Author: laurit

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

@@ -37,6 +37,9 @@ class AdviceTransformer {
   private static final Type OBJECT_TYPE = Type.getType(Object.class);
   private static final Type OBJECT_ARRAY_TYPE = Type.getType(Object[].class);
 
+  static final Type ADVICE_ON_METHOD_ENTER = Type.getType(Advice.OnMethodEnter.class);
+  static final Type ADVICE_ON_METHOD_EXIT = Type.getType(Advice.OnMethodExit.class);
+
   static byte[] transform(byte[] bytes) {
     ClassReader cr = new ClassReader(bytes);
     ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
@@ -236,7 +239,6 @@ private static List<AdviceLocal> getLocals(MethodNode source) {
     return result;
   }
 
-  static final Type ADVICE_ON_METHOD_ENTER = Type.getType(Advice.OnMethodEnter.class);
   private static final Type ADVICE_ASSIGN_RETURNED_TO_RETURNED =
       Type.getType(Advice.AssignReturned.ToReturned.class);
   private static final Type ADVICE_ASSIGN_RETURNED_TO_ARGUMENTS =
@@ -266,8 +268,6 @@ private static boolean isEnterAdvice(MethodNode source) {
     return hasAnnotation(source, ADVICE_ON_METHOD_ENTER);
   }
 
-  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);
   }

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

@@ -11,6 +11,7 @@
 import io.opentelemetry.javaagent.tooling.bytebuddy.ExceptionHandlers;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.util.function.BiFunction;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -65,7 +66,32 @@ public void applyAdviceToMethod(
 
   private ClassFileLocator getAdviceLocator(ClassLoader classLoader) {
     ClassFileLocator classFileLocator = ClassFileLocator.ForClassLoader.of(classLoader);
-    return transformAdvice ? new AdviceLocator(classFileLocator) : classFileLocator;
+    return new AdviceLocator(
+        classFileLocator,
+        (slashClassName, bytes) -> {
+          if (transformAdvice) {
+            // rewrite inline advice to a from accepted by indy instrumentation
+            return transformAdvice(slashClassName, bytes);
+          } else {
+            // verify that advice does not call VirtualField.find
+            // NOTE: this check is here to help converting the advice for the indy instrumentation,
+            // it can be removed once the conversion is completed
+            VirtualFieldChecker.check(bytes);
+            return bytes;
+          }
+        });
+  }
+
+  private static byte[] transformAdvice(String slashClassName, byte[] bytes) {
+    byte[] result = AdviceTransformer.transform(bytes);
+    if (result != null) {
+      dump(slashClassName, result);
+      InstrumentationModuleClassLoader.bytecodeOverride.put(
+          slashClassName.replace('/', '.'), result);
+    } else {
+      result = bytes;
+    }
+    return result;
   }
 
   @Override
@@ -79,15 +105,17 @@ public AgentBuilder.Identified.Extendable getAgentBuilder() {
 
   private static class AdviceLocator implements ClassFileLocator {
     private final ClassFileLocator delegate;
+    private final BiFunction<String, byte[], byte[]> transform;
 
-    AdviceLocator(ClassFileLocator delegate) {
+    AdviceLocator(ClassFileLocator delegate, BiFunction<String, byte[], byte[]> transform) {
       this.delegate = delegate;
+      this.transform = transform;
     }
 
     @Override
     public Resolution locate(String name) throws IOException {
       Resolution resolution = delegate.locate(name);
-      return new AdviceTransformingResolution(name, resolution);
+      return new AdviceTransformingResolution(name, resolution, transform);
     }
 
     @Override
@@ -99,10 +127,13 @@ public void close() throws IOException {
   private static class AdviceTransformingResolution implements Resolution {
     private final String name;
     private final Resolution delegate;
+    private final BiFunction<String, byte[], byte[]> transform;
 
-    AdviceTransformingResolution(String name, Resolution delegate) {
+    AdviceTransformingResolution(
+        String name, Resolution delegate, BiFunction<String, byte[], byte[]> transform) {
       this.name = name;
       this.delegate = delegate;
+      this.transform = transform;
     }
 
     @Override
@@ -113,14 +144,7 @@ public boolean isResolved() {
     @Override
     public byte[] resolve() {
       byte[] bytes = delegate.resolve();
-      byte[] result = AdviceTransformer.transform(bytes);
-      if (result != null) {
-        dump(name, result);
-        InstrumentationModuleClassLoader.bytecodeOverride.put(name.replace('/', '.'), result);
-      } else {
-        result = bytes;
-      }
-      return result;
+      return transform.apply(name, bytes);
     }
   }
 

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+import static io.opentelemetry.javaagent.tooling.instrumentation.indy.AdviceTransformer.ADVICE_ON_METHOD_ENTER;
+import static io.opentelemetry.javaagent.tooling.instrumentation.indy.AdviceTransformer.ADVICE_ON_METHOD_EXIT;
+import static io.opentelemetry.javaagent.tooling.instrumentation.indy.AdviceTransformer.hasAnnotation;
+
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.AsmApi;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.MethodNode;
+
+/**
+ * Check that advice does not call {@link VirtualField#find(Class, Class)}. For inline advice {@link
+ * VirtualField#find(Class, Class)} calls in advice are rewritten for efficiency. In non-inline
+ * advice we don't do such rewriting, we expect users to keep the result of {@link
+ * VirtualField#find(Class, Class)} in a static field.
+ */
+class VirtualFieldChecker {
+
+  private static final Type VIRTUAL_FIELD_TYPE = Type.getType(VirtualField.class);
+
+  static void check(byte[] bytes) {
+    ClassReader cr = new ClassReader(bytes);
+    ClassNode classNode = new ClassNode();
+    cr.accept(classNode, ClassReader.SKIP_FRAMES | ClassReader.SKIP_DEBUG);
+
+    String dotClassName = Type.getObjectType(classNode.name).getClassName();
+    classNode.methods.forEach(m -> checkMethod(m, dotClassName));
+  }
+
+  private static void checkMethod(MethodNode methodNode, String dotClassName) {
+    if (!hasAnnotation(methodNode, ADVICE_ON_METHOD_ENTER)
+        && !hasAnnotation(methodNode, ADVICE_ON_METHOD_EXIT)) {
+      return;
+    }
+
+    methodNode.accept(
+        new MethodVisitor(AsmApi.VERSION, null) {
+          @Override
+          public void visitMethodInsn(
+              int opcode, String owner, String name, String descriptor, boolean isInterface) {
+            if (opcode == Opcodes.INVOKESTATIC
+                && VIRTUAL_FIELD_TYPE.getInternalName().equals(owner)
+                && "find".equals(name)) {
+              throw new IllegalStateException(
+                  "Found usage of VirtualField.find in advice "
+                      + dotClassName
+                      + "."
+                      + methodNode.name);
+            }
+          }
+        });
+  }
+
+  private VirtualFieldChecker() {}
+}