Replace use of ByteBuddy's ClassInjector when redefining modules to add reads

Author: mcculls

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java

@@ -1,14 +1,17 @@
 package datadog.trace.agent.tooling;
 
 import static datadog.trace.bootstrap.AgentClassLoading.INJECTING_HELPERS;
+import static java.util.Arrays.asList;
 
 import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
+import datadog.trace.util.JDK9ModuleAccess;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
+import java.lang.reflect.AnnotatedElement;
 import java.security.CodeSource;
 import java.security.ProtectionDomain;
-import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -20,7 +23,6 @@
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.dynamic.loading.ClassInjector;
-import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
 import net.bytebuddy.utility.JavaModule;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -40,9 +42,9 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
   private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
 
   private final Map<ClassLoader, Boolean> injectedClassLoaders =
-      Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
+      Collections.synchronizedMap(new WeakHashMap<>());
 
-  private final List<WeakReference<Object>> helperModules = new CopyOnWriteArrayList<>();
+  private final List<WeakReference<AnnotatedElement>> helperModules = new CopyOnWriteArrayList<>();
 
   /**
    * Construct HelperInjector.
@@ -71,7 +73,7 @@ public HelperInjector(
     this.requestingName = requestingName;
     this.adviceShader = adviceShader;
 
-    this.helperClassNames = new LinkedHashSet<>(Arrays.asList(helperClassNames));
+    this.helperClassNames = new LinkedHashSet<>(asList(helperClassNames));
   }
 
   public HelperInjector(
@@ -132,12 +134,10 @@ public DynamicType.Builder<?> transform(
           final Map<String, byte[]> classnameToBytes = getHelperMap();
           final Map<String, Class<?>> classes = injectClassLoader(classLoader, classnameToBytes);
 
-          // All datadog helper classes are in the unnamed module
-          // And there's exactly one unnamed module per classloader
-          // Use the module of the first class for convenience
+          // all datadog helper classes are in the unnamed module
+          // and there's exactly one unnamed module per classloader
           if (JavaModule.isSupported()) {
-            final JavaModule javaModule = JavaModule.ofType(classes.values().iterator().next());
-            helperModules.add(new WeakReference<>(javaModule.unwrap()));
+            helperModules.add(new WeakReference<>(JDK9ModuleAccess.getUnnamedModule(classLoader)));
           }
 
           // forcibly initialize any eager helpers
@@ -177,43 +177,42 @@ private Map<String, Class<?>> injectClassLoader(
       final ClassLoader classLoader, final Map<String, byte[]> classnameToBytes) {
     INJECTING_HELPERS.begin();
     try {
-      ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
-      return new ClassInjector.UsingReflection(classLoader, protectionDomain)
-          .injectRaw(classnameToBytes);
+      if (useAgentCodeSource) {
+        ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
+        return new ClassInjector.UsingReflection(classLoader, protectionDomain)
+            .injectRaw(classnameToBytes);
+      } else {
+        return new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
+      }
     } finally {
       INJECTING_HELPERS.end();
     }
   }
 
   private ProtectionDomain createProtectionDomain(final ClassLoader classLoader) {
-    if (useAgentCodeSource) {
-      CodeSource codeSource = HelperInjector.class.getProtectionDomain().getCodeSource();
-      return new ProtectionDomain(codeSource, null, classLoader, null);
-    } else {
-      return ClassLoadingStrategy.NO_PROTECTION_DOMAIN;
-    }
+    CodeSource codeSource = HelperInjector.class.getProtectionDomain().getCodeSource();
+    return new ProtectionDomain(codeSource, null, classLoader, null);
   }
 
   private void ensureModuleCanReadHelperModules(final JavaModule target) {
     if (JavaModule.isSupported() && target != JavaModule.UNSUPPORTED && target.isNamed()) {
-      for (final WeakReference<Object> helperModuleReference : helperModules) {
-        final Object realModule = helperModuleReference.get();
-        if (realModule != null) {
-          final JavaModule helperModule = JavaModule.of(realModule);
-
-          if (!target.canRead(helperModule)) {
-            log.debug("Adding module read from {} to {}", target, helperModule);
-            ClassInjector.UsingInstrumentation.redefineModule(
-                Utils.getInstrumentation(),
-                target,
-                Collections.singleton(helperModule),
-                Collections.<String, Set<JavaModule>>emptyMap(),
-                Collections.<String, Set<JavaModule>>emptyMap(),
-                Collections.<Class<?>>emptySet(),
-                Collections.<Class<?>, List<Class<?>>>emptyMap());
+      AnnotatedElement targetModule = (AnnotatedElement) target.unwrap();
+      Set<AnnotatedElement> extraReads = null;
+      for (final WeakReference<AnnotatedElement> helperModuleReference : helperModules) {
+        final AnnotatedElement helperModule = helperModuleReference.get();
+        if (helperModule != null) {
+          if (!JDK9ModuleAccess.canRead(targetModule, helperModule)) {
+            if (extraReads == null) {
+              extraReads = new HashSet<>();
+            }
+            extraReads.add(helperModule);
           }
         }
       }
+      if (extraReads != null) {
+        log.debug("Adding module reads from {} to {}", targetModule, extraReads);
+        JDK9ModuleAccess.addModuleReads(Utils.getInstrumentation(), targetModule, extraReads);
+      }
     }
   }
 }

internal-api/internal-api-9/src/main/java/datadog/trace/util/JDK9ModuleAccess.java

@@ -0,0 +1,41 @@
+package datadog.trace.util;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.emptySet;
+
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.AnnotatedElement;
+import java.util.Set;
+import javax.annotation.Nullable;
+
+/** Use standard API to work with JPMS modules on Java9+. */
+@SuppressWarnings("Since15")
+public final class JDK9ModuleAccess {
+
+  /** Retrieves a class-loader's unnamed module. */
+  public static @Nullable AnnotatedElement getUnnamedModule(ClassLoader cl) {
+    try {
+      return cl.getUnnamedModule();
+    } catch (Throwable e) {
+      return null; // not available
+    }
+  }
+
+  /** Returns {@code true} if the first module can read the second module. */
+  public static boolean canRead(AnnotatedElement module, AnnotatedElement anotherModule) {
+    return ((java.lang.Module) module).canRead((java.lang.Module) anotherModule);
+  }
+
+  /** Adds extra module reads to the given module. */
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  public static void addModuleReads(
+      Instrumentation inst, AnnotatedElement module, Set<AnnotatedElement> extraReads) {
+    inst.redefineModule(
+        (java.lang.Module) module,
+        (Set) extraReads,
+        emptyMap(),
+        emptyMap(),
+        emptySet(),
+        emptyMap());
+  }
+}