Enable bootstrap instrumentation and helper injection.

Andrew Kent · tylerbenson · commit 4fffb61a32fb · 2018-03-09T14:31:48.000+10:00
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
@@ -17,6 +17,11 @@
 
 @Slf4j
 public class AgentInstaller {
+  private static volatile Instrumentation INSTRUMENTATION;
+
+  public static Instrumentation getInstrumentation() {
+    return INSTRUMENTATION;
+  }
 
   public static ResettableClassFileTransformer installBytebuddyAgent(final Instrumentation inst) {
     return installBytebuddyAgent(inst, new AgentBuilder.Listener[0]);
@@ -30,6 +35,7 @@ public static ResettableClassFileTransformer installBytebuddyAgent(final Instrum
    */
   public static ResettableClassFileTransformer installBytebuddyAgent(
       final Instrumentation inst, final AgentBuilder.Listener... listeners) {
+    INSTRUMENTATION = inst;
     AgentBuilder agentBuilder =
         new AgentBuilder.Default()
             .disableClassFormatChanges()
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java
@@ -13,6 +13,7 @@ public class ClassLoaderMatcher {
   public static final String[] BOOTSTRAP_PACKAGE_PREFIXES = {
     "io.opentracing", "datadog.slf4j", "datadog.trace"
   };
+  public static final ClassLoader BOOTSTRAP_CLASSLOADER = null;
 
   /** A private constructor that must not be invoked. */
   private ClassLoaderMatcher() {
@@ -53,9 +54,9 @@ private SkipClassLoaderMatcher() {}
 
     @Override
     public boolean matches(ClassLoader target) {
-      if (null == target) {
-        // bootstrap instrumentation not supported yet.
-        return true;
+      if (target == BOOTSTRAP_CLASSLOADER) {
+        // Don't skip bootstrap loader
+        return false;
       }
       return shouldSkipClass(target) || shouldSkipInstance(target);
     }
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -1,5 +1,8 @@
 package datadog.trace.agent.tooling;
 
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
+
+import java.io.File;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
@@ -57,14 +60,14 @@ public DynamicType.Builder<?> transform(
       final TypeDescription typeDescription,
       final ClassLoader classLoader,
       final JavaModule module) {
-    if (helperClassNames.size() > 0 && classLoader != null) {
+    if (helperClassNames.size() > 0) {
       synchronized (this) {
         if (!injectedClassLoaders.contains(classLoader)) {
           try {
             final Map<TypeDescription, byte[]> helperMap = getHelperMap();
             final Set<String> existingClasses = new HashSet<>();
             final ClassLoader systemCL = ClassLoader.getSystemClassLoader();
-            if (!classLoader.equals(systemCL)) {
+            if (classLoader != BOOTSTRAP_CLASSLOADER && !classLoader.equals(systemCL)) {
               // Build a list of existing helper classes.
               for (final TypeDescription def : helperMap.keySet()) {
                 final String name = def.getName();
@@ -73,8 +76,16 @@ public DynamicType.Builder<?> transform(
                 }
               }
             }
-            new ClassInjector.UsingReflection(classLoader).inject(helperMap);
-            if (!classLoader.equals(systemCL)) {
+            if (classLoader == BOOTSTRAP_CLASSLOADER) {
+              ClassInjector.UsingInstrumentation.of(
+                      new File(System.getProperty("java.io.tmpdir")),
+                      ClassInjector.UsingInstrumentation.Target.BOOTSTRAP,
+                      AgentInstaller.getInstrumentation())
+                  .inject(helperMap);
+            } else {
+              new ClassInjector.UsingReflection(classLoader).inject(helperMap);
+            }
+            if (classLoader != BOOTSTRAP_CLASSLOADER && !classLoader.equals(systemCL)) {
               for (final TypeDescription def : helperMap.keySet()) {
                 // Ensure we didn't add any helper classes to the system CL.
                 final String name = def.getName();
@@ -91,7 +102,9 @@ public DynamicType.Builder<?> transform(
                     + ". Failed to inject helper classes into instance "
                     + classLoader
                     + " of type "
-                    + classLoader.getClass().getName(),
+                    + (classLoader == BOOTSTRAP_CLASSLOADER
+                        ? "<bootstrap>"
+                        : classLoader.getClass().getName()),
                 e);
             throw new RuntimeException(e);
           }
diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ClassLoaderMatcherTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ClassLoaderMatcherTest.groovy
@@ -29,6 +29,11 @@ class ClassLoaderMatcherTest extends Specification {
     !ClassLoaderMatcher.skipClassLoader().matches(emptyLoader)
   }
 
+  def "does not skip bootstrap classloader"() {
+    expect:
+    !ClassLoaderMatcher.skipClassLoader().matches(null)
+  }
+
   /*
    * A URLClassloader which only delegates java.* classes
    */
diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy
@@ -1,7 +1,12 @@
 package datadog.trace.agent.test
 
+import datadog.trace.agent.tooling.AgentInstaller
+
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER
+
 import datadog.trace.agent.tooling.HelperInjector
 import datadog.trace.agent.tooling.Utils
+import net.bytebuddy.agent.ByteBuddyAgent
 import spock.lang.Specification
 
 import java.lang.reflect.Method
@@ -31,6 +36,26 @@ class HelperInjectionTest extends Specification {
     emptyLoader?.close()
   }
 
+  def "helpers injected on bootstrap classloader"() {
+    setup:
+    ByteBuddyAgent.install()
+    AgentInstaller.installBytebuddyAgent(ByteBuddyAgent.getInstrumentation())
+    String helperClassName = HelperInjectionTest.getPackage().getName() + '.HelperClass'
+    HelperInjector injector = new HelperInjector(helperClassName)
+    URLClassLoader bootstrapChild = new URLClassLoader(new URL[0], (ClassLoader) null)
+
+    when:
+    bootstrapChild.loadClass(helperClassName)
+    then:
+    thrown ClassNotFoundException
+
+    when:
+    injector.transform(null, null, BOOTSTRAP_CLASSLOADER, null)
+    Class<?> helperClass = bootstrapChild.loadClass(helperClassName)
+    then:
+    helperClass.getClassLoader() == BOOTSTRAP_CLASSLOADER
+  }
+
   private static boolean isClassLoaded(String className, ClassLoader classLoader) {
     final Method findLoadedClassMethod = ClassLoader.getDeclaredMethod("findLoadedClass", String)
     try {
