Invoke premain method

Author: acrylic-style

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>

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 {