Define helper classes in loadClass (#5528)

* Define helper classes in loadClass similarly to regular classes

* fix test

* spotless

* address review comments

Author: laurit

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/ClassLoaderInstrumentation.java renamed to instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/BootDelegationInstrumentation.java

@@ -30,17 +30,17 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import org.slf4j.LoggerFactory;
 
-/*
- * Some class loaders do not delegate to their parent, so classes in those class loaders
- * will not be able to see classes in the bootstrap class loader.
+/**
+ * Some class loaders do not delegate to their parent, so classes in those class loaders will not be
+ * able to see classes in the bootstrap class loader.
  *
- * In particular, instrumentation on classes in those class loaders will not be able to see
- * the shaded OpenTelemetry API classes in the bootstrap class loader.
+ * <p>In particular, instrumentation on classes in those class loaders will not be able to see the
+ * shaded OpenTelemetry API classes in the bootstrap class loader.
  *
- * This instrumentation forces all class loaders to delegate to the bootstrap class loader
- * for the classes that we have put in the bootstrap class loader.
+ * <p>This instrumentation forces all class loaders to delegate to the bootstrap class loader for
+ * the classes that we have put in the bootstrap class loader.
  */
-public class ClassLoaderInstrumentation implements TypeInstrumentation {
+public class BootDelegationInstrumentation implements TypeInstrumentation {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
@@ -67,7 +67,7 @@ public void transform(TypeTransformer transformer) {
                             .and(takesArgument(1, boolean.class))))
             .and(isPublic().or(isProtected()))
             .and(not(isStatic())),
-        ClassLoaderInstrumentation.class.getName() + "$LoadClassAdvice");
+        BootDelegationInstrumentation.class.getName() + "$LoadClassAdvice");
   }
 
   public static class Holder {

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/ClassLoaderInstrumentationModule.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.instrumentation.internal.classloader;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -30,17 +29,11 @@ public boolean isHelperClass(String className) {
     return className.equals("io.opentelemetry.javaagent.tooling.Constants");
   }
 
-  @Override
-  public List<String> getAdditionalHelperClassNames() {
-    return singletonList(
-        "io.opentelemetry.javaagent.instrumentation.internal.classloader.DefineClassUtil");
-  }
-
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
-        new ClassLoaderInstrumentation(),
-        new ResourceInjectionInstrumentation(),
-        new DefineClassInstrumentation());
+        new BootDelegationInstrumentation(),
+        new LoadInjectedClassInstrumentation(),
+        new ResourceInjectionInstrumentation());
   }
 }

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/DefineClassInstrumentation.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.internal.classloader;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isProtected;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * This instrumentation inserts loading of our injected helper classes at the start of {@code
+ * ClassLoader.loadClass} method.
+ */
+public class LoadInjectedClassInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return extendsClass(named("java.lang.ClassLoader"));
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("loadClass"))
+            .and(
+                takesArguments(1)
+                    .and(takesArgument(0, String.class))
+                    .or(
+                        takesArguments(2)
+                            .and(takesArgument(0, String.class))
+                            .and(takesArgument(1, boolean.class))))
+            .and(isPublic().or(isProtected()))
+            .and(not(isStatic())),
+        LoadInjectedClassInstrumentation.class.getName() + "$LoadClassAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class LoadClassAdvice {
+
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static Class<?> onEnter(
+        @Advice.This ClassLoader classLoader, @Advice.Argument(0) String name) {
+      Class<?> helperClass = InjectedClassHelper.loadHelperClass(classLoader, name);
+      if (helperClass != null) {
+        return helperClass;
+      }
+
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Return(readOnly = false) Class<?> result, @Advice.Enter Class<?> loadedClass) {
+      if (loadedClass != null) {
+        result = loadedClass;
+      }
+    }
+  }
+}

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/DefineClassUtil.java

@@ -6,7 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.internal.lambda;
 
 import io.opentelemetry.javaagent.bootstrap.ClassFileTransformerHolder;
-import io.opentelemetry.javaagent.bootstrap.InjectedHelperClassDetector;
+import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper;
 import java.lang.instrument.ClassFileTransformer;
 
 /** Helper class for transforming lambda class bytes. */
@@ -31,7 +31,7 @@ private static boolean isJava9() {
   @SuppressWarnings("unused")
   public static byte[] transform(byte[] classBytes, String slashClassName, Class<?> targetClass) {
     // Skip transforming lambdas of injected helper classes.
-    if (InjectedHelperClassDetector.isHelperClass(targetClass)) {
+    if (InjectedClassHelper.isHelperClass(targetClass)) {
       return classBytes;
     }