Reduce unsafe usages (#14855)

Author: laurit

instrumentation/executors/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/executors/ThreadPoolExecutorTest.java

@@ -10,7 +10,7 @@
 
 import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldInstalledMarker;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldInstalledMarker;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.LinkedBlockingQueue;

instrumentation/internal/internal-class-loader/compile-stub/build.gradle.kts

@@ -0,0 +1,3 @@
+plugins {
+  id("otel.java-conventions")
+}

instrumentation/internal/internal-class-loader/compile-stub/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/stub/ClassLoader.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.internal.classloader.stub;
+
+import java.security.ProtectionDomain;
+
+/**
+ * A placeholder for java.lang.ClassLoader to allow compilation of advice classes that invoke
+ * protected methods of ClassLoader (like defineClass and findLoadedClass). During the build we'll
+ * use shadow plugin to replace reference to this class with the real java.lang.ClassLoader.
+ */
+@SuppressWarnings("JavaLangClash")
+public abstract class ClassLoader {
+  public abstract Class<?> findLoadedClass(String name);
+
+  public abstract Class<?> defineClass(
+      String name, byte[] b, int off, int len, ProtectionDomain protectionDomain);
+}

instrumentation/internal/internal-class-loader/javaagent-integration-tests/build.gradle.kts

@@ -6,5 +6,5 @@ dependencies {
   compileOnly("org.apache.commons:commons-lang3:3.12.0")
   testImplementation("org.apache.commons:commons-lang3:3.12.0")
 
-  testInstrumentation(project(":instrumentation:internal:internal-class-loader:javaagent"))
+  testInstrumentation(project(":instrumentation:internal:internal-class-loader:javaagent", configuration = "shaded"))
 }

instrumentation/internal/internal-class-loader/javaagent/build.gradle.kts

@@ -1,10 +1,13 @@
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
+
 plugins {
   id("otel.javaagent-instrumentation")
 }
 
 dependencies {
   compileOnly(project(":javaagent-bootstrap"))
   compileOnly(project(":javaagent-tooling"))
+  compileOnly(project(":instrumentation:internal:internal-class-loader:compile-stub"))
 
   testImplementation(project(":javaagent-bootstrap"))
 
@@ -21,3 +24,32 @@ dependencies {
   testImplementation("org.eclipse.platform:org.eclipse.osgi:3.13.200")
   testImplementation("org.apache.felix:org.apache.felix.framework:6.0.2")
 }
+
+val shadedJar by tasks.registering(ShadowJar::class) {
+  from(zipTree(tasks.jar.get().archiveFile))
+  archiveClassifier.set("shaded")
+}
+
+tasks {
+  withType(ShadowJar::class) {
+    relocate("io.opentelemetry.javaagent.instrumentation.internal.classloader.stub", "java.lang")
+  }
+
+  assemble {
+    dependsOn(shadedJar)
+  }
+}
+
+// Create a consumable configuration for the shaded jar. We can't use the "shadow" configuration
+// because that is taken by the agent-testing.jar
+configurations {
+  consumable("shaded") {
+    attributes {
+      attribute(LibraryElements.LIBRARY_ELEMENTS_ATTRIBUTE, objects.named("shaded"))
+    }
+  }
+}
+
+artifacts {
+  add("shaded", shadedJar)
+}

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

@@ -57,10 +57,37 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class LoadClassAdvice {
 
+    // Class loader stub is shaded back to the real class loader class. It is used to call protected
+    // method from the advice that the compiler won't let us call directly. During runtime it is
+    // fine since this code is inlined into subclasses of ClassLoader that can call protected
+    // methods.
     @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
     public static Class<?> onEnter(
-        @Advice.This ClassLoader classLoader, @Advice.Argument(0) String name) {
-      return InjectedClassHelper.loadHelperClass(classLoader, name);
+        @Advice.This java.lang.ClassLoader classLoader,
+        @Advice.This
+            io.opentelemetry.javaagent.instrumentation.internal.classloader.stub.ClassLoader
+                classLoaderStub,
+        @Advice.Argument(0) String name) {
+      InjectedClassHelper.HelperClassInfo helperClassInfo =
+          InjectedClassHelper.getHelperClassInfo(classLoader, name);
+      if (helperClassInfo != null) {
+        Class<?> clazz = classLoaderStub.findLoadedClass(name);
+        if (clazz != null) {
+          return clazz;
+        }
+        try {
+          byte[] bytes = helperClassInfo.getClassBytes();
+          return classLoaderStub.defineClass(
+              name, bytes, 0, bytes.length, helperClassInfo.getProtectionDomain());
+        } catch (LinkageError error) {
+          clazz = classLoaderStub.findLoadedClass(name);
+          if (clazz != null) {
+            return clazz;
+          }
+          throw error;
+        }
+      }
+      return null;
     }
 
     @AssignReturned.ToReturned

instrumentation/internal/internal-lambda/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/internal/lambda/LambdaInstrumentationTest.java

@@ -7,7 +7,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldInstalledMarker;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldInstalledMarker;
 import org.junit.jupiter.api.Test;
 
 class LambdaInstrumentationTest {

instrumentation/internal/internal-reflection/javaagent-integration-tests/src/test/java/ReflectionTest.java

@@ -7,8 +7,8 @@
 
 import instrumentation.TestHelperClass;
 import io.opentelemetry.javaagent.bootstrap.InstrumentationProxy;
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldAccessorMarker;
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldInstalledMarker;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldAccessorMarker;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldInstalledMarker;
 import java.io.ObjectStreamClass;
 import java.io.Serializable;
 import java.lang.reflect.Field;

instrumentation/internal/internal-reflection/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/reflection/ReflectionHelper.java

@@ -6,9 +6,9 @@
 package io.opentelemetry.javaagent.instrumentation.internal.reflection;
 
 import io.opentelemetry.javaagent.bootstrap.InstrumentationProxy;
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldAccessorMarker;
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldDetector;
-import io.opentelemetry.javaagent.bootstrap.VirtualFieldInstalledMarker;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldAccessorMarker;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldDetector;
+import io.opentelemetry.javaagent.bootstrap.field.VirtualFieldInstalledMarker;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.ArrayList;

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/InjectedClassHelper.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.bootstrap;
 
+import java.security.ProtectionDomain;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
 import java.util.function.Function;
@@ -37,21 +38,27 @@ public static boolean isHelperClass(ClassLoader classLoader, String className) {
     return helperClassDetector.test(classLoader, className);
   }
 
-  private static volatile BiFunction<ClassLoader, String, Class<?>> helperClassLoader;
+  private static volatile BiFunction<ClassLoader, String, HelperClassInfo> helperClassInfo;
 
-  public static void internalSetHelperClassLoader(
-      BiFunction<ClassLoader, String, Class<?>> helperClassLoader) {
-    if (InjectedClassHelper.helperClassLoader != null) {
+  public static void internalSetHelperClassInfo(
+      BiFunction<ClassLoader, String, HelperClassInfo> helperClassInfo) {
+    if (InjectedClassHelper.helperClassInfo != null) {
       // Only possible by misuse of this API, just ignore.
       return;
     }
-    InjectedClassHelper.helperClassLoader = helperClassLoader;
+    InjectedClassHelper.helperClassInfo = helperClassInfo;
   }
 
-  public static Class<?> loadHelperClass(ClassLoader classLoader, String className) {
-    if (helperClassLoader == null) {
+  public static HelperClassInfo getHelperClassInfo(ClassLoader classLoader, String className) {
+    if (helperClassInfo == null) {
       return null;
     }
-    return helperClassLoader.apply(classLoader, className);
+    return helperClassInfo.apply(classLoader, className);
+  }
+
+  public interface HelperClassInfo {
+    byte[] getClassBytes();
+
+    ProtectionDomain getProtectionDomain();
   }
 }