Merge pull request #190 from DataDog/tyler/fix-injection-order

Ensure helper classes are injected in a fixed order

Author: tylerbenson

dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java

@@ -9,8 +9,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.lang.reflect.Field;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
 import java.net.URL;
 import java.util.UUID;
 import java.util.concurrent.Callable;
@@ -19,15 +17,6 @@
 import java.util.jar.Manifest;
 
 public class TestUtils {
-  private static Method findLoadedClassMethod = null;
-
-  static {
-    try {
-      findLoadedClassMethod = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
-    } catch (NoSuchMethodException | SecurityException e) {
-      throw new IllegalStateException(e);
-    }
-  }
 
   public static void registerOrReplaceGlobalTracer(final Tracer tracer) {
     try {
@@ -65,18 +54,6 @@ public static <T extends Object> Object runUnderTrace(
     }
   }
 
-  public static boolean isClassLoaded(final String className, final ClassLoader classLoader) {
-    try {
-      findLoadedClassMethod.setAccessible(true);
-      final Class<?> loadedClass = (Class<?>) findLoadedClassMethod.invoke(classLoader, className);
-      return null != loadedClass && loadedClass.getClassLoader() == classLoader;
-    } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
-      throw new IllegalStateException(e);
-    } finally {
-      findLoadedClassMethod.setAccessible(false);
-    }
-  }
-
   /**
    * Create a temporary jar on the filesystem with the bytes of the given classes.
    *

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

@@ -3,8 +3,9 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
 import lombok.extern.slf4j.Slf4j;
@@ -26,15 +27,17 @@ public class HelperInjector implements Transformer {
    * Construct HelperInjector.
    *
    * @param helperClassNames binary names of the helper classes to inject. These class names must be
-   *     resolvable by the classloader returned by DDAdvice#getAgentClassLoader()
+   *     resolvable by the classloader returned by DDAdvice#getAgentClassLoader(). Classes are
+   *     injected in the order provided. This is important if there is interdependency between
+   *     helper classes that requires them to be injected in a specific order.
    */
   public HelperInjector(final String... helperClassNames) {
-    this.helperClassNames = new HashSet<>(Arrays.asList(helperClassNames));
+    this.helperClassNames = new LinkedHashSet<>(Arrays.asList(helperClassNames));
   }
 
   private synchronized Map<TypeDescription, byte[]> getHelperMap() throws IOException {
     if (helperMap == null) {
-      helperMap = new HashMap<>(helperClassNames.size());
+      helperMap = new LinkedHashMap<>(helperClassNames.size());
       for (final String helperName : helperClassNames) {
         final ClassFileLocator locator =
             ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader());
@@ -58,7 +61,29 @@ public DynamicType.Builder<?> transform(
       synchronized (this) {
         if (!injectedClassLoaders.contains(classLoader)) {
           try {
-            new ClassInjector.UsingReflection(classLoader).inject(getHelperMap());
+            final Map<TypeDescription, byte[]> helperMap = getHelperMap();
+            final Set<String> existingClasses = new HashSet<>();
+            final ClassLoader systemCL = ClassLoader.getSystemClassLoader();
+            if (!classLoader.equals(systemCL)) {
+              // Build a list of existing helper classes.
+              for (final TypeDescription def : helperMap.keySet()) {
+                final String name = def.getName();
+                if (Utils.isClassLoaded(name, systemCL)) {
+                  existingClasses.add(name);
+                }
+              }
+            }
+            new ClassInjector.UsingReflection(classLoader).inject(helperMap);
+            if (!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();
+                if (!existingClasses.contains(name) && Utils.isClassLoaded(name, systemCL)) {
+                  throw new IllegalStateException(
+                      "Class was erroneously loaded on the System classloader: " + name);
+                }
+              }
+            }
           } catch (final Exception e) {
             log.error("Failed to inject helper classes into " + classLoader, e);
             throw new RuntimeException(e);

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

@@ -1,6 +1,19 @@
 package datadog.trace.agent.tooling;
 
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+
 public class Utils {
+  private static Method findLoadedClassMethod = null;
+
+  static {
+    try {
+      findLoadedClassMethod = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
+    } catch (NoSuchMethodException | SecurityException e) {
+      throw new IllegalStateException(e);
+    }
+  }
+
   /** Return the classloader the core agent is running on. */
   public static ClassLoader getAgentClassLoader() {
     return AgentInstaller.class.getClassLoader();
@@ -11,5 +24,17 @@ public static String getResourceName(final String className) {
     return className.replace('.', '/') + ".class";
   }
 
+  public static boolean isClassLoaded(final String className, final ClassLoader classLoader) {
+    try {
+      findLoadedClassMethod.setAccessible(true);
+      final Class<?> loadedClass = (Class<?>) findLoadedClassMethod.invoke(classLoader, className);
+      return null != loadedClass && loadedClass.getClassLoader() == classLoader;
+    } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
+      throw new IllegalStateException(e);
+    } finally {
+      findLoadedClassMethod.setAccessible(false);
+    }
+  }
+
   private Utils() {}
 }