feat: java command detects agent lib automatically

teletha · Jan 29, 2024 · b51d1ed · b51d1ed
1 parent 222c51a
commit b51d1ed
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/src/main/java/bee/Bee.java b/src/main/java/bee/Bee.java
@@ -302,7 +302,7 @@ public void execute() {
      * @param tasks A list of task commands
      */
     public static void main(String... tasks) {
-        if (tasks.length == 0) tasks = new String[] {"bun"};
+        if (tasks.length == 0) tasks = new String[] {"test"};
 
         System.exit(new Bee().execute(BeeOption.parse(tasks)));
     }

diff --git a/src/main/java/bee/util/Java.java b/src/main/java/bee/util/Java.java
@@ -13,6 +13,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileNotFoundException;
+import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.PrintWriter;
@@ -25,6 +26,9 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
+import java.util.jar.Attributes;
+import java.util.jar.JarFile;
+import java.util.jar.Manifest;
 import java.util.stream.Collectors;
 
 import javax.management.MBeanServer;
@@ -239,6 +243,20 @@ public void run(Class<? extends JVM> mainClass, Object... arguments) {
         if (classpaths.size() != 0) {
             command.add("-cp");
             command.add(classpaths.stream().map(Path::toString).collect(Collectors.joining(File.pathSeparator)));
+
+            I.signal(classpaths).map(Locator::file).take(psychopath.File::isFile).take(file -> {
+                try (JarFile jar = new JarFile(file.asJavaFile())) {
+                    Manifest manifest = jar.getManifest();
+                    if (manifest != null) {
+                        Attributes main = manifest.getMainAttributes();
+                        String clazz = main.getValue("Premain-Class");
+                        return clazz != null;
+                    }
+                } catch (IOException e) {
+                    // ignore
+                }
+                return false;
+            }).to(x -> command.add("-javaagent:\"" + x.path() + '"'));
         }
 
         if (enableAssertion) {