Added ClassTransformers back!

Author: RealMangorage

src/main/java/module-info.java

@@ -2,5 +2,8 @@
     requires jdk.unsupported;
     requires java.scripting;
 
-    exports org.mangorage.bootstrap;
+    opens org.mangorage.bootstrap;
+    exports org.mangorage.bootstrap.api;
+
+    uses org.mangorage.bootstrap.api.IClassTransformer;
 }

src/main/java/org/mangorage/bootstrap/Bootstrap.java

@@ -46,10 +46,10 @@ public static void main(String[] args) throws IOException {
         final var moduleLayerController = ModuleLayer.defineModules(moduleCfg, List.of(ModuleLayer.boot()), s -> moduleCl);
         final var moduleLayer = moduleLayerController.layer();
 
-
-        System.out.println("Imagiine we worked!");
         Thread.currentThread().setContextClassLoader(moduleCl);
 
+        moduleCl.loadTransformers();
+
         callMain("org.mangorage.entrypoint.MangoBotCore", args, moduleLayer.findModule("org.mangorage.mangobotcore").get());
     }
 }

src/main/java/org/mangorage/bootstrap/api/IClassTransformer.java

@@ -0,0 +1,7 @@
+package org.mangorage.bootstrap.api;
+
+public interface IClassTransformer {
+    TransformResult transform(byte[] classData);
+
+    String getName();
+}

src/main/java/org/mangorage/bootstrap/api/TransformResult.java

@@ -0,0 +1,3 @@
+package org.mangorage.bootstrap.api;
+
+public record TransformResult(byte[] classData, TransformerFlag flag) {}

src/main/java/org/mangorage/bootstrap/api/TransformerFlag.java

@@ -0,0 +1,11 @@
+package org.mangorage.bootstrap.api;
+
+public enum TransformerFlag {
+    NO_REWRITE,
+    SIMPLE_REWRITE,
+    FULL_REWRITE;
+
+    public TransformResult of(byte[] classData) {
+        return new TransformResult(classData, this);
+    }
+}

src/main/java/org/mangorage/bootstrap/internal/ClassTransformers.java

@@ -0,0 +1,78 @@
+package org.mangorage.bootstrap.internal;
+
+import org.mangorage.bootstrap.api.IClassTransformer;
+import org.mangorage.bootstrap.api.TransformResult;
+import org.mangorage.bootstrap.api.TransformerFlag;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+
+public final class ClassTransformers {
+    private final Map<String, Class<?>> classes = new HashMap<>(); // Transformed Class's
+    private final List<IClassTransformer> transformers = new ArrayList<>(); // Transformer's
+    private final ClassLoader loader;
+
+    public ClassTransformers(ClassLoader loader) {
+        this.loader = loader;
+    }
+
+    public void add(String name, Class<?> clz) {
+        classes.put(name, clz);
+    }
+
+    public void add(IClassTransformer transformer) {
+        transformers.add(transformer);
+    }
+
+    public boolean isEmpty() {
+        return transformers.isEmpty();
+    }
+
+    private byte[] getClassBytes(String clazz) {
+        try {
+            String className = clazz.replace('.', '/');
+            String classFileName = className + ".class";
+
+            try (var is = loader.getResourceAsStream(classFileName)) {
+                if (is != null) return is.readAllBytes();
+            }
+        } catch (IOException e) {
+            return null;
+        }
+        return null;
+    }
+
+    public byte[] transform(String name) {
+        byte[] originalClassData = getClassBytes(name);
+
+        AtomicReference<TransformResult> result = new AtomicReference<>(TransformerFlag.NO_REWRITE.of(originalClassData));
+        AtomicReference<IClassTransformer> _transformer = new AtomicReference<>();
+
+        for (IClassTransformer transformer : transformers) {
+            result.set(transformer.transform(originalClassData));
+            if (result.get().flag() != TransformerFlag.NO_REWRITE) {
+                _transformer.set(transformer);
+                break;
+            }
+        }
+
+        if (result.get().flag() != TransformerFlag.NO_REWRITE && _transformer.get() != null) {
+            System.out.println("%s Transformed %s".formatted(_transformer.get().getName(), name));
+            return result.get().classData();
+        }
+
+        return null;
+    }
+
+    public boolean containsClass(String name) {
+        return classes.containsKey(name);
+    }
+
+    public Class<?> getClazz(String string) {
+        return classes.get(string);
+    }
+}

src/main/java/org/mangorage/bootstrap/internal/MangoLoader.java

@@ -1,5 +1,6 @@
 package org.mangorage.bootstrap.internal;
 
+import org.mangorage.bootstrap.api.IClassTransformer;
 
 import java.io.IOException;
 import java.lang.module.ModuleReader;
@@ -9,6 +10,7 @@
 import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.ServiceLoader;
 import java.util.Set;
 
 public final class MangoLoader extends URLClassLoader {
@@ -17,6 +19,7 @@ public final class MangoLoader extends URLClassLoader {
         ClassLoader.registerAsParallelCapable();
     }
 
+    private ClassTransformers transformers = new ClassTransformers(this);
     private final Map<String, LoadedModule> moduleMap = new HashMap<>();
     private final Map<String, LoadedModule> localPackageToModule = new HashMap<>();
 
@@ -41,6 +44,58 @@ public MangoLoader(URL[] urls, Set<ResolvedModule> modules, ClassLoader parent)
         });
     }
 
+    public void loadTransformers() {
+        ServiceLoader.load(IClassTransformer.class)
+                .stream()
+                .forEach(provider -> {
+                    transformers.add(provider.get());
+                });
+    }
+
+    @Override
+    protected Class<?> findClass(String name) throws ClassNotFoundException {
+        if (transformers == null || transformers.isEmpty())
+            return super.findClass(name);
+
+        if (transformers.containsClass(name))
+            return transformers.getClazz(name);
+
+        byte[] originalBytes = getClassBytes(name);
+
+        if (originalBytes == null) {
+            throw new ClassNotFoundException("Failed to load original class bytes for " + name);
+        }
+
+        byte[] arr = transformers.transform(name);
+        if (arr != null) {
+            Class<?> clz = defineClass(name, arr);
+            transformers.add(name, clz);
+            return clz;
+        }
+
+        return super.findClass(name);
+    }
+
+    private byte[] getClassBytes(String clazz) {
+        try {
+            String className = clazz.replace('.', '/');
+            String classFileName = className + ".class";
+
+            try (var is = getResourceAsStream(classFileName)) {
+                if (is != null) return is.readAllBytes();
+            }
+        } catch (IOException e) {
+            return null;
+        }
+        return null;
+    }
+
+    private Class<?> defineClass(String name, byte[] bytes) {
+        return super.defineClass(name, bytes, 0, bytes.length);
+    }
+
+
+
     @Override
     protected URL findResource(String moduleName, String name) throws IOException {
         final var loadedModule = moduleMap.get(moduleName);
@@ -58,8 +113,11 @@ protected URL findResource(String moduleName, String name) throws IOException {
     protected Class<?> findClass(String moduleName, String name) {
         Class<?> c = null;
         LoadedModule loadedModule = findLoadedModule(name);
-        if (loadedModule != null && loadedModule.getModuleReference().descriptor().name().equals(moduleName))
+        if (loadedModule != null && loadedModule.getModuleReference().descriptor().name().equals(moduleName)) {
             c = defineClass(name, loadedModule);
+        } else if (loadedModule != null) {
+            throw new IllegalArgumentException("Expected Class '%s' in module '%s', instead was in '%s'".formatted(name, moduleName, loadedModule.getModuleReference().descriptor().name()));
+        }
         return c;
     }