Code modifiers and test extension

build.gradle

@@ -111,6 +111,10 @@ dependencies {
     api 'org.projectlombok:lombok:1.18.12'
     annotationProcessor 'org.projectlombok:lombok:1.18.12'
 
+    // Code modification
+    api "org.ow2.asm:asm:6.2.1"
+    api "org.ow2.asm:asm-tree:6.2.1"
+
     // Path finding
     api 'com.github.MadMartian:hydrazine-path-finding:1.3.1'
 

src/main/java/fr/themode/demo/Main.java

@@ -36,6 +36,7 @@ public class Main {
         commandManager.register(new ShutdownCommand());
         commandManager.register(new TeleportCommand());
 
+
         StorageManager storageManager = MinecraftServer.getStorageManager();
         storageManager.defineDefaultStorageSystem(FileStorageSystem::new);
 

src/main/java/fr/themode/demo/MainWrapper.java

@@ -1,11 +0,0 @@
-package fr.themode.demo;
-
-import net.minestom.server.Bootstrap;
-
-public class MainWrapper {
-
-    public static void main(String[] args) {
-        Bootstrap.bootstrap("fr.themode.demo.Main", args);
-    }
-
-}

src/main/java/net/minestom/server/Bootstrap.java

@@ -13,6 +13,12 @@ public class Bootstrap {
     public static void bootstrap(String mainClassFullName, String[] args) {
         try {
             ClassLoader classLoader = new MinestomOverwriteClassLoader(Bootstrap.class.getClassLoader());
+
+            // ensure extensions are loaded when starting the server
+            Class<?> serverClass = classLoader.loadClass("net.minestom.server.MinecraftServer");
+            Method init = serverClass.getMethod("init");
+            init.invoke(null);
+
             Class<?> mainClass = classLoader.loadClass(mainClassFullName);
             Method main = mainClass.getDeclaredMethod("main", String[].class);
             main.invoke(null, new Object[] { args });

src/main/java/net/minestom/server/extensions/ExtensionManager.java

@@ -1,18 +1,13 @@
 package net.minestom.server.extensions;
 
-import com.google.gson.Gson;
+import com.google.gson.*;
-import com.google.gson.JsonObject;
-import com.google.gson.JsonParser;
 import lombok.extern.slf4j.Slf4j;
 import net.minestom.server.extras.selfmodification.MinestomOverwriteClassLoader;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.LoggerFactory;
 
-import java.io.File;
+import java.io.*;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
@@ -24,9 +19,11 @@ import java.util.zip.ZipFile;
 
 @Slf4j
 public class ExtensionManager {
-    private final Map<URL, URLClassLoader> extensionLoaders = new HashMap<>();
+    private final Map<String, URLClassLoader> extensionLoaders = new HashMap<>();
     private final Map<String, Extension> extensions = new HashMap<>();
     private final File extensionFolder = new File("extensions");
+    private final static String INDEV_CLASSES_FOLDER = "minestom.extension.indevfolder.classes";
+    private final static String INDEV_RESOURCES_FOLDER = "minestom.extension.indevfolder.resources";
 
     public ExtensionManager() {
     }
@@ -42,20 +39,29 @@ public class ExtensionManager {
         List<DiscoveredExtension> discoveredExtensions = discoverExtensions();
         setupCodeModifiers(discoveredExtensions);
 
-        for (DiscoveredExtension extension : discoveredExtensions) {
+        for (DiscoveredExtension discoveredExtension : discoveredExtensions) {
             URLClassLoader loader;
-            File file = extension.jarFile;
+            URL[] urls = new URL[discoveredExtension.files.length];
             try {
-                URL url = file.toURI().toURL();
+                for (int i = 0; i < urls.length; i++) {
-                loader = loadJar(url);
+                    urls[i] = discoveredExtension.files[i].toURI().toURL();
-                extensionLoaders.put(url, loader);
+                }
+                loader = newClassLoader(urls);
             } catch (MalformedURLException e) {
-                log.error(String.format("Failed to get URL for file %s.", file.getPath()));
+                log.error("Failed to get URL.", e);
                 return;
             }
             InputStream extensionInputStream = loader.getResourceAsStream("extension.json");
             if (extensionInputStream == null) {
-                log.error(String.format("Failed to find extension.json in the file '%s'.", file.getPath()));
+                StringBuilder urlsString = new StringBuilder();
+                for (int i = 0; i < urls.length; i++) {
+                    URL url = urls[i];
+                    if(i != 0) {
+                        urlsString.append(" ; ");
+                    }
+                    urlsString.append("'").append(url.toString()).append("'");
+                }
+                log.error(String.format("Failed to find extension.json in the urls '%s'.", urlsString));
                 return;
             }
             JsonObject extensionDescription = JsonParser.parseReader(new InputStreamReader(extensionInputStream)).getAsJsonObject();
@@ -63,6 +69,8 @@ public class ExtensionManager {
             String mainClass = extensionDescription.get("entrypoint").getAsString();
             String extensionName = extensionDescription.get("name").getAsString();
 
+            extensionLoaders.put(extensionName, loader);
+
             if (extensions.containsKey(extensionName.toLowerCase())) {
                 log.error(String.format("An extension called '%s' has already been registered.", extensionName));
                 return;
@@ -150,7 +158,22 @@ public class ExtensionManager {
                     InputStreamReader reader = new InputStreamReader(f.getInputStream(f.getEntry("extension.json")))) {
 
                 DiscoveredExtension extension = new DiscoveredExtension();
-                extension.jarFile = file;
+                extension.files = new File[]{file};
+                extension.description = gson.fromJson(reader, JsonObject.class);
+                extensions.add(extension);
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        }
+
+        // this allows developers to have their extension discovered while working on it, without having to build a jar and put in the extension folder
+        if(System.getProperty(INDEV_CLASSES_FOLDER) != null && System.getProperty(INDEV_RESOURCES_FOLDER) != null) {
+            log.info("Found indev folders for extension. Adding to list of discovered extensions.");
+            String extensionClasses = System.getProperty(INDEV_CLASSES_FOLDER);
+            String extensionResources = System.getProperty(INDEV_RESOURCES_FOLDER);
+            try(InputStreamReader reader = new InputStreamReader(new FileInputStream(new File(extensionResources, "extension.json")))) {
+                DiscoveredExtension extension = new DiscoveredExtension();
+                extension.files = new File[] { new File(extensionClasses), new File(extensionResources) };
                 extension.description = gson.fromJson(reader, JsonObject.class);
                 extensions.add(extension);
             } catch (IOException e) {
@@ -163,11 +186,11 @@ public class ExtensionManager {
     /**
      * Loads a URL into the classpath.
      *
-     * @param url {@link URL} (usually a JAR) that should be loaded.
+     * @param urls {@link URL} (usually a JAR) that should be loaded.
      */
     @NotNull
-    public URLClassLoader loadJar(@NotNull URL url) {
+    public URLClassLoader newClassLoader(@NotNull URL[] urls) {
-        return URLClassLoader.newInstance(new URL[]{url}, ExtensionManager.class.getClassLoader());
+        return URLClassLoader.newInstance(urls, ExtensionManager.class.getClassLoader());
     }
 
     @NotNull
@@ -186,7 +209,7 @@ public class ExtensionManager {
     }
 
     @NotNull
-    public Map<URL, URLClassLoader> getExtensionLoaders() {
+    public Map<String, URLClassLoader> getExtensionLoaders() {
         return new HashMap<>(extensionLoaders);
     }
 
@@ -196,28 +219,30 @@ public class ExtensionManager {
     private void setupCodeModifiers(List<DiscoveredExtension> extensions) {
         ClassLoader cl = getClass().getClassLoader();
         if(!(cl instanceof MinestomOverwriteClassLoader)) {
-            log.warning("Current class loader is not a MinestomOverwriteClassLoader, but "+cl+". This disables code modifiers (Mixin support is therefore disabled)");
+            log.warn("Current class loader is not a MinestomOverwriteClassLoader, but "+cl+". This disables code modifiers (Mixin support is therefore disabled)");
             return;
         }
         MinestomOverwriteClassLoader modifiableClassLoader = (MinestomOverwriteClassLoader)cl;
         log.info("Start loading code modifiers...");
         for(DiscoveredExtension extension : extensions) {
             try {
-                if(extension.description.has("codeModifier")) {
+                if(extension.description.has("codeModifiers")) {
-                    String codeModifierClass = extension.description.get("codeModifier").getAsString();
+                    JsonArray codeModifierClasses = extension.description.getAsJsonArray("codeModifiers");
-                    modifiableClassLoader.loadModifier(extension.jarFile, codeModifierClass);
+                    for(JsonElement elem : codeModifierClasses) {
+                        modifiableClassLoader.loadModifier(extension.files, elem.getAsString());
+                    }
                 }
                 // TODO: special support for mixins
             } catch (Exception e) {
                 e.printStackTrace();
-                log.error("Failed to load code modifier for extension "+extension.jarFile, e);
+                log.error("Failed to load code modifier for extension in files: "+Arrays.toString(extension.files), e);
             }
         }
         log.info("Done loading code modifiers.");
     }
 
     private class DiscoveredExtension {
-        private File jarFile;
+        private File[] files;
         private JsonObject description;
     }
 }

src/main/java/net/minestom/server/extras/selfmodification/CodeModifier.java

@@ -0,0 +1,22 @@
+package net.minestom.server.extras.selfmodification;
+
+import org.objectweb.asm.tree.ClassNode;
+
+/**
+ * Will be called by {@link MinestomOverwriteClassLoader} to transform classes at load-time
+ */
+public abstract class CodeModifier {
+    /**
+     * Must return true iif the class node has been modified
+     * @param source
+     * @return
+     */
+    public abstract boolean transform(ClassNode source);
+
+    /**
+     * Beginning of the class names to transform.
+     * 'null' is allowed to transform any class, but not recommended
+     * @return
+     */
+    public abstract String getNamespace();
+}

src/main/java/net/minestom/server/extras/selfmodification/MinestomOverwriteClassLoader.java

@@ -1,15 +1,61 @@
 package net.minestom.server.extras.selfmodification;
 
+import lombok.extern.slf4j.Slf4j;
+import org.jetbrains.annotations.NotNull;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.tree.ClassNode;
+
 import java.io.File;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Set;
 
-// TODO: register code modifiers
+@Slf4j
 public class MinestomOverwriteClassLoader extends URLClassLoader {
 
+    /**
+     * Classes that cannot be loaded/modified by this classloader.
+     * Will go through parent class loader
+     */
+    private static final Set<String> protectedClasses = Set.of(
+            "net.minestom.server.extras.selfmodification.CodeModifier",
+            "net.minestom.server.extras.selfmodification.MinestomOverwriteClassLoader"
+    );
+    private final URLClassLoader asmClassLoader;
+
+    // TODO: replace by tree to optimize lookup times. We can use the fact that package names are split with '.' to allow for fast lookup
+    // TODO: eg. Node("java", Node("lang"), Node("io")). Loading "java.nio.Channel" would apply modifiers from "java", but not "java.io" or "java.lang".
+    // TODO: that's an example, please don't modify standard library classes. And this classloader should not let you do it because it first asks the platform classloader
+
+    // TODO: priorities?
+    private List<CodeModifier> modifiers = new LinkedList<>();
+    private final Method findParentLoadedClass;
+    private final Class<?> loadedCodeModifier;
+
     public MinestomOverwriteClassLoader(ClassLoader parent) {
         super("Minestom ClassLoader", loadURLs(), parent);
+        try {
+            findParentLoadedClass = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
+            findParentLoadedClass.setAccessible(true);
+        } catch (NoSuchMethodException e) {
+            e.printStackTrace();
+            throw new Error("Failed to access ClassLoader#findLoadedClass", e);
+        }
+
+        try {
+            loadedCodeModifier = loadClass("net.minestom.server.extras.selfmodification.CodeModifier");
+        } catch (ClassNotFoundException e) {
+            e.printStackTrace();
+            throw new Error("Failed to access CodeModifier class.");
+        }
+
+        asmClassLoader = newChild(new URL[0]);
     }
 
     private static URL[] loadURLs() {
@@ -48,17 +94,27 @@ public class MinestomOverwriteClassLoader extends URLClassLoader {
 
         try {
             Class<?> systemClass = ClassLoader.getPlatformClassLoader().loadClass(name);
+            log.trace("System class: "+systemClass);
             return systemClass;
         } catch (ClassNotFoundException e) {
             try {
+                // check if parent already loaded the class
+                Class<?> loadedByParent = (Class<?>) findParentLoadedClass.invoke(getParent(), name);
+                if(loadedByParent != null) {
+                    log.trace("Already found in parent: "+loadedByParent);
+                    return super.loadClass(name, resolve);
+                }
+
+                if(isProtected(name)) {
+                    log.trace("Protected: "+name);
+                    return super.loadClass(name, resolve);
+                }
+
                 String path = name.replace(".", "/") + ".class";
                 byte[] bytes = getResourceAsStream(path).readAllBytes();
-                Class<?> defined = defineClass(name, bytes, 0, bytes.length);
+                return transformAndLoad(name, bytes, resolve);
-                if(resolve) {
+            } catch (Exception ex) {
-                    resolveClass(defined);
+                log.trace("Fail to load class, resorting to parent loader: "+name, ex);
-                }
-                return defined;
-            } catch (Exception ioException) {
                 // fail to load class, let parent load
                 // this forbids code modification, but at least it will load
                 return super.loadClass(name, resolve);
@@ -66,7 +122,63 @@ public class MinestomOverwriteClassLoader extends URLClassLoader {
         }
     }
 
-    public void loadModifier(File originFile, String codeModifierClass) {
+    private boolean isProtected(String name) {
-        throw new UnsupportedOperationException("TODO");
+        return protectedClasses.contains(name);
+    }
+
+    private Class<?> transformAndLoad(String name, byte[] bytes, boolean resolve) throws ClassNotFoundException {
+        ClassReader reader = new ClassReader(bytes);
+        ClassNode node = new ClassNode();
+        reader.accept(node, 0);
+        boolean modified = false;
+        synchronized (modifiers) {
+            for(CodeModifier modifier : modifiers) {
+                boolean shouldModify = modifier.getNamespace() == null || name.startsWith(modifier.getNamespace());
+                if(shouldModify) {
+                    modified |= modifier.transform(node);
+                }
+            }
+        }
+        if(modified) {
+            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES) {
+                @Override
+                protected ClassLoader getClassLoader() {
+                    return asmClassLoader;
+                }
+            };
+            node.accept(writer);
+            bytes = writer.toByteArray();
+        }
+        Class<?> defined = defineClass(name, bytes, 0, bytes.length);
+        log.trace("Loaded with code modifiers: "+name);
+        if(resolve) {
+            resolveClass(defined);
+        }
+        return defined;
+    }
+
+    @NotNull
+    public URLClassLoader newChild(@NotNull URL[] urls) {
+        return URLClassLoader.newInstance(urls, this);
+    }
+
+    public void loadModifier(File[] originFiles, String codeModifierClass) {
+        URL[] urls = new URL[originFiles.length];
+        try {
+            for (int i = 0; i < originFiles.length; i++) {
+                urls[i] = originFiles[i].toURI().toURL();
+            }
+            URLClassLoader loader = newChild(urls);
+            Class<?> modifierClass = loader.loadClass(codeModifierClass);
+            if(CodeModifier.class.isAssignableFrom(modifierClass)) {
+                CodeModifier modifier = (CodeModifier) modifierClass.getDeclaredConstructor().newInstance();
+                synchronized (modifiers) {
+                    log.warn("Added Code modifier: "+modifier);
+                    modifiers.add(modifier);
+                }
+            }
+        } catch (MalformedURLException | ClassNotFoundException | InvocationTargetException | InstantiationException | IllegalAccessException | NoSuchMethodException e) {
+            e.printStackTrace();
+        }
     }
 }

src/test/java/testextension/TestExtension.java

@@ -0,0 +1,15 @@
+package testextension;
+
+import net.minestom.server.extensions.Extension;
+
+public class TestExtension extends Extension {
+    @Override
+    public void initialize() {
+        System.out.println("Hello from extension!");
+    }
+
+    @Override
+    public void terminate() {
+
+    }
+}

src/test/java/testextension/TestExtensionLauncher.java

@@ -0,0 +1,13 @@
+package testextension;
+
+import net.minestom.server.Bootstrap;
+
+// To launch with VM arguments:
+// -Dminestom.extension.indevfolder.classes=build/classes/java/test/ -Dminestom.extension.indevfolder.resources=build/resources/test/
+public class TestExtensionLauncher {
+
+    public static void main(String[] args) {
+        Bootstrap.bootstrap("fr.themode.demo.MainDemo", args);
+    }
+
+}

src/test/java/testextension/TestModifier.java

@@ -0,0 +1,37 @@
+package testextension;
+
+import net.minestom.server.extras.selfmodification.CodeModifier;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.*;
+
+import java.util.List;
+
+public class TestModifier extends CodeModifier implements Opcodes {
+    @Override
+    public boolean transform(ClassNode source) {
+        if(source.name.equals("net/minestom/server/instance/InstanceContainer")) {
+            System.out.println("Modifying code source of "+source.name);
+            MethodNode constructor = findConstructor(source.methods);
+            constructor.instructions.insert(constructor.instructions.getFirst(), buildInjectionCode());
+            return true;
+        }
+        return false;
+    }
+
+    private InsnList buildInjectionCode() {
+        InsnList list = new InsnList();
+        list.add(new FieldInsnNode(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;"));
+        list.add(new LdcInsnNode("Hello from modified code!!"));
+        list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false));
+        return list;
+    }
+
+    private MethodNode findConstructor(List<MethodNode> methods) {
+        return methods.stream().filter(m -> m.name.equals("<init>")).findFirst().orElseThrow();
+    }
+
+    @Override
+    public String getNamespace() {
+        return "net.minestom.server";
+    }
+}

src/test/resources/extension.json

@@ -0,0 +1,7 @@
+{
+  "entrypoint": "testextension.TestExtension",
+  "name": "Test extension",
+  "codeModifiers": [
+    "testextension.TestModifier"
+  ]
+}