From 6da27bdf9926b08a6e949f212e57c24b9a50a44d Mon Sep 17 00:00:00 2001
From: Gravita <gravita@gravit.pro>
Date: Sat, 20 Mar 2021 00:57:24 +0700
Subject: [PATCH] [FEATURE] ServerWrapperSetup Agent

---
 .../pro/gravit/launcher/server/ServerAgent.java   |  5 ++---
 .../pro/gravit/launcher/server/ServerWrapper.java |  2 --
 .../launcher/server/setup/ServerWrapperSetup.java | 15 +++++++++++++--
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerAgent.java b/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerAgent.java
index f5756b73..00333325 100644
--- a/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerAgent.java
+++ b/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerAgent.java
@@ -17,7 +17,6 @@
 
 public class ServerAgent {
     public static final Boolean isAutoloadLibraries = Boolean.getBoolean(System.getProperty("serverwrapper,agentlibrariesload", "false"));
-    public static final Boolean isAgentProxy = Boolean.getBoolean(System.getProperty("serverwrapper,agentproxy", "false"));
     public static Instrumentation inst = null;
     private static boolean isAgentStarted = false;
 
@@ -47,8 +46,8 @@ public static void premain(String agentArgument, Instrumentation instrumentation
             Path libraries = Paths.get("libraries");
             if (IOHelper.exists(libraries)) loadLibraries(libraries);
         }
-        if (isAgentProxy) {
-            String proxyClassName = System.getProperty("serverwrapper.agentproxyclass");
+        String proxyClassName = System.getProperty("serverwrapper.agentproxy", null);
+        if (proxyClassName != null) {
             Class<?> proxyClass;
             try {
                 proxyClass = Class.forName(proxyClassName);
diff --git a/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerWrapper.java b/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerWrapper.java
index e590f7bf..f783aaf5 100644
--- a/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerWrapper.java
+++ b/ServerWrapper/src/main/java/pro/gravit/launcher/server/ServerWrapper.java
@@ -249,7 +249,6 @@ public Config getDefaultConfig() {
         newConfig.stopOnError = true;
         newConfig.reconnectCount = 10;
         newConfig.reconnectSleep = 1000;
-        newConfig.websocket = new WebSocketConf();
         newConfig.address = "ws://localhost:9274/api";
         newConfig.env = LauncherConfig.LauncherEnvironment.STD;
         return newConfig;
@@ -261,7 +260,6 @@ public static final class Config {
         public String projectname;
         public String address;
         public String serverName;
-        public WebSocketConf websocket;
         public int reconnectCount;
         public int reconnectSleep;
         public boolean customClassPath;
diff --git a/ServerWrapper/src/main/java/pro/gravit/launcher/server/setup/ServerWrapperSetup.java b/ServerWrapper/src/main/java/pro/gravit/launcher/server/setup/ServerWrapperSetup.java
index 9bcd7118..d4526e96 100644
--- a/ServerWrapper/src/main/java/pro/gravit/launcher/server/setup/ServerWrapperSetup.java
+++ b/ServerWrapper/src/main/java/pro/gravit/launcher/server/setup/ServerWrapperSetup.java
@@ -29,11 +29,13 @@ public void run() throws IOException {
         String jarName = commands.commandHandler.readLine();
         Path jarPath = Paths.get(jarName);
         String mainClassName;
+        String agentClassName;
         try (JarFile file = new JarFile(jarPath.toFile())) {
             URL jarURL = jarPath.toUri().toURL();
             urlClassLoader = new PublicURLClassLoader(new URL[]{jarURL});
             LogHelper.info("Check server jar MainClass");
             mainClassName = file.getManifest().getMainAttributes().getValue("Main-Class");
+            agentClassName = file.getManifest().getMainAttributes().getValue("Premain-Class");
             if (mainClassName == null) {
                 LogHelper.error("Main-Class not found in MANIFEST");
                 return;
@@ -46,6 +48,9 @@ public void run() throws IOException {
             }
         }
         LogHelper.info("Found MainClass %s", mainClassName);
+        if(agentClassName != null) {
+            LogHelper.info("Found PremainClass %s", agentClassName);
+        }
         System.out.println("Print your server name:");
         wrapper.config.serverName = commands.commandHandler.readLine();
         System.out.println("Print launchserver websocket host( ws://host:port/api ):");
@@ -82,14 +87,20 @@ public void run() throws IOException {
         ServerWrapper.modulesManager.invokeEvent(new ServerWrapperSetupEvent(this));
         try (Writer writer = IOHelper.newWriter(startScript)) {
             if (JVMHelper.OS_TYPE == JVMHelper.OS.LINUX) {
-                writer.append("#!/bin/sh\n\n");
+                writer.append("#!/bin/bash\n\n");
             }
-            writer.append("java ");
+            writer.append(IOHelper.resolveJavaBin(Paths.get(System.getProperty("java.home"))).toAbsolutePath().toString());
+            writer.append(" ");
             if (mainClassName.contains("bungee")) {
                 LogHelper.info("Found BungeeCord mainclass. Modules dir change to modules_srv");
                 writer.append(JVMHelper.jvmProperty("serverwrapper.modulesDir", "modules_srv"));
                 writer.append(" ");
             }
+            if(agentClassName != null) {
+                writer.append("-javaagent:ServerWrapper.jar ");
+                writer.append("-Dserverwrapper.agentproxy=".concat(agentClassName));
+                writer.append(" ");
+            }
             //More args
             writer.append("-cp ");
             String pathServerWrapper = IOHelper.getCodeSource(ServerWrapper.class).getFileName().toString();