Skip to content

Commit

Permalink
Extend class injector API.
Browse files Browse the repository at this point in the history
  • Loading branch information
@raphw
raphw committed Sep 18, 2018
1 parent c921de9 commit 944ea90
Show file tree
Hide file tree
Showing 5 changed files with 183 additions and 96 deletions.
252 changes: 159 additions & 93 deletions byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java
View file
Expand Up @@ -26,6 +26,7 @@
import java.security.PrivilegedAction; import java.security.PrivilegedAction;
import java.security.ProtectionDomain; import java.security.ProtectionDomain;
import java.util.HashMap; import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map; import java.util.Map;
import java.util.jar.JarEntry; import java.util.jar.JarEntry;
import java.util.jar.JarFile; import java.util.jar.JarFile;
Expand All @@ -48,6 +49,13 @@ public interface ClassInjector {
*/ */
boolean ALLOW_EXISTING_TYPES = false; boolean ALLOW_EXISTING_TYPES = false;


/**
* Indicates if this class injector is available on the current VM.
*
* @return {@code true} if this injector is available on the current VM.
*/
boolean isAlive();

/** /**
* Injects the given types into the represented class loader. * Injects the given types into the represented class loader.
* *
Expand All @@ -56,11 +64,41 @@ public interface ClassInjector {
*/ */
Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types); Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types);


/**
* Injects the given types into the represented class loader using a mapping from name to binary representation.
*
* @param types The types to load via injection.
* @return The loaded types that were passed as arguments.
*/
Map<String, Class<?>> injectRaw(Map<? extends String, byte[]> types);

/**
* An abstract base implementation of a class injector.
*/
abstract class AbstractBase implements ClassInjector {

/**
* {@inheritDoc}
*/
public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types) {
Map<String, byte[]> binaryRepresentations = new LinkedHashMap<String, byte[]>();
for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) {
binaryRepresentations.put(entry.getKey().getName(), entry.getValue());
}
Map<String, Class<?>> loadedTypes = injectRaw(binaryRepresentations);
Map<TypeDescription, Class<?>> result = new LinkedHashMap<TypeDescription, Class<?>>();
for (TypeDescription typeDescription : types.keySet()) {
result.put(typeDescription, loadedTypes.get(typeDescription.getName()));
}
return result;
}
}

/** /**
* A class injector that uses reflective method calls. * A class injector that uses reflective method calls.
*/ */
@HashCodeAndEqualsPlugin.Enhance @HashCodeAndEqualsPlugin.Enhance
class UsingReflection implements ClassInjector { class UsingReflection extends AbstractBase {


/** /**
* The dispatcher to use for accessing a class loader via reflection. * The dispatcher to use for accessing a class loader via reflection.
Expand Down Expand Up @@ -134,38 +172,26 @@ public UsingReflection(ClassLoader classLoader,
} }


/** /**
* Indicates if this class injection is available on the current VM. * {@inheritDoc}
*
* @return {@code true} if this class injection is available.
*/
public static boolean isAvailable() {
return DISPATCHER.isAvailable();
}

/**
* Creates a class injector for the system class loader.
*
* @return A class injector for the system class loader.
*/ */
public static ClassInjector ofSystemClassLoader() { public boolean isAlive() {
return new UsingReflection(ClassLoader.getSystemClassLoader()); return isAvailable();
} }


/** /**
* {@inheritDoc} * {@inheritDoc}
*/ */
public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types) { public Map<String, Class<?>> injectRaw(Map<? extends String, byte[]> types) {
Dispatcher dispatcher = DISPATCHER.initialize(); Dispatcher dispatcher = DISPATCHER.initialize();
Map<TypeDescription, Class<?>> loadedTypes = new HashMap<TypeDescription, Class<?>>(); Map<String, Class<?>> result = new HashMap<String, Class<?>>();
for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) { for (Map.Entry<? extends String, byte[]> entry : types.entrySet()) {
String typeName = entry.getKey().getName(); synchronized (dispatcher.getClassLoadingLock(classLoader, entry.getKey())) {
synchronized (dispatcher.getClassLoadingLock(classLoader, typeName)) { Class<?> type = dispatcher.findClass(classLoader, entry.getKey());
Class<?> type = dispatcher.findClass(classLoader, typeName);
if (type == null) { if (type == null) {
int packageIndex = typeName.lastIndexOf('.'); int packageIndex = entry.getKey().lastIndexOf('.');
if (packageIndex != -1) { if (packageIndex != -1) {
String packageName = typeName.substring(0, packageIndex); String packageName = entry.getKey().substring(0, packageIndex);
PackageDefinitionStrategy.Definition definition = packageDefinitionStrategy.define(classLoader, packageName, typeName); PackageDefinitionStrategy.Definition definition = packageDefinitionStrategy.define(classLoader, packageName, entry.getKey());
if (definition.isDefined()) { if (definition.isDefined()) {
Package definedPackage = dispatcher.getPackage(classLoader, packageName); Package definedPackage = dispatcher.getPackage(classLoader, packageName);
if (definedPackage == null) { if (definedPackage == null) {
Expand All @@ -183,14 +209,32 @@ public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte
} }
} }
} }
type = dispatcher.defineClass(classLoader, typeName, entry.getValue(), protectionDomain); type = dispatcher.defineClass(classLoader, entry.getKey(), entry.getValue(), protectionDomain);
} else if (forbidExisting) { } else if (forbidExisting) {
throw new IllegalStateException("Cannot inject already loaded type: " + type); throw new IllegalStateException("Cannot inject already loaded type: " + type);
} }
loadedTypes.put(entry.getKey(), type); result.put(entry.getKey(), type);
} }
} }
return loadedTypes; return result;
}

/**
* Indicates if this class injection is available on the current VM.
*
* @return {@code true} if this class injection is available.
*/
public static boolean isAvailable() {
return DISPATCHER.isAvailable();
}

/**
* Creates a class injector for the system class loader.
*
* @return A class injector for the system class loader.
*/
public static ClassInjector ofSystemClassLoader() {
return new UsingReflection(ClassLoader.getSystemClassLoader());
} }


/** /**
Expand Down Expand Up @@ -1243,7 +1287,7 @@ public Package definePackage(ClassLoader classLoader,
* </p> * </p>
*/ */
@HashCodeAndEqualsPlugin.Enhance @HashCodeAndEqualsPlugin.Enhance
class UsingLookup implements ClassInjector { class UsingLookup extends AbstractBase {


/** /**
* The dispatcher to interacting with method handles. * The dispatcher to interacting with method handles.
Expand Down Expand Up @@ -1286,15 +1330,6 @@ public static UsingLookup of(Object lookup) {
return new UsingLookup(DISPATCHER.dropLookupMode(lookup, Opcodes.ACC_PRIVATE)); return new UsingLookup(DISPATCHER.dropLookupMode(lookup, Opcodes.ACC_PRIVATE));
} }


/**
* Checks if the current VM is capable of defining classes using a method handle lookup.
*
* @return {@code true} if the current VM is capable of defining classes using a lookup.
*/
public static boolean isAvailable() {
return DISPATCHER.isAlive();
}

/** /**
* Returns the lookup type this injector is based upon. * Returns the lookup type this injector is based upon.
* *
Expand All @@ -1317,15 +1352,32 @@ public UsingLookup in(Class<?> type) {
/** /**
* {@inheritDoc} * {@inheritDoc}
*/ */
public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types) { public boolean isAlive() {
Map<TypeDescription, Class<?>> loaded = new HashMap<TypeDescription, Class<?>>(); return isAvailable();
for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) { }
if (!entry.getKey().isSamePackage(TypeDescription.ForLoadedType.of(lookupType()))) {
/**
* {@inheritDoc}
*/
public Map<String, Class<?>> injectRaw(Map<? extends String, byte[]> types) {
Map<String, Class<?>> result = new HashMap<String, Class<?>>();
for (Map.Entry<? extends String, byte[]> entry : types.entrySet()) {
int index = entry.getKey().lastIndexOf('.');
if (index == -1 || !entry.getKey().substring(0, index).equals(TypeDescription.ForLoadedType.of(lookupType()).getPackage().getName())) {
throw new IllegalArgumentException(entry.getKey() + " must be defined in the same package as " + lookup); throw new IllegalArgumentException(entry.getKey() + " must be defined in the same package as " + lookup);
} }
loaded.put(entry.getKey(), DISPATCHER.defineClass(lookup, entry.getValue())); result.put(entry.getKey(), DISPATCHER.defineClass(lookup, entry.getValue()));
} }
return loaded; return result;
}

/**
* Checks if the current VM is capable of defining classes using a method handle lookup.
*
* @return {@code true} if the current VM is capable of defining classes using a lookup.
*/
public static boolean isAvailable() {
return DISPATCHER.isAlive();
} }


/** /**
Expand Down Expand Up @@ -1601,7 +1653,7 @@ public Class<?> defineClass(Object lookup, byte[] binaryRepresentation) {
* </p> * </p>
*/ */
@HashCodeAndEqualsPlugin.Enhance @HashCodeAndEqualsPlugin.Enhance
class UsingUnsafe implements ClassInjector { class UsingUnsafe extends AbstractBase {


/** /**
* The dispatcher to use. * The dispatcher to use.
Expand Down Expand Up @@ -1645,6 +1697,33 @@ public UsingUnsafe(ClassLoader classLoader, ProtectionDomain protectionDomain) {
this.protectionDomain = protectionDomain; this.protectionDomain = protectionDomain;
} }


/**
* {@inheritDoc}
*/
public boolean isAlive() {
return isAvailable();
}

/**
* {@inheritDoc}
*/
public Map<String, Class<?>> injectRaw(Map<? extends String, byte[]> types) {
Dispatcher dispatcher = DISPATCHER.initialize();
Map<String, Class<?>> result = new HashMap<String, Class<?>>();
synchronized (classLoader == null
? BOOTSTRAP_LOADER_LOCK
: classLoader) {
for (Map.Entry<? extends String, byte[]> entry : types.entrySet()) {
try {
result.put(entry.getKey(), Class.forName(entry.getKey(), false, classLoader));
} catch (ClassNotFoundException ignored) {
result.put(entry.getKey(), dispatcher.defineClass(classLoader, entry.getKey(), entry.getValue(), protectionDomain));
}
}
}
return result;
}

/** /**
* Checks if unsafe class injection is available on the current VM. * Checks if unsafe class injection is available on the current VM.
* *
Expand Down Expand Up @@ -1672,26 +1751,6 @@ public static ClassInjector ofClassPath() {
return new UsingUnsafe(ClassLoader.getSystemClassLoader()); return new UsingUnsafe(ClassLoader.getSystemClassLoader());
} }


/**
* {@inheritDoc}
*/
public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types) {
Dispatcher dispatcher = DISPATCHER.initialize();
Map<TypeDescription, Class<?>> loaded = new HashMap<TypeDescription, Class<?>>();
synchronized (classLoader == null
? BOOTSTRAP_LOADER_LOCK
: classLoader) {
for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) {
try {
loaded.put(entry.getKey(), Class.forName(entry.getKey().getName(), false, classLoader));
} catch (ClassNotFoundException ignored) {
loaded.put(entry.getKey(), dispatcher.defineClass(classLoader, entry.getKey().getName(), entry.getValue(), protectionDomain));
}
}
}
return loaded;
}

/** /**
* A dispatcher for using {@code sun.misc.Unsafe}. * A dispatcher for using {@code sun.misc.Unsafe}.
*/ */
Expand Down Expand Up @@ -1863,7 +1922,7 @@ public Dispatcher initialize() {
* or the system class path. * or the system class path.
*/ */
@HashCodeAndEqualsPlugin.Enhance @HashCodeAndEqualsPlugin.Enhance
class UsingInstrumentation implements ClassInjector { class UsingInstrumentation extends AbstractBase {


/** /**
* The jar file name extension. * The jar file name extension.
Expand Down Expand Up @@ -1900,27 +1959,6 @@ class UsingInstrumentation implements ClassInjector {
*/ */
private final RandomString randomString; private final RandomString randomString;


/**
* Creates an instrumentation-based class injector.
*
* @param folder The folder to be used for storing jar files.
* @param target A representation of the target path to which classes are to be appended.
* @param instrumentation The instrumentation to use for appending to the class path or the boot path.
* @return An appropriate class injector that applies instrumentation.
*/
public static ClassInjector of(File folder, Target target, Instrumentation instrumentation) {
return new UsingInstrumentation(folder, target, instrumentation, new RandomString());
}

/**
* Returns {@code true} if this class injector is available on this VM.
*
* @return {@code true} if this class injector is available on this VM.
*/
public static boolean isAvailable() {
return DISPATCHER.isAlive();
}

/** /**
* Creates an instrumentation-based class injector. * Creates an instrumentation-based class injector.
* *
Expand All @@ -1939,37 +1977,65 @@ protected UsingInstrumentation(File folder,
this.randomString = randomString; this.randomString = randomString;
} }


/**
* Creates an instrumentation-based class injector.
*
* @param folder The folder to be used for storing jar files.
* @param target A representation of the target path to which classes are to be appended.
* @param instrumentation The instrumentation to use for appending to the class path or the boot path.
* @return An appropriate class injector that applies instrumentation.
*/
public static ClassInjector of(File folder, Target target, Instrumentation instrumentation) {
return new UsingInstrumentation(folder, target, instrumentation, new RandomString());
}

/** /**
* {@inheritDoc} * {@inheritDoc}
*/ */
public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types) { public boolean isAlive() {
return isAvailable();
}

/**
* {@inheritDoc}
*/
public Map<String, Class<?>> injectRaw(Map<? extends String, byte[]> types) {
File jarFile = new File(folder, JAR + randomString.nextString() + "." + JAR); File jarFile = new File(folder, JAR + randomString.nextString() + "." + JAR);
try { try {
if (!jarFile.createNewFile()) { if (!jarFile.createNewFile()) {
throw new IllegalStateException("Cannot create file " + jarFile); throw new IllegalStateException("Cannot create file " + jarFile);
} }
JarOutputStream jarOutputStream = new JarOutputStream(new FileOutputStream(jarFile)); JarOutputStream jarOutputStream = new JarOutputStream(new FileOutputStream(jarFile));
try { try {
for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) { for (Map.Entry<? extends String, byte[]> entry : types.entrySet()) {
jarOutputStream.putNextEntry(new JarEntry(entry.getKey().getInternalName() + CLASS_FILE_EXTENSION)); jarOutputStream.putNextEntry(new JarEntry(entry.getKey().replace('.', '/') + CLASS_FILE_EXTENSION));
jarOutputStream.write(entry.getValue()); jarOutputStream.write(entry.getValue());
} }
} finally { } finally {
jarOutputStream.close(); jarOutputStream.close();
} }
target.inject(instrumentation, new JarFile(jarFile)); target.inject(instrumentation, new JarFile(jarFile));
Map<TypeDescription, Class<?>> loaded = new HashMap<TypeDescription, Class<?>>(); Map<String, Class<?>> result = new HashMap<String, Class<?>>();
for (TypeDescription typeDescription : types.keySet()) { for (String name : types.keySet()) {
loaded.put(typeDescription, Class.forName(typeDescription.getName(), false, ClassLoader.getSystemClassLoader())); result.put(name, Class.forName(name, false, ClassLoader.getSystemClassLoader()));
} }
return loaded; return result;
} catch (IOException exception) { } catch (IOException exception) {
throw new IllegalStateException("Cannot write jar file to disk", exception); throw new IllegalStateException("Cannot write jar file to disk", exception);
} catch (ClassNotFoundException exception) { } catch (ClassNotFoundException exception) {
throw new IllegalStateException("Cannot load injected class", exception); throw new IllegalStateException("Cannot load injected class", exception);
} }
} }


/**
* Returns {@code true} if this class injector is available on this VM.
*
* @return {@code true} if this class injector is available on this VM.
*/
public static boolean isAvailable() {
return DISPATCHER.isAlive();
}

/** /**
* A dispatcher to interact with the instrumentation API. * A dispatcher to interact with the instrumentation API.
*/ */
Expand Down

0 comments on commit 944ea90

Please sign in to comment.