From 944ea9059d6453f454ad76d834b3efaabe774afc Mon Sep 17 00:00:00 2001
From: Rafael Winterhalter <rafael.wth@gmail.com>
Date: Tue, 18 Sep 2018 21:57:03 +0200
Subject: [PATCH] Extend class injector API.

---
 .../dynamic/loading/ClassInjector.java        | 252 +++++++++++-------
 ...ClassInjectorUsingInstrumentationTest.java |   9 +
 .../loading/ClassInjectorUsingLookupTest.java |  13 +-
 .../ClassInjectorUsingReflectionTest.java     |   4 +-
 .../loading/ClassInjectorUsingUnsafeTest.java |   1 +
 5 files changed, 183 insertions(+), 96 deletions(-)

diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java b/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java
index 3f8750e5a98..e8cbfb3a644 100644
--- a/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java
+++ b/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java
@@ -26,6 +26,7 @@
 import java.security.PrivilegedAction;
 import java.security.ProtectionDomain;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
@@ -48,6 +49,13 @@ public interface ClassInjector {
      */
     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.
      *
@@ -56,11 +64,41 @@ public interface ClassInjector {
      */
     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.
      */
     @HashCodeAndEqualsPlugin.Enhance
-    class UsingReflection implements ClassInjector {
+    class UsingReflection extends AbstractBase {
 
         /**
          * The dispatcher to use for accessing a class loader via reflection.
@@ -134,38 +172,26 @@ public UsingReflection(ClassLoader classLoader,
         }
 
         /**
-         * 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.
+         * {@inheritDoc}
          */
-        public static ClassInjector ofSystemClassLoader() {
-            return new UsingReflection(ClassLoader.getSystemClassLoader());
+        public boolean isAlive() {
+            return isAvailable();
         }
 
         /**
          * {@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();
-            Map<TypeDescription, Class<?>> loadedTypes = new HashMap<TypeDescription, Class<?>>();
-            for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) {
-                String typeName = entry.getKey().getName();
-                synchronized (dispatcher.getClassLoadingLock(classLoader, typeName)) {
-                    Class<?> type = dispatcher.findClass(classLoader, typeName);
+            Map<String, Class<?>> result = new HashMap<String, Class<?>>();
+            for (Map.Entry<? extends String, byte[]> entry : types.entrySet()) {
+                synchronized (dispatcher.getClassLoadingLock(classLoader, entry.getKey())) {
+                    Class<?> type = dispatcher.findClass(classLoader, entry.getKey());
                     if (type == null) {
-                        int packageIndex = typeName.lastIndexOf('.');
+                        int packageIndex = entry.getKey().lastIndexOf('.');
                         if (packageIndex != -1) {
-                            String packageName = typeName.substring(0, packageIndex);
-                            PackageDefinitionStrategy.Definition definition = packageDefinitionStrategy.define(classLoader, packageName, typeName);
+                            String packageName = entry.getKey().substring(0, packageIndex);
+                            PackageDefinitionStrategy.Definition definition = packageDefinitionStrategy.define(classLoader, packageName, entry.getKey());
                             if (definition.isDefined()) {
                                 Package definedPackage = dispatcher.getPackage(classLoader, packageName);
                                 if (definedPackage == null) {
@@ -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) {
                         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());
         }
 
         /**
@@ -1243,7 +1287,7 @@ public Package definePackage(ClassLoader classLoader,
      * </p>
      */
     @HashCodeAndEqualsPlugin.Enhance
-    class UsingLookup implements ClassInjector {
+    class UsingLookup extends AbstractBase {
 
         /**
          * The dispatcher to interacting with method handles.
@@ -1286,15 +1330,6 @@ public static UsingLookup of(Object lookup) {
             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.
          *
@@ -1317,15 +1352,32 @@ public UsingLookup in(Class<?> type) {
         /**
          * {@inheritDoc}
          */
-        public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte[]> types) {
-            Map<TypeDescription, Class<?>> loaded = new HashMap<TypeDescription, Class<?>>();
-            for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) {
-                if (!entry.getKey().isSamePackage(TypeDescription.ForLoadedType.of(lookupType()))) {
+        public boolean isAlive() {
+            return isAvailable();
+        }
+
+        /**
+         * {@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);
                 }
-                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();
         }
 
         /**
@@ -1601,7 +1653,7 @@ public Class<?> defineClass(Object lookup, byte[] binaryRepresentation) {
      * </p>
      */
     @HashCodeAndEqualsPlugin.Enhance
-    class UsingUnsafe implements ClassInjector {
+    class UsingUnsafe extends AbstractBase {
 
         /**
          * The dispatcher to use.
@@ -1645,6 +1697,33 @@ public UsingUnsafe(ClassLoader classLoader, 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.
          *
@@ -1672,26 +1751,6 @@ public static ClassInjector ofClassPath() {
             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}.
          */
@@ -1863,7 +1922,7 @@ public Dispatcher initialize() {
      * or the system class path.
      */
     @HashCodeAndEqualsPlugin.Enhance
-    class UsingInstrumentation implements ClassInjector {
+    class UsingInstrumentation extends AbstractBase {
 
         /**
          * The jar file name extension.
@@ -1900,27 +1959,6 @@ class UsingInstrumentation implements ClassInjector {
          */
         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.
          *
@@ -1939,10 +1977,29 @@ protected UsingInstrumentation(File folder,
             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}
          */
-        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);
             try {
                 if (!jarFile.createNewFile()) {
@@ -1950,19 +2007,19 @@ public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte
                 }
                 JarOutputStream jarOutputStream = new JarOutputStream(new FileOutputStream(jarFile));
                 try {
-                    for (Map.Entry<? extends TypeDescription, byte[]> entry : types.entrySet()) {
-                        jarOutputStream.putNextEntry(new JarEntry(entry.getKey().getInternalName() + CLASS_FILE_EXTENSION));
+                    for (Map.Entry<? extends String, byte[]> entry : types.entrySet()) {
+                        jarOutputStream.putNextEntry(new JarEntry(entry.getKey().replace('.', '/') + CLASS_FILE_EXTENSION));
                         jarOutputStream.write(entry.getValue());
                     }
                 } finally {
                     jarOutputStream.close();
                 }
                 target.inject(instrumentation, new JarFile(jarFile));
-                Map<TypeDescription, Class<?>> loaded = new HashMap<TypeDescription, Class<?>>();
-                for (TypeDescription typeDescription : types.keySet()) {
-                    loaded.put(typeDescription, Class.forName(typeDescription.getName(), false, ClassLoader.getSystemClassLoader()));
+                Map<String, Class<?>> result = new HashMap<String, Class<?>>();
+                for (String name : types.keySet()) {
+                    result.put(name, Class.forName(name, false, ClassLoader.getSystemClassLoader()));
                 }
-                return loaded;
+                return result;
             } catch (IOException exception) {
                 throw new IllegalStateException("Cannot write jar file to disk", exception);
             } catch (ClassNotFoundException exception) {
@@ -1970,6 +2027,15 @@ public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte
             }
         }
 
+        /**
+         * 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.
          */
diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingInstrumentationTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingInstrumentationTest.java
index 3891333547f..9ab0c4fdefe 100644
--- a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingInstrumentationTest.java
+++ b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingInstrumentationTest.java
@@ -63,4 +63,13 @@ public void testSystemInjection() throws Exception {
         assertThat(types.get(dynamicType.getTypeDescription()).getName(), is(name));
         assertThat(types.get(dynamicType.getTypeDescription()).getClassLoader(), is(ClassLoader.getSystemClassLoader()));
     }
+
+    @Test
+    @AgentAttachmentRule.Enforce
+    public void testAvailable() {
+        assertThat(ClassInjector.UsingInstrumentation.isAvailable(), is(true));
+        assertThat(ClassInjector.UsingInstrumentation.of(folder,
+                        ClassInjector.UsingInstrumentation.Target.SYSTEM,
+                        ByteBuddyAgent.install()).isAlive(), is(true));
+    }
 }
diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingLookupTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingLookupTest.java
index 31b310abb6d..571f11438d1 100644
--- a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingLookupTest.java
+++ b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingLookupTest.java
@@ -1,10 +1,12 @@
 package net.bytebuddy.dynamic.loading;
 
 import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.agent.ByteBuddyAgent;
 import net.bytebuddy.description.modifier.Ownership;
 import net.bytebuddy.description.modifier.Visibility;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.implementation.MethodCall;
+import net.bytebuddy.test.utility.AgentAttachmentRule;
 import net.bytebuddy.test.utility.JavaVersionRule;
 import org.junit.Before;
 import org.junit.Rule;
@@ -44,13 +46,13 @@ public void testIsAvailable() {
     @Test
     @JavaVersionRule.Enforce(9)
     public void testLookupType() throws Exception {
-        assertThat(ClassInjector.UsingLookup.of(type.getDeclaredMethod("lookup").invoke(null)).lookupType(), is((Object) type));
+        assertThat(ClassInjector.UsingLookup.of(type.getMethod("lookup").invoke(null)).lookupType(), is((Object) type));
     }
 
     @Test
     @JavaVersionRule.Enforce(9)
     public void testLookupInjection() throws Exception {
-        ClassInjector injector = ClassInjector.UsingLookup.of(type.getDeclaredMethod("lookup").invoke(null));
+        ClassInjector injector = ClassInjector.UsingLookup.of(type.getMethod("lookup").invoke(null));
         DynamicType dynamicType = new ByteBuddy()
                 .subclass(Object.class)
                 .name("net.bytebuddy.test.Bar")
@@ -70,4 +72,11 @@ public void testLookupInjectionPropagate() throws Exception {
         assertThat(injector.inject(Collections.singletonMap(dynamicType.getTypeDescription(), dynamicType.getBytes()))
                 .get(dynamicType.getTypeDescription()).getName(), is("net.bytebuddy.test.Bar"));
     }
+
+    @Test
+    @JavaVersionRule.Enforce(9)
+    public void testAvailable() throws Exception {
+        assertThat(ClassInjector.UsingLookup.isAvailable(), is(true));
+        assertThat(ClassInjector.UsingLookup.of(type.getMethod("lookup").invoke(null)).isAlive(), is((Object) true));
+    }
 }
diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingReflectionTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingReflectionTest.java
index b3489949a6e..4651795626c 100644
--- a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingReflectionTest.java
+++ b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingReflectionTest.java
@@ -199,7 +199,8 @@ public void testFaultyDispatcherGetLock() throws Exception {
     public void testInjectionOrderNoPrematureAuxiliaryInjection() throws Exception {
         ClassLoader classLoader = new ByteArrayClassLoader(ClassLoadingStrategy.BOOTSTRAP_LOADER,
                 ClassFileLocator.ForClassLoader.readToNames(Bar.class, Interceptor.class));
-        Class<?> type = new ByteBuddy().rebase(Bar.class)
+        Class<?> type = new ByteBuddy()
+                .rebase(Bar.class)
                 .method(named(BAR))
                 .intercept(MethodDelegation.to(Interceptor.class)).make()
                 .load(classLoader, ClassLoadingStrategy.Default.INJECTION)
@@ -211,6 +212,7 @@ public void testInjectionOrderNoPrematureAuxiliaryInjection() throws Exception {
     @ClassReflectionInjectionAvailableRule.Enforce
     public void testAvailability() throws Exception {
         assertThat(ClassInjector.UsingReflection.isAvailable(), is(true));
+        assertThat(new ClassInjector.UsingReflection(ClassLoader.getSystemClassLoader()).isAlive(), is(true));
         assertThat(new ClassInjector.UsingReflection.Dispatcher.Initializable.Unavailable(null).isAvailable(), is(false));
     }
 
diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingUnsafeTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingUnsafeTest.java
index ac154db3f6e..e2eb9105e89 100644
--- a/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingUnsafeTest.java
+++ b/byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/loading/ClassInjectorUsingUnsafeTest.java
@@ -42,6 +42,7 @@ public void testUnsafeInjection() throws Exception {
     @ClassUnsafeInjectionAvailableRule.Enforce
     public void testAvailability() throws Exception {
         assertThat(ClassInjector.UsingUnsafe.isAvailable(), is(true));
+        assertThat(new ClassInjector.UsingUnsafe(ClassLoader.getSystemClassLoader()).isAlive(), is(true));
         assertThat(new ClassInjector.UsingUnsafe.Dispatcher.Disabled(null).isAvailable(), is(false));
     }