new feature: management transformer

Signed-off-by: pengzhile <pengzhile@gmail.com>

README.md

@@ -1,4 +1,4 @@
-# ja-netfilter v2.2.3
+# ja-netfilter v2.3.0
 
 ### A javaagent framework
 

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>com.ja-netfilter</groupId>
     <artifactId>ja-netfilter</artifactId>
-    <version>2.2.3</version>
+    <version>2.3.0</version>
 
     <name>ja-netfilter</name>
     <description>A javaagent framework</description>

src/main/java/com/janetfilter/core/Dispatcher.java

@@ -13,12 +13,17 @@ public final class Dispatcher implements ClassFileTransformer {
     private final Set<String> classSet = new TreeSet<>();
     private final Map<String, List<MyTransformer>> transformerMap = new HashMap<>();
     private final List<MyTransformer> globalTransformers = new ArrayList<>();
+    private final List<MyTransformer> manageTransformers = new ArrayList<>();
 
     public Dispatcher(Environment environment) {
         this.environment = environment;
     }
 
-    public synchronized void addTransformer(MyTransformer transformer) {
+    public void addTransformer(MyTransformer transformer) {
+        if (null == transformer) {
+            return;
+        }
+
         if (environment.isAttachMode() && !transformer.attachMode()) {
             DebugInfo.debug("Transformer: " + transformer.getClass().getName() + " is set to not load in attach mode, ignored.");
             return;
@@ -29,16 +34,23 @@ public final class Dispatcher implements ClassFileTransformer {
             return;
         }
 
-        String className = transformer.getHookClassName();
-        if (null == className) {
-            globalTransformers.add(transformer);
-            return;
+        synchronized (this) {
+            String className = transformer.getHookClassName();
+            if (null == className) {
+                globalTransformers.add(transformer);
+
+                if (transformer.isManager()) {
+                    manageTransformers.add(transformer);
+                }
+
+                return;
+            }
+
+            classSet.add(className.replace('/', '.'));
+            List<MyTransformer> transformers = transformerMap.computeIfAbsent(className, k -> new ArrayList<>());
+
+            transformers.add(transformer);
         }
-
-        classSet.add(className.replace('/', '.'));
-        List<MyTransformer> transformers = transformerMap.computeIfAbsent(className, k -> new ArrayList<>());
-
-        transformers.add(transformer);
     }
 
     public void addTransformers(List<MyTransformer> transformers) {
@@ -64,42 +76,40 @@ public final class Dispatcher implements ClassFileTransformer {
     }
 
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) throws IllegalClassFormatException {
-        do {
-            if (null == className) {
-                break;
+        if (null == className) {
+            return classFileBuffer;
+        }
+
+        List<MyTransformer> transformers = transformerMap.get(className);
+        List<MyTransformer> globalTransformers = null == transformers ? this.manageTransformers : this.globalTransformers;
+
+        int order = 0;
+
+        try {
+            for (MyTransformer transformer : globalTransformers) {
+                transformer.before(className, classFileBuffer);
             }
 
-            List<MyTransformer> transformers = transformerMap.get(className);
-            if (null == transformers) {
-                break;
+            for (MyTransformer transformer : globalTransformers) {
+                classFileBuffer = transformer.preTransform(className, classFileBuffer, order++);
             }
 
-            int order = 0;
-
-            try {
-                for (MyTransformer transformer : globalTransformers) {
-                    transformer.before(className, classFileBuffer);
-                }
-
-                for (MyTransformer transformer : globalTransformers) {
-                    classFileBuffer = transformer.preTransform(className, classFileBuffer, order++);
-                }
-
+            if (null != transformers) {
                 for (MyTransformer transformer : transformers) {
                     classFileBuffer = transformer.transform(className, classFileBuffer, order++);
                 }
-
-                for (MyTransformer transformer : globalTransformers) {
-                    classFileBuffer = transformer.postTransform(className, classFileBuffer, order++);
-                }
-
-                for (MyTransformer transformer : globalTransformers) {
-                    transformer.after(className, classFileBuffer);
-                }
-            } catch (Throwable e) {
-                DebugInfo.error("Transform class failed: " + className, e);
             }
-        } while (false);
+
+            for (MyTransformer transformer : globalTransformers) {
+                classFileBuffer = transformer.postTransform(className, classFileBuffer, order++);
+            }
+
+            for (MyTransformer transformer : globalTransformers) {
+                transformer.after(className, classFileBuffer);
+            }
+        } catch (Throwable e) {
+            DebugInfo.error("Transform class failed: " + className, e);
+        }
 
         return classFileBuffer;
     }

src/main/java/com/janetfilter/core/Environment.java

@@ -4,6 +4,7 @@ import com.janetfilter.core.utils.ProcessUtils;
 import com.janetfilter.core.utils.StringUtils;
 
 import java.io.File;
+import java.lang.instrument.Instrumentation;
 
 public final class Environment {
     private final String pid;
@@ -18,11 +19,14 @@ public final class Environment {
     private final String disabledPluginSuffix;
     private final boolean attachMode;
 
-    public Environment(File agentFile, boolean attachMode) {
-        this(agentFile, null, attachMode);
+    private final Instrumentation instrumentation;
+
+    public Environment(Instrumentation instrumentation, File agentFile, boolean attachMode) {
+        this(instrumentation, agentFile, null, attachMode);
     }
 
-    public Environment(File agentFile, String app, boolean attachMode) {
+    public Environment(Instrumentation instrumentation, File agentFile, String app, boolean attachMode) {
+        this.instrumentation = instrumentation;
         this.agentFile = agentFile;
         baseDir = agentFile.getParentFile();
 
@@ -93,6 +97,10 @@ public final class Environment {
         return !attachMode;
     }
 
+    public Instrumentation getInstrumentation() {
+        return instrumentation;
+    }
+
     @Override
     public String toString() {
         return "Environment: {" +

src/main/java/com/janetfilter/core/Initializer.java

@@ -7,13 +7,14 @@ import java.lang.instrument.Instrumentation;
 import java.util.Set;
 
 public class Initializer {
-    public static void init(Instrumentation inst, Environment environment) {
+    public static void init(Environment environment) {
         DebugInfo.useFile(environment.getLogsDir());
         DebugInfo.info(environment.toString());
 
         Dispatcher dispatcher = new Dispatcher(environment);
-        new PluginManager(inst, dispatcher, environment).loadPlugins();
+        new PluginManager(dispatcher, environment).loadPlugins();
 
+        Instrumentation inst = environment.getInstrumentation();
         inst.addTransformer(dispatcher, true);
         inst.setNativeMethodPrefix(dispatcher, environment.getNativePrefix());
 

src/main/java/com/janetfilter/core/Launcher.java

@@ -12,7 +12,7 @@ import java.util.jar.JarFile;
 
 public class Launcher {
     public static final String ATTACH_ARG = "--attach";
-    public static final String VERSION = "v2.2.3";
+    public static final String VERSION = "v2.3.0";
 
     private static boolean loaded = false;
 
@@ -78,7 +78,7 @@ public class Launcher {
             return;
         }
 
-        Initializer.init(inst, new Environment(agentFile, args, attachMode)); // for some custom UrlLoaders
+        Initializer.init(new Environment(inst, agentFile, args, attachMode)); // for some custom UrlLoaders
     }
 
     private static void printUsage() {

src/main/java/com/janetfilter/core/plugin/MyTransformer.java

@@ -20,6 +20,16 @@ public interface MyTransformer {
         return true;
     }
 
+    /**
+     * for global transformers only
+     * whether it is a management transformer
+     *
+     * @return return true to handle the transform of all classes
+     */
+    default boolean isManager() {
+        return false;
+    }
+
     /**
      * for global transformers only
      */

src/main/java/com/janetfilter/core/plugin/PluginManager.java

@@ -22,8 +22,8 @@ public final class PluginManager {
     private final Dispatcher dispatcher;
     private final Environment environment;
 
-    public PluginManager(Instrumentation inst, Dispatcher dispatcher, Environment environment) {
-        this.inst = inst;
+    public PluginManager(Dispatcher dispatcher, Environment environment) {
+        this.inst = environment.getInstrumentation();
         this.dispatcher = dispatcher;
         this.environment = environment;
     }