Enable bootstrap instrumentation and helper injection.

This commit is contained in:
Andrew Kent 2018-02-19 17:08:55 -08:00 committed by Tyler Benson
parent f57faba5db
commit 4fffb61a32
5 changed files with 58 additions and 8 deletions

View File

@ -17,6 +17,11 @@ import net.bytebuddy.utility.JavaModule;
@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 class AgentInstaller {
*/
public static ResettableClassFileTransformer installBytebuddyAgent(
final Instrumentation inst, final AgentBuilder.Listener... listeners) {
INSTRUMENTATION = inst;
AgentBuilder agentBuilder =
new AgentBuilder.Default()
.disableClassFormatChanges()

View File

@ -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 @@ public class ClassLoaderMatcher {
@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);
}

View File

@ -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 class HelperInjector implements Transformer {
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 class HelperInjector implements Transformer {
}
}
}
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 class HelperInjector implements Transformer {
+ ". 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);
}

View File

@ -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
*/

View File

@ -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 {