reFix: optimizate agent attach performance

dongtai-agent/pom.xml

@@ -107,13 +107,6 @@
                             <goal>shade</goal>
                         </goals>
                         <configuration>
-                            <relocations>
-                                <relocation>
-                                    <pattern>org.apache.commons</pattern>
-                                    <shadedPattern>${shade-prefix}.org.apache.commons</shadedPattern>
-                                </relocation>
-                            </relocations>
-
                             <filters>
                                 <filter>
                                     <artifact>*:*</artifact>
@@ -131,6 +124,16 @@
                                     </excludes>
                                 </filter>
                             </filters>
+                            <relocations>
+                                <relocation>
+                                    <pattern>org.</pattern>
+                                    <shadedPattern>${shade-prefix}.org.</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>oshi.</pattern>
+                                    <shadedPattern>${shade-prefix}.oshi.</shadedPattern>
+                                </relocation>
+                            </relocations>
                         </configuration>
                     </execution>
                 </executions>

dongtai-agent/src/main/java/com/secnium/iast/agent/Agent.java

@@ -7,7 +7,7 @@
 import java.net.URL;
 import java.net.URLClassLoader;
 
-import com.secnium.iast.log.DongTaiLog;
+import io.dongtai.log.DongTaiLog;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.CommandLineParser;
 import org.apache.commons.cli.DefaultParser;
@@ -31,6 +31,10 @@ public static void main(String[] args) {
 
         attachOptions.addOption(build("p", "pid", "webserver process id"));
         attachOptions.addOption(build("m", "mode", "optional: install uninstall"));
+        attachOptions.addOption(build("debug", "debug", "optional: debug mode"));
+        attachOptions.addOption(build("app_name", "app_name", "optional: DongTai Application Name, default: ExampleApplication"));
+        attachOptions.addOption(build("app_create", "app_create", "optional: DongTai Application Auto Create, default: false"));
+        attachOptions.addOption(build("app_version", "app_version", "optional: DongTai Application Version, default: v1.0"));
 
         CommandLineParser parser = new DefaultParser();
         HelpFormatter formatter = new HelpFormatter();
@@ -41,15 +45,27 @@ public static void main(String[] args) {
             if (result.hasOption("p") && result.hasOption("m")) {
                 String pid = result.getOptionValue("p");
                 String mode = result.getOptionValue("m");
-                String attachArgs = null;
-                attachArgs = mode;
+                StringBuilder attachArgs = new StringBuilder();
+                attachArgs.append("mode=").append(mode);
+                if (result.hasOption("debug")) {
+                    attachArgs.append("&debug=").append(result.getOptionValue("debug"));
+                }
+                if (result.hasOption("app_create")) {
+                    attachArgs.append("&appCreate=").append(result.getOptionValue("app_create"));
+                }
+                if (result.hasOption("app_name")) {
+                    attachArgs.append("&appName=").append(result.getOptionValue("app_name"));
+                }
+                if (result.hasOption("app_version")) {
+                    attachArgs.append("&appVersion=").append(result.getOptionValue("app_version"));
+                }
 
                 String jdkVersion = getJdkVersion();
                 if ("1".equals(jdkVersion) && appendToolsPath()) {
-                    AttachLauncher.attach(pid, attachArgs);
+                    AttachLauncher.attach(pid, attachArgs.toString());
                     DongTaiLog.info("engine " + attachArgs + " successfully. pid: " + pid);
                 } else {
-                    AttachLauncher.attach(pid, attachArgs);
+                    AttachLauncher.attach(pid, attachArgs.toString());
                     DongTaiLog.info("engine " + attachArgs + " successfully. pid: " + pid);
                 }
             } else {

dongtai-agent/src/main/java/com/secnium/iast/agent/AgentLauncher.java

@@ -4,17 +4,20 @@
 import com.secnium.iast.agent.monitor.EngineMonitor;
 import com.secnium.iast.agent.monitor.MonitorDaemonThread;
 import com.secnium.iast.agent.report.AgentRegisterReport;
-import com.secnium.iast.log.DongTaiLog;
+import io.dongtai.log.DongTaiLog;
+
 import java.lang.instrument.Instrumentation;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * @author dongzhiyong@huoxian.cn
  */
 public class AgentLauncher {
 
-    private static final String LAUNCH_MODE_AGENT = "agent";
-    private static final String LAUNCH_MODE_ATTACH = "attach";
-    private static String LAUNCH_MODE;
+    public static final String LAUNCH_MODE_AGENT = "agent";
+    public static final String LAUNCH_MODE_ATTACH = "attach";
+    public static String LAUNCH_MODE;
 
     /**
      * install agent with premain
@@ -30,7 +33,6 @@ public static void premain(String args, Instrumentation inst) {
         LAUNCH_MODE = LAUNCH_MODE_AGENT;
         try {
             install(inst);
-            System.setProperty("protect.by.dongtai", "1");
         } catch (Exception e) {
             e.printStackTrace();
         }
@@ -43,8 +45,9 @@ public static void premain(String args, Instrumentation inst) {
      * @param inst inst
      */
     public static void agentmain(String args, Instrumentation inst) {
-        System.out.println(System.getProperty("protect.by.dongtai", null));
-        if ("uninstall".equals(args)) {
+        DongTaiLog.info(System.getProperty("protect.by.dongtai", "Current Application Run Without DongTai"));
+        Map<String, String> argsMap = parseArgs(args);
+        if ("uninstall".equals(argsMap.get("mode"))) {
             if (System.getProperty("protect.by.dongtai", null) == null) {
                 DongTaiLog.info("DongTai wasn't installed.");
                 return;
@@ -59,9 +62,19 @@ public static void agentmain(String args, Instrumentation inst) {
             }
             LAUNCH_MODE = LAUNCH_MODE_ATTACH;
             try {
-                Agent.appendToolsPath();
+                if (argsMap.containsKey("debug")) {
+                    System.setProperty("dongtai.debug", argsMap.get("debug"));
+                }
+                if (argsMap.containsKey("appCreate")) {
+                    System.setProperty("dongtai.app.create", argsMap.get("appCreate"));
+                }
+                if (argsMap.containsKey("appName")) {
+                    System.setProperty("dongtai.app.name", argsMap.get("appName"));
+                }
+                if (argsMap.containsKey("appVersion")) {
+                    System.setProperty("dongtai.app.version", argsMap.get("appVersion"));
+                }
                 install(inst);
-                System.setProperty("protect.with.dongtai", "1");
             } catch (Exception e) {
                 e.printStackTrace();
             }
@@ -85,6 +98,7 @@ public static synchronized void uninstall() {
      */
     private static void install(final Instrumentation inst) {
         IastProperties iastProperties = IastProperties.getInstance();
+        DongTaiLog.info("try to register agent to: " + iastProperties.getBaseUrl());
         Boolean send = AgentRegisterReport.send();
         if (send) {
             DongTaiLog.info("Agent has successfully registered with " + iastProperties.getBaseUrl());
@@ -95,7 +109,9 @@ private static void install(final Instrumentation inst) {
             } else {
                 EngineMonitor.isCoreRegisterStart = true;
             }
+            Runtime.getRuntime().addShutdownHook(new ShutdownThread());
             loadEngine(inst);
+            System.setProperty("protect.by.dongtai", "1");
         } else {
             DongTaiLog.error("Agent register failed. Start without DongTai IAST.");
         }
@@ -114,6 +130,15 @@ private static void loadEngine(final Instrumentation inst) {
         agentMonitorDaemonThread.setPriority(1);
         agentMonitorDaemonThread.setName("dongtai-monitor");
         agentMonitorDaemonThread.start();
-        Runtime.getRuntime().addShutdownHook(new ShutdownThread());
+    }
+
+    private static Map<String, String> parseArgs(String args) {
+        Map<String, String> argsMap = new HashMap<String, String>();
+        String[] argsItems = args.split("&");
+        for (String argsItem : argsItems) {
+            String[] argItems = argsItem.split("=");
+            argsMap.put(argItems[0], argItems[1]);
+        }
+        return argsMap;
     }
 }

dongtai-agent/src/main/java/com/secnium/iast/agent/AttachLauncher.java

@@ -1,10 +1,11 @@
 package com.secnium.iast.agent;
 
 import com.secnium.iast.agent.util.JavaVersionUtils;
-import com.secnium.iast.log.DongTaiLog;
+import io.dongtai.log.DongTaiLog;
 import com.sun.tools.attach.AgentLoadException;
 import com.sun.tools.attach.AttachNotSupportedException;
 import com.sun.tools.attach.VirtualMachine;
+
 import java.io.IOException;
 import java.util.Properties;
 

dongtai-agent/src/main/java/com/secnium/iast/agent/IastProperties.java

@@ -1,6 +1,7 @@
 package com.secnium.iast.agent;
 
-import com.secnium.iast.log.DongTaiLog;
+import io.dongtai.log.DongTaiLog;
+
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.FileOutputStream;
@@ -87,8 +88,7 @@ public void init(String path) throws ClassNotFoundException {
             ByteArrayInputStream inputStream = new ByteArrayInputStream(data);
             cfg.load(inputStream);
 
-            DongTaiLog.info(
-                    "The engine configuration file is initialized successfully. file is " + propertiesFile.toString());
+            DongTaiLog.info("DongTai configuration has initialized successfully. config: " + propertiesFile.toString());
         } catch (IOException e) {
             e.printStackTrace();
         }

dongtai-agent/src/main/java/com/secnium/iast/agent/ShutdownThread.java

@@ -1,7 +1,7 @@
 package com.secnium.iast.agent;
 
 import com.secnium.iast.agent.monitor.MonitorDaemonThread;
-import com.secnium.iast.log.DongTaiLog;
+import io.dongtai.log.DongTaiLog;
 
 public class ShutdownThread extends Thread {
 

dongtai-agent/src/main/java/com/secnium/iast/agent/manager/EngineManager.java

@@ -5,13 +5,10 @@
 import com.secnium.iast.agent.IastProperties;
 import com.secnium.iast.agent.report.AgentRegisterReport;
 import com.secnium.iast.agent.util.http.HttpClientUtils;
-import com.secnium.iast.log.DongTaiLog;
-import java.io.BufferedInputStream;
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.PrintWriter;
-import java.io.StringWriter;
+import io.dongtai.log.DongTaiLog;
+import org.json.JSONObject;
+
+import java.io.*;
 import java.lang.instrument.Instrumentation;
 import java.lang.management.ManagementFactory;
 import java.lang.reflect.InvocationTargetException;
@@ -128,22 +125,34 @@ private boolean downloadJarPackageToCacheFromUrl(String fileUrl, String fileName
             connection.setUseCaches(false);
             connection.setDoOutput(true);
 
-            BufferedInputStream in = new BufferedInputStream(connection.getInputStream());
-            final File classPath = new File(new File(fileName).getParent());
+            if (connection.getContentType().equals("application/json")) {
+                BufferedReader streamReader = new BufferedReader(new InputStreamReader(connection.getInputStream(), "UTF-8"));
+                StringBuilder responseStrBuilder = new StringBuilder();
+                String inputStr;
+                while ((inputStr = streamReader.readLine()) != null)
+                    responseStrBuilder.append(inputStr);
 
-            if (!classPath.mkdirs() && !classPath.exists()) {
-                DongTaiLog.info("Check or create local file cache path, path is " + classPath);
-            }
-            FileOutputStream fileOutputStream = new FileOutputStream(fileName);
-            byte[] dataBuffer = new byte[1024];
-            int bytesRead;
-            while ((bytesRead = in.read(dataBuffer, 0, 1024)) != -1) {
-                fileOutputStream.write(dataBuffer, 0, bytesRead);
+                JSONObject jsonObject = new JSONObject(responseStrBuilder.toString());
+                DongTaiLog.error("DongTai Core Package: {} download failed. response: {}", fileUrl, jsonObject);
+                return false;
+            } else {
+                BufferedInputStream in = new BufferedInputStream(connection.getInputStream());
+                final File classPath = new File(new File(fileName).getParent());
+
+                if (!classPath.mkdirs() && !classPath.exists()) {
+                    DongTaiLog.info("Check or create local file cache path, path is " + classPath);
+                }
+                FileOutputStream fileOutputStream = new FileOutputStream(fileName);
+                byte[] dataBuffer = new byte[1024];
+                int bytesRead;
+                while ((bytesRead = in.read(dataBuffer, 0, 1024)) != -1) {
+                    fileOutputStream.write(dataBuffer, 0, bytesRead);
+                }
+                in.close();
+                fileOutputStream.close();
+                DongTaiLog.info("The remote file " + fileUrl + " was successfully written to the local cache.");
+                status = true;
             }
-            in.close();
-            fileOutputStream.close();
-            DongTaiLog.info("The remote file " + fileUrl + " was successfully written to the local cache.");
-            status = true;
         } catch (Exception ignore) {
             DongTaiLog.error("The remote file " + fileUrl + " download failure, please check the dongtai-token.");
         }
@@ -164,20 +173,21 @@ public boolean updateEnginePackage() {
 
     public boolean downloadEnginePackage() {
         if (engineNotExist(getInjectPackageCachePath()) || engineNotExist(getEnginePackageCachePath())) {
-            DongTaiLog.info("Engine does not exist in local cache, the engine will be downloaded.");
             return updateEnginePackage();
         } else {
             return true;
         }
     }
 
     public boolean install() {
+        String spyPackage = EngineManager.getInjectPackageCachePath();
+        String corePackage = EngineManager.getEnginePackageCachePath();
         try {
-            JarFile file = new JarFile(new File(EngineManager.getInjectPackageCachePath()));
+            JarFile file = new JarFile(new File(spyPackage));
             inst.appendToBootstrapClassLoaderSearch(file);
             file.close();
             if (IAST_CLASS_LOADER == null) {
-                IAST_CLASS_LOADER = new IastClassLoader(EngineManager.getEnginePackageCachePath());
+                IAST_CLASS_LOADER = new IastClassLoader(corePackage);
             }
             classOfEngine = IAST_CLASS_LOADER.loadClass(ENGINE_ENTRYPOINT_CLASS);
             String agentPath = this.getClass().getProtectionDomain().getCodeSource().getLocation().getFile();
@@ -187,11 +197,12 @@ public boolean install() {
                             AgentRegisterReport.getAgentFlag(), inst, agentPath);
             return true;
         } catch (IOException e) {
-            DongTaiLog.error("DongTai engine start failed, please contact staff for help.");
+            DongTaiLog.error("DongTai engine start failed, Reason: dongtai-spy.jar or dongtai-core.jar open failed. path: \n\tdongtai-core.jar: " + corePackage + "\n\tdongtai-spy.jar: " + spyPackage);
         } catch (ClassNotFoundException e) {
-            DongTaiLog.error(" DongTai engine start failed, please contact staff for help.");
+            e.printStackTrace();
+            DongTaiLog.error("ClassNotFoundException: DongTai engine start failed, please contact staff for help.");
         } catch (Throwable throwable) {
-            DongTaiLog.error("DongTai engine start failed, please contact staff for help.");
+            DongTaiLog.error("Throwable: DongTai engine start failed, please contact staff for help.");
             throwable.printStackTrace();
         }
         return false;
@@ -210,10 +221,13 @@ public boolean start() {
             }
             return false;
         } catch (InvocationTargetException e) {
+            e.printStackTrace();
             DongTaiLog.error("DongTai engine start failed, please contact staff for help.");
         } catch (NoSuchMethodException e) {
+            e.printStackTrace();
             DongTaiLog.error("DongTai engine start failed, please contact staff for help.");
         } catch (IllegalAccessException e) {
+            e.printStackTrace();
             DongTaiLog.error("DongTai engine start failed, please contact staff for help.");
         } catch (Throwable throwable) {
             DongTaiLog.error("DongTai engine start failed, please contact staff for help.");
@@ -301,10 +315,8 @@ public synchronized boolean uninstall() {
      * @return true-引擎不存在；false-引擎存在
      */
     private boolean engineNotExist(final String jarPath) {
-        DongTaiLog.info("Check if the engine[" + jarPath + "] needs to be updated");
-
         if (properties.isDebug()) {
-            DongTaiLog.info("current mode: debug, load engine from " + jarPath);
+            DongTaiLog.info("current mode: debug, try to read package from " + jarPath);
             File tempFile = new File(jarPath);
             return !tempFile.exists();
         } else {

dongtai-agent/src/main/java/com/secnium/iast/agent/middlewarerecognition/IServer.java

@@ -6,7 +6,7 @@
  * @author dongzhiyong@huoxian.cn
  */
 public interface IServer {
-    boolean isMatch(RuntimeMXBean paramRuntimeMXBean);
+    boolean isMatch(RuntimeMXBean paramRuntimeMXBean, ClassLoader loader);
 
     String getName();