Merge pull request #587 from Nizernizer/feature/config-pool

Feature/config pool
HXSecurity · Sep 7, 2023 · 4937755 · 4937755
2 parents bf4ee17 + bb7b271
commit 4937755
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 13 deletions.
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
@@ -34,6 +34,10 @@ public class IastProperties {
         put("uuid_path", PropertyConstant.PROPERTY_UUID_PATH);
         put("disabled_plugins", PropertyConstant.PROPERTY_DISABLED_PLUGINS);
         put("disabled_features", PropertyConstant.PROPERTY_DISABLED_FEATURES);
+        put("pool_capacity", PropertyConstant.PROPERTY_POOL_CAPACITY);
+        put("pool_size", PropertyConstant.PROPERTY_POOL_SIZE);
+        put("pool_max_size", PropertyConstant.PROPERTY_POOL_MAX_SIZE);
+        put("pool_keepalive", PropertyConstant.PROPERTY_POOL_KEEPALIVE);
     }};
 
     private static IastProperties instance;

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/constants/PropertyConstant.java b/dongtai-common/src/main/java/io/dongtai/iast/common/constants/PropertyConstant.java
@@ -33,4 +33,8 @@ public class PropertyConstant {
     public static final String PROPERTY_DISABLED_PLUGINS = "dongtai.disabled.plugins";
     public static final String PROPERTY_DISABLED_FEATURES = "dongtai.disabled.features";
     public static final String PROPERTY_TAINT_LENGTH = "dongtai.taint.length";
+    public static final String PROPERTY_POOL_CAPACITY = "dongtai.pool.capacity";
+    public static final String PROPERTY_POOL_SIZE = "dongtai.pool.size";
+    public static final String PROPERTY_POOL_MAX_SIZE = "dongtai.pool.max.size";
+    public static final String PROPERTY_POOL_KEEPALIVE = "dongtai.pool.keepalive";
 }
diff --git a/...ore/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilder.java b/...ore/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/PolicyBuilder.java
@@ -146,6 +146,7 @@ public static void buildValidator(Policy policy, PolicyNodeType type, JSONObject
         setInheritable(node, validatorNode);
         List<String[]> tags = parseTags(node, validatorNode);
         validatorNode.setTags(tags.get(0));
+        parseFlags(node, validatorNode);
         policy.addValidator(validatorNode);
     }
 

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/service/ThreadPools.java b/dongtai-core/src/main/java/io/dongtai/iast/core/service/ThreadPools.java
@@ -2,6 +2,7 @@
 
 import io.dongtai.iast.common.constants.AgentConstant;
 import io.dongtai.iast.core.replay.HttpRequestReplay;
+import io.dongtai.iast.core.utils.PropertyUtils;
 
 import java.util.concurrent.*;
 
@@ -10,39 +11,37 @@
  */
 public class ThreadPools {
 
-    private static final ExecutorService METHOD_REPORT_THREAD = new ThreadPoolExecutor(0, 5, 10L, TimeUnit.SECONDS,
-            new LinkedBlockingQueue<Runnable>(5120), new ThreadFactory() {
+    private static final PropertyUtils propertyUtils = PropertyUtils.getInstance();
+    private static final Integer poolSize = propertyUtils.getPoolSize();
+    private static final Integer poolMaxSize = propertyUtils.getPoolMaxSize();
+    private static final Integer poolKeepAlive = propertyUtils.getPoolKeepalive();
+    private static final Integer poolCapacity = propertyUtils.getPoolCapacity();
+
+    private static final ExecutorService METHOD_REPORT_THREAD = new ThreadPoolExecutor(poolSize, poolMaxSize, poolKeepAlive, TimeUnit.SECONDS,
+            new LinkedBlockingQueue<>(poolCapacity), new ThreadFactory() {
         @Override
         public Thread newThread(Runnable r) {
             return new Thread(r, AgentConstant.THREAD_NAME_PREFIX_CORE + "VulReport-" + r.hashCode());
         }
     });
 
     private static final ExecutorService COMMON_REPORT_THREAD = new ThreadPoolExecutor(0, 5, 10L, TimeUnit.SECONDS,
-            new LinkedBlockingQueue<Runnable>(10000), new ThreadFactory() {
+            new LinkedBlockingQueue<>(1024), new ThreadFactory() {
         @Override
         public Thread newThread(Runnable r) {
             return new Thread(r, AgentConstant.THREAD_NAME_PREFIX_CORE + "Report-" + r.hashCode());
         }
     });
 
     private static final ExecutorService REPLAY_REQUEST_THREAD = new ThreadPoolExecutor(0, 1, 10L, TimeUnit.SECONDS,
-            new LinkedBlockingQueue<Runnable>(1024), new ThreadFactory() {
+            new LinkedBlockingQueue<>(1024), new ThreadFactory() {
         @Override
         public Thread newThread(Runnable r) {
             return new Thread(r, AgentConstant.THREAD_NAME_PREFIX_CORE + "VulReplay-" + r.hashCode());
         }
 
     });
 
-    private static final ExecutorService LIMIT_REPORT_THREAD = new ThreadPoolExecutor(0, 5, 10L, TimeUnit.SECONDS,
-            new LinkedBlockingQueue<Runnable>(5120), new ThreadFactory() {
-        @Override
-        public Thread newThread(Runnable r) {
-            return new Thread(r, AgentConstant.THREAD_NAME_PREFIX_CORE + "LimitReport-" + r.hashCode());
-        }
-    });
-
 
     public static void execute(Runnable r) {
         COMMON_REPORT_THREAD.execute(r);

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/utils/PropertyUtils.java b/dongtai-core/src/main/java/io/dongtai/iast/core/utils/PropertyUtils.java
@@ -36,9 +36,17 @@ public class PropertyUtils {
     private final String propertiesFilePath;
 
     public static final Integer DEFAULT_TAINT_TO_STRING_CHAR_LIMIT = 1024;
+    public static final Integer DEFAULT_POOL_CAPACITY = 4096;
+    public static final Integer DEFAULT_POOL_SIZE = 0;
+    public static final Integer DEFAULT_POOL_MAX_SIZE = 10;
+    public static final Integer DEFAULT_POOL_KEEPALIVE = 10;
 
     // 污点转换为字符串的时候字符数长度限制
     private Integer taintToStringCharLimit = DEFAULT_TAINT_TO_STRING_CHAR_LIMIT;
+    private Integer poolCapacity;
+    private Integer poolSize;
+    private Integer poolMaxSize;
+    private Integer poolKeepalive;
 
     public static PropertyUtils getInstance(String propertiesFilePath) throws DongTaiPropertyConfigException, DongTaiEnvConfigException {
         if (null == instance) {
@@ -78,7 +86,7 @@ private void init() throws DongTaiPropertyConfigException, DongTaiEnvConfigExcep
 
         // 初始化一些参数
         this.initTaintToStringCharLimit();
-
+        this.initPool();
     }
 
     public static String getTmpDir() {
@@ -244,6 +252,34 @@ public static Integer getTaintToStringCharLimit() {
         return instance.taintToStringCharLimit;
     }
 
+    public Integer getPoolCapacity() {
+        if (instance == null) {
+            return DEFAULT_POOL_CAPACITY;
+        }
+        return instance.poolCapacity;
+    }
+
+    public Integer getPoolSize() {
+        if (instance == null) {
+            return DEFAULT_POOL_SIZE;
+        }
+        return instance.poolSize;
+    }
+
+    public Integer getPoolMaxSize() {
+        if (instance == null) {
+            return DEFAULT_POOL_MAX_SIZE;
+        }
+        return instance.poolMaxSize;
+    }
+
+    public Integer getPoolKeepalive() {
+        if (instance == null) {
+            return DEFAULT_POOL_KEEPALIVE;
+        }
+        return instance.poolKeepalive;
+    }
+
     /**
      * 初始化taintToStringCharLimit参数的值
      *
@@ -274,6 +310,36 @@ public void initTaintToStringCharLimit() throws DongTaiPropertyConfigException,
 
     }
 
+    private void initPool() throws DongTaiPropertyConfigException, DongTaiEnvConfigException {
+        this.poolCapacity = parseAndSetProperty(PropertyConstant.PROPERTY_POOL_CAPACITY, DEFAULT_POOL_CAPACITY);
+        this.poolSize = parseAndSetProperty(PropertyConstant.PROPERTY_POOL_SIZE, DEFAULT_POOL_SIZE);
+        this.poolMaxSize = parseAndSetProperty(PropertyConstant.PROPERTY_POOL_MAX_SIZE, DEFAULT_POOL_MAX_SIZE);
+        this.poolKeepalive = parseAndSetProperty(PropertyConstant.PROPERTY_POOL_KEEPALIVE, DEFAULT_POOL_KEEPALIVE);
+    }
+
+    private Integer parseAndSetProperty(String propertyKey,Integer defaultValue) throws DongTaiPropertyConfigException, DongTaiEnvConfigException {
+        String propertyStr = cfg.getProperty(propertyKey);
+        Integer value = defaultValue;
+        if (!StringUtils.isBlank(propertyStr)) {
+            value = Integer.parseInt(propertyStr.trim());
+            if (value <= 0) {
+                throw new DongTaiPropertyConfigException("The value of parameter " + propertyKey
+                        + " value " + propertyStr + " in your configuration file " + this.propertiesFilePath + " is illegal, such as passing a number greater than 1");
+            }
+        }
+
+        // 2. 然后从环境变量中读取
+        propertyStr = System.getProperty(propertyKey);
+        if (!StringUtils.isBlank(propertyStr)) {
+            value = Integer.parseInt(propertyStr.trim());
+            if (value <= 0) {
+                throw new DongTaiEnvConfigException("The value of this parameter " + propertyKey
+                        + " value " + propertyStr + " in your environment variables is illegal, such as passing an number greater than 1");
+            }
+        }
+        return value;
+    }
+
     /**
      * Property文件配置错误
      */