Allow disabling muzzle checks for specific methods

conventions/build.gradle.kts

@@ -42,8 +42,6 @@ dependencies {
   implementation("com.google.guava:guava:31.1-jre")
   implementation("gradle.plugin.com.google.protobuf:protobuf-gradle-plugin:0.8.18")
   implementation("gradle.plugin.com.github.johnrengelman:shadow:7.1.2")
-  implementation("org.ow2.asm:asm:9.4")
-  implementation("org.ow2.asm:asm-tree:9.4")
   implementation("org.apache.httpcomponents:httpclient:4.5.13")
   implementation("org.gradle:test-retry-gradle-plugin:1.4.1")
   implementation("org.owasp:dependency-check-gradle:7.3.2")

dependencyManagement/build.gradle.kts

@@ -40,6 +40,7 @@ val autoServiceVersion = "1.0.1"
 val autoValueVersion = "1.10.1"
 val errorProneVersion = "2.16"
 val byteBuddyVersion = "1.12.19"
+val asmVersion = "9.4"
 val jmhVersion = "1.36"
 val mockitoVersion = "4.9.0"
 val slf4jVersion = "2.0.5"
@@ -57,6 +58,8 @@ val CORE_DEPENDENCIES = listOf(
   "net.bytebuddy:byte-buddy-dep:${byteBuddyVersion}",
   "net.bytebuddy:byte-buddy-agent:${byteBuddyVersion}",
   "net.bytebuddy:byte-buddy-gradle-plugin:${byteBuddyVersion}",
+  "org.ow2.asm:asm:${asmVersion}",
+  "org.ow2.asm:asm-tree:${asmVersion}",
   "org.openjdk.jmh:jmh-core:${jmhVersion}",
   "org.openjdk.jmh:jmh-generator-bytecode:${jmhVersion}",
   "org.mockito:mockito-core:${mockitoVersion}",

instrumentation/spring/spring-kafka-2.7/library/build.gradle.kts

@@ -3,6 +3,8 @@ plugins {
 }
 
 dependencies {
+  compileOnly(project(":muzzle"))
+
   compileOnly("com.google.auto.value:auto-value-annotations")
   annotationProcessor("com.google.auto.value:auto-value")
 

instrumentation/spring/spring-kafka-2.7/library/src/main/java/io/opentelemetry/instrumentation/spring/kafka/v2_7/InstrumentedRecordInterceptor.java

@@ -9,9 +9,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
-import java.lang.invoke.MethodHandle;
-import java.lang.invoke.MethodHandles;
-import java.lang.invoke.MethodType;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
 import javax.annotation.Nullable;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -23,22 +21,6 @@ final class InstrumentedRecordInterceptor<K, V> implements RecordInterceptor<K,
       VirtualField.find(ConsumerRecord.class, Context.class);
   private static final VirtualField<ConsumerRecord<?, ?>, State<ConsumerRecord<?, ?>>> stateField =
       VirtualField.find(ConsumerRecord.class, State.class);
-  private static final MethodHandle interceptRecord;
-
-  static {
-    MethodHandle interceptRecordHandle;
-    try {
-      interceptRecordHandle =
-          MethodHandles.lookup()
-              .findVirtual(
-                  RecordInterceptor.class,
-                  "intercept",
-                  MethodType.methodType(ConsumerRecord.class, ConsumerRecord.class));
-    } catch (NoSuchMethodException | IllegalAccessException e) {
-      interceptRecordHandle = null;
-    }
-    interceptRecord = interceptRecordHandle;
-  }
 
   private final Instrumenter<ConsumerRecord<?, ?>, Void> processInstrumenter;
   @Nullable private final RecordInterceptor<K, V> decorated;
@@ -50,25 +32,13 @@ final class InstrumentedRecordInterceptor<K, V> implements RecordInterceptor<K,
     this.decorated = decorated;
   }
 
-  @SuppressWarnings({
-    "deprecation",
-    "unchecked"
-  }) // implementing deprecated method (removed in 3.0) for better compatibility
+  @NoMuzzle
+  @SuppressWarnings(
+      "deprecation") // implementing deprecated method (removed in 3.0) for better compatibility
   @Override
   public ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record) {
-    if (interceptRecord == null) {
-      return null;
-    }
     start(record);
-    if (decorated == null) {
-      return null;
-    }
-    try {
-      return (ConsumerRecord<K, V>) interceptRecord.invoke(decorated, record);
-    } catch (Throwable e) {
-      rethrow(e);
-      return null; // unreachable
-    }
+    return decorated == null ? record : decorated.intercept(record);
   }
 
   @Override
@@ -77,11 +47,6 @@ public ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record, Consumer<K, V
     return decorated == null ? record : decorated.intercept(record, consumer);
   }
 
-  @SuppressWarnings("unchecked")
-  private static <E extends Throwable> void rethrow(Throwable e) throws E {
-    throw (E) e;
-  }
-
   private void start(ConsumerRecord<K, V> record) {
     Context parentContext = getParentContext(record);
 

muzzle/build.gradle.kts

@@ -14,6 +14,7 @@ dependencies {
   annotationProcessor("com.google.auto.value:auto-value")
 
   api("net.bytebuddy:byte-buddy-dep")
+  implementation("org.ow2.asm:asm-tree")
 
   implementation(project(":javaagent-bootstrap"))
   implementation(project(":instrumentation-api"))

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/NoMuzzle.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.muzzle;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/** Skip muzzle checks for methods annotated with this annotation. */
+@Retention(RetentionPolicy.CLASS)
+@Target(ElementType.METHOD)
+public @interface NoMuzzle {}

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/ReferenceCollectingClassVisitor.java

@@ -28,6 +28,8 @@
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.AnnotationNode;
+import org.objectweb.asm.tree.MethodNode;
 
 /** Visit a class and collect all references made by the visited class. */
 // Additional things we could check
@@ -244,11 +246,33 @@ public MethodVisitor visitMethod(
               .build());
     }
 
+    MethodVisitor methodVisitor =
+        super.visitMethod(access, name, descriptor, signature, exceptions);
+    MethodVisitor methodNode =
+        new MethodNode(Opcodes.ASM9, access, name, descriptor, signature, exceptions) {
+          @Override
+          public void visitEnd() {
+            super.visitEnd();
+
+            boolean skip = false;
+            if (invisibleAnnotations != null) {
+              for (AnnotationNode annotationNode : invisibleAnnotations) {
+                if (Type.getDescriptor(NoMuzzle.class).equals(annotationNode.desc)) {
+                  skip = true;
+                  break;
+                }
+              }
+            }
+            MethodVisitor target =
+                skip ? methodVisitor : new AdviceReferenceMethodVisitor(methodVisitor);
+            if (target != null) {
+              accept(target);
+            }
+          }
+        };
     // Additional references we could check
     // - Classes in signature (return type, params) and visible from this package
-    return new AdviceReferenceMethodVisitor(
-        new VirtualFieldCollectingMethodVisitor(
-            super.visitMethod(access, name, descriptor, signature, exceptions)));
+    return new VirtualFieldCollectingMethodVisitor(methodNode);
   }
 
   private static VisibilityFlag computeVisibilityFlag(int access) {