Invoke premain method

AzisabaNetwork · Dec 12, 2021 · 90ab4d2 · 90ab4d2
1 parent 90e8e5c
commit 90ab4d2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/pom.xml b/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>net.azisaba</groupId>
     <artifactId>Log4j2Fix</artifactId>
-    <version>1.0.1</version>
+    <version>1.0.2</version>
 
     <properties>
         <maven.compiler.source>8</maven.compiler.source>

diff --git a/src/main/java/net/azisaba/log4j2Fix/Log4j2Fix.java b/src/main/java/net/azisaba/log4j2Fix/Log4j2Fix.java
@@ -17,6 +17,9 @@
 import java.util.zip.ZipFile;
 
 public class Log4j2Fix {
+    private static String arg;
+    private static Instrumentation inst;
+
     public static void main(String[] args) throws IOException {
         transformClasses();
         List<String> arguments = new ArrayList<>(Arrays.asList(args));
@@ -25,6 +28,7 @@ public static void main(String[] args) throws IOException {
             return;
         }
         String main = arguments.remove(0);
+        String premain = null;
         File file = new File(main);
         if (file.exists()) {
             System.out.println("Using " + file.getAbsolutePath() + " for classpath");
@@ -45,7 +49,8 @@ public static void main(String[] args) throws IOException {
                     if (read.startsWith("Main-Class: ")) {
                         main = read.replace("Main-Class: ", "");
                         found = true;
-                        break;
+                    } else if (read.startsWith("Premain-Class")) {
+                        premain = read.replace("Premain-Class: " , "");
                     }
                 }
                 reader.close();
@@ -58,6 +63,16 @@ public static void main(String[] args) throws IOException {
                 }
             }
         }
+        if (premain != null) {
+            try {
+                Class<?> clazz = Class.forName(premain);
+                Method m = clazz.getMethod("premain", String.class, Instrumentation.class);
+                m.invoke(null, arg, inst);
+            } catch (ReflectiveOperationException e) {
+                System.err.println("Failed to invoke premain method of class " + premain);
+                e.printStackTrace();
+            }
+        }
         try {
             Class<?> clazz = Class.forName(main);
             Method m = clazz.getMethod("main", String[].class);
@@ -70,10 +85,14 @@ public static void main(String[] args) throws IOException {
 
     public static void agentmain(String args, Instrumentation instrumentation) throws IOException {
         transformClasses();
+        arg = args;
+        inst = instrumentation;
     }
 
     public static void premain(String args, Instrumentation instrumentation) throws IOException {
         transformClasses();
+        arg = args;
+        inst = instrumentation;
     }
 
     public static void transformClasses() throws IOException {