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

Ensure helper classes are injected in a fixed order
This commit is contained in:
Tyler Benson 2018-01-16 15:53:13 -05:00 committed by GitHub
commit 1bb01c256b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 55 additions and 28 deletions

View File

@ -9,8 +9,6 @@ import java.io.FileOutputStream;
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.JarOutputStream;
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 class TestUtils {
}
}
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.
*

View File

@ -3,8 +3,9 @@ package datadog.trace.agent.tooling;
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 class HelperInjector implements Transformer {
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);

View File

@ -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 class Utils {
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() {}
}