Don't load classes on agent classloader during injection
MODIFY
dd-java-agent/tooling/src/main/java/dd/trace/HelperInjector.java
dd-java-agent/tooling/src/main/java/dd/trace/HelperInjector.java
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);
}

MODIFY
dd-java-agent/tooling/src/test/groovy/dd/test/HelperInjectionTest.groovy
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()
}
}

MODIFY
dd-java-agent/tooling/src/test/groovy/dd/test/TestUtils.groovy
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"

Don't load classes on agent classloader during injection