Merge pull request #360 from lostsnow/feature/config-disable-log-coll…

…ector Feature/config disable log collector
HXSecurity · Sep 22, 2022 · 9fa60af · 9fa60af
2 parents 646e139 + 04fc655
commit 9fa60af
Show file tree

Hide file tree

Showing 7 changed files with 113 additions and 88 deletions.
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java
@@ -1,11 +1,12 @@
 package io.dongtai.iast.agent;
 
-import java.io.*;
-import java.util.Arrays;
-
 import io.dongtai.iast.agent.util.FileUtils;
 import org.apache.commons.cli.*;
 
+import java.io.File;
+import java.io.IOException;
+import java.util.Arrays;
+
 /**
  * @author dongzhiyong@huoxian.cn
  */
@@ -31,6 +32,7 @@ private static String[] parseAgentArgs(String[] args) throws ParseException {
         attachOptions.addOption(build("server_package", "server_package", "optional: DongTai core package download way."));
         attachOptions.addOption(build("log_level", "log_level", "optional: DongTai agent log print level."));
         attachOptions.addOption(build("log_path", "log_path", "optional: DongTai agent log print path."));
+        attachOptions.addOption(build("disable_log_collector", "disable_log_collector", "optional: DongTai agent disable log collector."));
 
         CommandLineParser parser = new DefaultParser();
         HelpFormatter formatter = new HelpFormatter();
@@ -81,6 +83,9 @@ private static String[] parseAgentArgs(String[] args) throws ParseException {
             if (result.hasOption("log_path")) {
                 attachArgs.append("&logPath=").append(result.getOptionValue("log_path"));
             }
+            if (result.hasOption("disable_log_collector")) {
+                attachArgs.append("&disableLogCollector=").append(result.getOptionValue("disable_log_collector"));
+            }
             return new String[]{pid, attachArgs.toString()};
         } else {
             formatter.printHelp("java -jar agent.jar", attachOptions, true);

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/AgentLauncher.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/AgentLauncher.java
@@ -4,17 +4,13 @@
 import io.dongtai.iast.agent.monitor.MonitorDaemonThread;
 import io.dongtai.iast.agent.monitor.impl.EngineMonitor;
 import io.dongtai.iast.agent.report.AgentRegisterReport;
-import io.dongtai.iast.agent.util.FileUtils;
 import io.dongtai.iast.agent.util.ThreadUtils;
 import io.dongtai.log.DongTaiLog;
 
 import java.io.File;
-import java.io.IOException;
 import java.lang.instrument.Instrumentation;
 import java.util.*;
 
-import static io.dongtai.iast.agent.Agent.*;
-
 /**
  * @author dongzhiyong@huoxian.cn
  */
@@ -23,9 +19,6 @@ public class AgentLauncher {
     public static final String LAUNCH_MODE_AGENT = "agent";
     public static final String LAUNCH_MODE_ATTACH = "attach";
     public static String LAUNCH_MODE;
-    private static String FLUENT_FILE;
-    private static String FLUENT_FILE_CONF;
-    public static Process fluent;
 
     static {
         /**
@@ -113,6 +106,9 @@ public static void agentmain(String args, Instrumentation inst) {
             if (argsMap.containsKey("logPath")) {
                 System.setProperty("dongtai.log.path", argsMap.get("logPath"));
             }
+            if (argsMap.containsKey("disableLogCollector")) {
+                System.setProperty("dongtai.disable.log-collector", argsMap.get("disableLogCollector"));
+            }
         } catch (Exception e) {
             DongTaiLog.error(e);
         }
@@ -162,7 +158,7 @@ private static void install(final Instrumentation inst) {
         IastProperties.getInstance();
         Boolean send = AgentRegisterReport.send();
         if (send) {
-            extractFluent();
+            LogCollector.extractFluent();
             DongTaiLog.info("Agent registered successfully.");
             Boolean agentStat = AgentRegisterReport.agentStat();
             if (!agentStat) {
@@ -179,46 +175,6 @@ private static void install(final Instrumentation inst) {
         }
     }
 
-    private static void doFluent() {
-        String[] execution = {
-                "nohup",
-                FLUENT_FILE,
-                "-c",
-                FLUENT_FILE_CONF
-        };
-        try {
-            fluent = Runtime.getRuntime().exec(execution);
-            Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
-                public void run() {
-                    fluent.destroy();
-                }
-            }));
-        } catch (IOException e) {
-            DongTaiLog.error(e);
-        }
-    }
-
-    private static void extractFluent() {
-        try {
-            if (!isMacOs() && !isWindows()) {
-                FLUENT_FILE = System.getProperty("java.io.tmpdir.dongtai") + "iast" + File.separator + "fluent";
-                FileUtils.getResourceToFile("bin/fluent", FLUENT_FILE);
-
-                FLUENT_FILE_CONF = System.getProperty("java.io.tmpdir.dongtai") + "iast" + File.separator + "fluent.conf";
-                FileUtils.getResourceToFile("bin/fluent.conf", FLUENT_FILE_CONF);
-                FileUtils.confReplace(FLUENT_FILE_CONF);
-                if ((new File(FLUENT_FILE)).setExecutable(true)) {
-                    DongTaiLog.info("fluent extract success.");
-                } else {
-                    DongTaiLog.info("fluent extract failure. please set execute permission, file: {}", FLUENT_FILE);
-                }
-                doFluent();
-            }
-        } catch (IOException e) {
-            DongTaiLog.error(e);
-        }
-    }
-
     private static void loadEngine(final Instrumentation inst) {
         EngineManager engineManager = EngineManager.getInstance(inst, LAUNCH_MODE, EngineManager.getPID());
         MonitorDaemonThread daemonThread = new MonitorDaemonThread(engineManager);

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/IastProperties.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/IastProperties.java
@@ -3,9 +3,7 @@
 import io.dongtai.iast.agent.util.FileUtils;
 import io.dongtai.log.DongTaiLog;
 
-import java.io.File;
-import java.io.IOException;
-import java.io.InputStream;
+import java.io.*;
 import java.util.Properties;
 
 /**
@@ -37,6 +35,7 @@ public class IastProperties {
     private String customCoreJarUrl;
     private String customSpyJarUrl;
     private String customApiJarUrl;
+    private Boolean disableLogCollector;
 
     public static IastProperties getInstance() {
         if (null == instance) {
@@ -251,6 +250,15 @@ public String getFallbackVersion() {
         return fallbackVersion;
     }
 
+    public Boolean getDisableLogCollector() {
+        if (disableLogCollector == null) {
+            String disable = System.getProperty("dongtai.disable.log-collector",
+                    cfg.getProperty("dongtai.disable.log-collector", "false"));
+            disableLogCollector = "true".equals(disable);
+        }
+        return disableLogCollector;
+    }
+
     public boolean isDebug() {
         return "true".equalsIgnoreCase(getDebugFlag());
     }

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/LogCollector.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/LogCollector.java
@@ -0,0 +1,73 @@
+package io.dongtai.iast.agent;
+
+import io.dongtai.iast.agent.util.FileUtils;
+import io.dongtai.log.DongTaiLog;
+
+import java.io.File;
+import java.io.IOException;
+
+import static io.dongtai.iast.agent.Agent.*;
+
+public class LogCollector {
+    private static String FLUENT_FILE;
+    private static String FLUENT_FILE_CONF;
+    private static Process fluent;
+
+    public static void extractFluent() {
+        if (IastProperties.getInstance().getDisableLogCollector()) {
+            return;
+        }
+        try {
+            if (!isMacOs() && !isWindows()) {
+                FLUENT_FILE = System.getProperty("java.io.tmpdir.dongtai") + "iast" + File.separator + "fluent";
+                FileUtils.getResourceToFile("bin/fluent", FLUENT_FILE);
+
+                FLUENT_FILE_CONF = System.getProperty("java.io.tmpdir.dongtai") + "iast" + File.separator + "fluent.conf";
+                FileUtils.getResourceToFile("bin/fluent.conf", FLUENT_FILE_CONF);
+                FileUtils.confReplace(FLUENT_FILE_CONF);
+                if (!(new File(FLUENT_FILE)).setExecutable(true)) {
+                    DongTaiLog.info("fluent setExecutable failure. please set execute permission, file: {}", FLUENT_FILE);
+                }
+                doFluent();
+            }
+        } catch (IOException e) {
+            DongTaiLog.error("fluent extract failure", e);
+        }
+    }
+
+    public static void doFluent() {
+        if (IastProperties.getInstance().getDisableLogCollector()) {
+            return;
+        }
+        String[] execution = {
+                "nohup",
+                FLUENT_FILE,
+                "-c",
+                FLUENT_FILE_CONF
+        };
+        try {
+            fluent = Runtime.getRuntime().exec(execution);
+            DongTaiLog.info("fluent process started");
+            Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
+                public void run() {
+                    stopFluent();
+                }
+            }));
+        } catch (IOException e) {
+            DongTaiLog.error("fluent process start failed", e);
+        }
+    }
+
+    public static void stopFluent() {
+        if (fluent == null) {
+            return;
+        }
+        try {
+            fluent.destroy();
+            DongTaiLog.info("fluent process stopped");
+        } catch (Exception ignored) {
+        } finally {
+            fluent = null;
+        }
+    }
+}
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/manager/EngineManager.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/manager/EngineManager.java
@@ -1,8 +1,6 @@
 package io.dongtai.iast.agent.manager;
 
-import io.dongtai.iast.agent.AgentLauncher;
-import io.dongtai.iast.agent.IastClassLoader;
-import io.dongtai.iast.agent.IastProperties;
+import io.dongtai.iast.agent.*;
 import io.dongtai.iast.agent.middlewarerecognition.ServerDetect;
 import io.dongtai.iast.agent.middlewarerecognition.tomcat.AbstractTomcat;
 import io.dongtai.iast.agent.monitor.MonitorDaemonThread;
@@ -19,9 +17,7 @@
 import java.lang.instrument.Instrumentation;
 import java.lang.management.ManagementFactory;
 import java.lang.reflect.InvocationTargetException;
-import java.net.HttpURLConnection;
-import java.net.Proxy;
-import java.net.URL;
+import java.net.*;
 import java.util.jar.JarFile;
 
 /**
@@ -474,12 +470,7 @@ public synchronized boolean uninstall() {
             uninstallObject();
             MonitorDaemonThread.isExit = true;
         }
-        try {
-            AgentLauncher.fluent.destroy();
-        }catch (Exception ignored){
-        }finally {
-            AgentLauncher.fluent=null;
-        }
+        LogCollector.stopFluent();
         return true;
     }
 

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/report/AgentRegisterReport.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/report/AgentRegisterReport.java
@@ -8,13 +8,12 @@
 import io.dongtai.iast.agent.util.base64.Base64Encoder;
 import io.dongtai.iast.agent.util.http.HttpClientUtils;
 import io.dongtai.log.DongTaiLog;
+import org.json.JSONArray;
+import org.json.JSONObject;
 
 import java.net.*;
 import java.util.Enumeration;
 
-import org.json.JSONArray;
-import org.json.JSONObject;
-
 /**
  * @author dongzhiyong@huoxian.cn
  */
@@ -225,11 +224,11 @@ public void register(){
                 setAgentData(responseRaw);
             }
         } catch (NullPointerException e) {
-            DongTaiLog.error("Agent registration to {} failed, Token: {}, Reason: {}", IastProperties.getInstance().getBaseUrl(), IastProperties.getInstance().getIastServerToken(), e.getMessage());
+            DongTaiLog.error("Agent registration to {} failed, Token: {}, Reason: {}",
+                    IastProperties.getInstance().getBaseUrl(), IastProperties.getInstance().getIastServerToken(), e.getMessage());
         } catch (Exception e) {
-            DongTaiLog.error(e);
             DongTaiLog.error("Agent registration to {} failed 10 seconds later, cause: {}, token: {}",
-                    IastProperties.getInstance().getBaseUrl(), e.getMessage(), IastProperties.getInstance().getIastServerToken());
+                    IastProperties.getInstance().getBaseUrl(), e.toString(), IastProperties.getInstance().getIastServerToken());
         }
     }
 
@@ -252,11 +251,12 @@ private void setAgentData(StringBuilder responseRaw) {
                 JSONObject data = (JSONObject) responseObj.get("data");
                 agentId = (Integer) data.get("id");
                 coreRegisterStart = (Integer) data.get("coreAutoStart");
-            }else {
-                DongTaiLog.error("Register msg: "+ responseRaw);
+            } else {
+                DongTaiLog.error("Register msg: " + responseRaw);
             }
-        }catch (Exception e){
-            DongTaiLog.error("DongTai server no response.");
+        } catch (Exception e) {
+            DongTaiLog.error("Parse {} register response failed: {}",
+                    IastProperties.getInstance().getBaseUrl(), e.toString());
         }
     }
 

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/util/http/HttpClientUtils.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/util/http/HttpClientUtils.java
@@ -3,21 +3,13 @@
 import io.dongtai.iast.agent.IastProperties;
 import io.dongtai.log.DongTaiLog;
 
-import java.io.BufferedReader;
-import java.io.InputStream;
-import java.io.InputStreamReader;
-import java.net.HttpURLConnection;
-import java.net.InetSocketAddress;
-import java.net.Proxy;
-import java.net.URL;
+import javax.net.ssl.*;
+import java.io.*;
+import java.net.*;
 import java.nio.charset.Charset;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.zip.GZIPOutputStream;
-import javax.net.ssl.HostnameVerifier;
-import javax.net.ssl.HttpsURLConnection;
-import javax.net.ssl.SSLContext;
-import javax.net.ssl.TrustManager;
 
 /**
  * @author dongzhiyong@huoxian.cn
@@ -113,7 +105,7 @@ private static StringBuilder sendRequest(HttpMethods method, String baseUrl, Str
                 response.append('\r');
             }
             rd.close();
-            DongTaiLog.debug("dongtai upload url is {}, request is {} ,response is {}", urlStr, data, response.toString());
+            DongTaiLog.trace("dongtai upload url is {}, request is {} ,response is {}", urlStr, data, response.toString());
             return response;
         } catch (Exception e){
             DongTaiLog.error(e);