Fix Agent and WideningTransformer and ClassPatcher

AccessWidener/src/main/java/de/steamwar/WideningTransformer.java

@@ -19,17 +19,9 @@
 
 package de.steamwar;
 
-import org.objectweb.asm.ClassReader;
-import org.objectweb.asm.ClassWriter;
-
 import java.lang.instrument.ClassFileTransformer;
-import java.lang.instrument.Instrumentation;
-import java.lang.instrument.UnmodifiableClassException;
 import java.security.ProtectionDomain;
-import java.util.*;
-import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.logging.Logger;
-import java.util.stream.Collectors;
+import java.util.List;
 
 /**
  * Transforms class bytecode to apply access widening rules.
@@ -39,111 +31,19 @@ import java.util.stream.Collectors;
  */
 public class WideningTransformer implements ClassFileTransformer {
 
-    private static final Logger LOG = Logger.getLogger("WidenerTransformer");
+    private final ClassPatcher patcher;
 
-    private final Instrumentation instrumentation;
-
-    /**
-     * All entries collected across all plugins. Thread-safe for concurrent plugin loads.
-     */
-    private final List<AccessWidenerEntry> entries = new CopyOnWriteArrayList<>();
-
-    /**
-     * ClassLoaders we have already scanned, to avoid re-scanning.
-     */
-    private final Set<ClassLoader> scannedLoaders = Collections.synchronizedSet(Collections.newSetFromMap(new IdentityHashMap<>()));
-
-    public WideningTransformer(Instrumentation instrumentation) {
-        this.instrumentation = instrumentation;
+    public WideningTransformer(List<AccessWidenerEntry> entries) {
+        patcher = new ClassPatcher(entries);
     }
 
-    /**
-     * Add entries — used during initial scan before the transformer is registered.
-     */
-    public void addEntries(List<AccessWidenerEntry> newEntries) {
-        entries.addAll(newEntries);
-    }
-
-    // -------------------------------------------------------------------------
-    // ClassFileTransformer
-    // -------------------------------------------------------------------------
-
     @Override
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
-
-        // Check for new plugin classloaders we haven't seen yet
-        if (loader != null && AccessWidenerScanner.isPluginClassLoader(loader) && scannedLoaders.add(loader)) {
-            onNewPluginClassLoader(loader);
-        }
-
-        // No entries for this class — skip transformation entirely
-        if (className == null) return classfileBuffer;
-        boolean relevant = entries.stream().anyMatch(e -> e.targets(className));
-        if (!relevant) return classfileBuffer;
-
-        // Apply widening via ASM
-        try {
-            ClassReader cr = new ClassReader(classfileBuffer);
-            ClassWriter cw = new ClassWriter(cr, 0);
-            cr.accept(new ClassTransformer(cw, className, entries), 0);
-            return cw.toByteArray();
-        } catch (Exception e) {
-            LOG.warning("[AccessWidener] Failed to transform " + className + ": " + e.getMessage());
+        byte[] result = patcher.patch(className, classfileBuffer);
+        if (result == null) {
             return classfileBuffer;
-        }
-    }
-
-    // -------------------------------------------------------------------------
-    // New ClassLoader detection
-    // -------------------------------------------------------------------------
-
-    private void onNewPluginClassLoader(ClassLoader loader) {
-        List<AccessWidenerEntry> newEntries = AccessWidenerScanner.scanClassLoader(loader);
-        if (newEntries.isEmpty()) return;
-
-        entries.addAll(newEntries);
-        LOG.info("[AccessWidener] Picked up " + newEntries.size() + " new entr" + (newEntries.size() == 1 ? "y" : "ies") + " from " + loader);
-
-        // Retransform already-loaded classes that are now targeted
-        scheduleRetransform(newEntries);
-    }
-
-    /**
-     * Retransformation cannot be called from within a transform() call,
-     * so we dispatch it to a short-lived daemon thread.
-     */
-    private void scheduleRetransform(List<AccessWidenerEntry> newEntries) {
-        Set<String> targets = newEntries.stream().map(e -> e.target().replace('/', '.')).collect(Collectors.toSet());
-
-        Thread t = new Thread(() -> retransform(targets), "access-widener-retransform");
-        t.setDaemon(true);
-        t.start();
-    }
-
-    private void retransform(Set<String> dotNames) {
-        List<Class> toRetransform = Arrays.stream(instrumentation.getAllLoadedClasses())
-                .filter(c -> dotNames.contains(c.getName()))
-                .filter(instrumentation::isModifiableClass)
-                .collect(Collectors.toList());
-
-        if (toRetransform.isEmpty()) return;
-
-        LOG.info("[AccessWidener] Retransforming " + toRetransform.size() + " class(es).");
-        List<String> failed = new ArrayList<>();
-
-        for (Class<?> clazz : toRetransform) {
-            try {
-                instrumentation.retransformClasses(clazz);
-            } catch (UnmodifiableClassException | UnsupportedOperationException e) {
-                failed.add(clazz.getName());
-            }
-        }
-
-        if (!failed.isEmpty()) {
-            LOG.warning("[AccessWidener] The following classes were already loaded before the agent"
-                    + " and cannot be retransformed. Add -javaagent: to your start script"
-                    + " to fix this:");
-            failed.forEach(name -> LOG.warning("  - " + name));
+        } else {
+            return result;
         }
     }
 }