Don't load classes on agent classloader during injection

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

@@ -1,5 +1,6 @@
 package dd.trace;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -30,14 +31,15 @@ public class HelperInjector implements Transformer {
     this.helperClassNames = new HashSet<String>(Arrays.asList(helperClassNames));
   }
 
-  private synchronized Map<TypeDescription, byte[]> getHelperMap() throws ClassNotFoundException {
+  private synchronized Map<TypeDescription, byte[]> getHelperMap() throws IOException {
     if (helperMap == null) {
       helperMap = new HashMap<TypeDescription, byte[]>(helperClassNames.size());
       for (String helperName : helperClassNames) {
-        Class<?> helper = DDAdvice.getAgentClassLoader().loadClass(helperName);
-        helperMap.put(
-            new TypeDescription.ForLoadedType(helper),
-            ClassFileLocator.ForClassLoader.read(helper).resolve());
+        final ClassFileLocator locator =
+            ClassFileLocator.ForClassLoader.of(DDAdvice.getAgentClassLoader());
+        final byte[] classBytes = locator.locate(helperName).resolve();
+        final TypeDescription typeDesc = new TypeDescription.Latent(helperName, 0, null);
+        helperMap.put(typeDesc, classBytes);
       }
     }
     return helperMap;
@@ -54,9 +56,9 @@ public class HelperInjector implements Transformer {
         if (!injectedClassLoaders.contains(classLoader)) {
           try {
             new ClassInjector.UsingReflection(classLoader).inject(getHelperMap());
-          } catch (ClassNotFoundException cnfe) {
-            log.error("Failed to inject helper classes into " + classLoader, cnfe);
-            throw new RuntimeException(cnfe);
+          } catch (Exception e) {
+            log.error("Failed to inject helper classes into " + classLoader, e);
+            throw new RuntimeException(e);
           }
           injectedClassLoaders.add(classLoader);
         }

dd-java-agent/tooling/src/test/groovy/dd/test/HelperInjectionTest.groovy

@@ -2,6 +2,8 @@ package dd.test
 
 import static dd.test.TestUtils.createJarWithClasses
 
+import dd.trace.DDAdvice
+import dd.trace.HelperInjector
 import java.lang.reflect.Method
 import net.bytebuddy.agent.ByteBuddyAgent
 import net.bytebuddy.agent.builder.AgentBuilder
@@ -24,26 +26,36 @@ class HelperInjectionTest extends Specification {
 
   def "helpers injected to non-delegating classloader"() {
     setup:
-    URL[] classpath = [createJarWithClasses(TestInstrumentation.ClassToInstrument)]
+    String helperClassName = TestInstrumentation.getName() + '$HelperClass'
+    String instrumentationClassName = TestInstrumentation.getName() + '$ClassToInstrument'
+    HelperInjector injector = new HelperInjector(TestInstrumentation.getName() + '$HelperClass')
+    URLClassLoader emptyLoader = new URLClassLoader(new URL[0], (ClassLoader)null)
+    injector.transform(null, null, emptyLoader, null)
+    // injecting into emptyLoader should not load on agent's classloader
+    assert !TestUtils.isClassLoaded(helperClassName, DDAdvice.getAgentClassLoader())
+    assert TestUtils.isClassLoaded(helperClassName, emptyLoader)
+
+    URL[] classpath = [createJarWithClasses(instrumentationClassName)]
     URLClassLoader classloader = new URLClassLoader(classpath, (ClassLoader)null)
 
     when:
-    classloader.loadClass(TestInstrumentation.HelperClass.getName())
+    classloader.loadClass(helperClassName)
     then:
     thrown ClassNotFoundException
 
     when:
-    Class<?> instrumentedClass = classloader.loadClass(TestInstrumentation.ClassToInstrument.getName())
+    Class<?> instrumentedClass = classloader.loadClass(instrumentationClassName)
     Method instrumentedMethod = instrumentedClass.getMethod("isInstrumented")
     then:
     instrumentedMethod.invoke(null)
 
     when:
-    classloader.loadClass(TestInstrumentation.HelperClass.getName())
+    classloader.loadClass(helperClassName)
     then:
     noExceptionThrown()
 
     cleanup:
     classloader?.close()
+    emptyLoader?.close()
   }
 }

dd-java-agent/tooling/src/test/groovy/dd/test/TestUtils.groovy

@@ -4,6 +4,7 @@ import dd.trace.Instrumenter
 import io.opentracing.ActiveSpan
 import io.opentracing.Tracer
 import io.opentracing.util.GlobalTracer
+import java.lang.reflect.Method
 import java.util.jar.JarEntry
 import java.util.jar.JarOutputStream
 import java.util.jar.Manifest
@@ -55,6 +56,17 @@ class TestUtils {
     }
   }
 
+  private static Method findLoadedClassMethod = ClassLoader.getDeclaredMethod("findLoadedClass", String)
+
+  static boolean isClassLoaded(String className, ClassLoader classLoader) {
+    try {
+      findLoadedClassMethod.setAccessible(true)
+      return null != findLoadedClassMethod.invoke(classLoader, className)
+    } finally {
+      findLoadedClassMethod.setAccessible(false)
+    }
+  }
+
   /** com.foo.Bar -> com/foo/Bar.class */
   static String getResourceName(Class<?> clazz) {
     return clazz.getName().replace('.', '/') + ".class"