Merge pull request #614 from 15911075183ma/feat-ratelimiter-server

Feat ratelimiter server
HXSecurity · Dec 18, 2023 · 449797c · 449797c
2 parents 9c40484 + bdfa411
commit 449797c
Show file tree

Hide file tree

Showing 20 changed files with 444 additions and 36 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
@@ -40,6 +40,7 @@ private static String[] parseAgentArgs(String[] args) throws ParseException {
         attachOptions.addOption(build("log_disable_collector", "log_disable_collector", "optional: DongTai agent disable log collector."));
         attachOptions.addOption(build("disabled_plugins", "disabled_plugins", "optional: DongTai agent disable plugins."));
         attachOptions.addOption(build("disabled_features", "disabled_features", "optional: DongTai agent disable features."));
+        attachOptions.addOption(build("rate_caps", "rate_caps", "optional: the maximum speed of the interface example： 100"));
 
         CommandLineParser parser = new DefaultParser();
         HelpFormatter formatter = new HelpFormatter();

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
@@ -5,7 +5,9 @@
 import io.dongtai.iast.agent.report.AgentRegisterReport;
 import io.dongtai.iast.common.constants.AgentConstant;
 import io.dongtai.iast.common.scope.ScopeManager;
-import io.dongtai.iast.common.state.*;
+import io.dongtai.iast.common.state.AgentState;
+import io.dongtai.iast.common.state.State;
+import io.dongtai.iast.common.state.StateCause;
 import io.dongtai.log.DongTaiLog;
 import io.dongtai.log.ErrorCode;
 

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
@@ -40,6 +40,7 @@ public class IastProperties {
         put("pool_size", PropertyConstant.PROPERTY_POOL_SIZE);
         put("pool_max_size", PropertyConstant.PROPERTY_POOL_MAX_SIZE);
         put("pool_keepalive", PropertyConstant.PROPERTY_POOL_KEEPALIVE);
+        put("rate_caps", PropertyConstant.THE_UPPER_LIMIT_OF_THE_INTERFACE_RATE);
     }};
 
     private static IastProperties instance;

diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/ConfigMonitor.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/ConfigMonitor.java
@@ -9,6 +9,7 @@
 import io.dongtai.iast.common.config.ConfigKey;
 import io.dongtai.iast.common.constants.AgentConstant;
 import io.dongtai.iast.common.constants.ApiPath;
+import io.dongtai.iast.common.utils.limit.InterfaceRateLimiterUtil;
 import io.dongtai.log.DongTaiLog;
 import io.dongtai.log.ErrorCode;
 
@@ -48,6 +49,37 @@ private void updateConfig() {
         if (logLevel != null) {
             DongTaiLog.setLevel(DongTaiLog.parseLevel(logLevel));
         }
+
+        //获取是否开启qps限流
+        Boolean enableQpsRate = ConfigBuilder.getInstance().get(ConfigKey.ENABLE_QPS_RATE_LIMIT);
+        if (enableQpsRate) {
+            int qpsRateLimit = ConfigBuilder.getInstance().get(ConfigKey.QPS_RATE_LIMIT);
+            if (qpsRateLimit <= 0){
+                DongTaiLog.error("qpsRateLimit the value cannot be less than 0");
+                qpsRateLimit =  100;
+                DongTaiLog.error("qpsRateLimit revert to 100");
+
+            }
+            int tokenBucketPoolSize = ConfigBuilder.getInstance().get(ConfigKey.TOKEN_BUCKET_POOL_SIZE);
+            if (tokenBucketPoolSize <= 0){
+                DongTaiLog.error("tokenBucketPoolSize the value cannot be less than 0");
+                tokenBucketPoolSize=5000;
+                DongTaiLog.error("tokenBucketPoolSize revert to 5000");
+
+            }
+            //判断是否已经开启，如果已经开启，则更新数据
+            if (InterfaceRateLimiterUtil.getRateLimiterState()) {
+                InterfaceRateLimiterUtil.updateTheData(qpsRateLimit, tokenBucketPoolSize);
+            } else {
+                //初始化令牌池工具，设置池大小和速率
+                InterfaceRateLimiterUtil.initializeInstance(qpsRateLimit, tokenBucketPoolSize);
+            }
+        }else {
+            if(InterfaceRateLimiterUtil.getRateLimiterState()){
+                InterfaceRateLimiterUtil.turnOffTheRateLimiter();
+            }
+        }
+
     }
 
     @Override

diff --git a/dongtai-agent/src/main/resources/iast.properties b/dongtai-agent/src/main/resources/iast.properties
@@ -11,4 +11,5 @@ dongtai.app.template=0
 iast.proxy.enable=false
 iast.proxy.host=
 iast.proxy.port=
-iast.server.mode=local
+iast.server.mode=local
+rate.caps=100
diff --git a/dongtai-common/pom.xml b/dongtai-common/pom.xml
@@ -38,6 +38,12 @@
             <artifactId>fastjson2</artifactId>
             <version>${fastjson2.version}</version>
         </dependency>
+
+        <dependency>
+            <groupId>com.github.vladimir-bukhtoyarov</groupId>
+            <artifactId>bucket4j-core</artifactId>
+            <version>${bucket4j.version}</version>
+        </dependency>
     </dependencies>
 
     <build>

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigBuilder.java b/dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigBuilder.java
@@ -28,6 +28,12 @@ private ConfigBuilder() {
                 Config.<String>create(ConfigKey.LOGGER_LEVEL));
         this.configMap.put(ConfigKey.VALIDATED_SINK,
                 Config.<Boolean>create(ConfigKey.VALIDATED_SINK).setDefaultValue(false));
+        this.configMap.put(ConfigKey.ENABLE_QPS_RATE_LIMIT,
+                Config.<Boolean>create(ConfigKey.ENABLE_QPS_RATE_LIMIT).setDefaultValue(false));
+        this.configMap.put(ConfigKey.QPS_RATE_LIMIT,
+                Config.<Integer>create(ConfigKey.QPS_RATE_LIMIT).setDefaultValue(100));
+        this.configMap.put(ConfigKey.TOKEN_BUCKET_POOL_SIZE,
+                Config.<Integer>create(ConfigKey.TOKEN_BUCKET_POOL_SIZE).setDefaultValue(5000));
     }
 
     public static ConfigBuilder getInstance() {
@@ -68,6 +74,9 @@ public void update(JSONObject config) {
         updateBool(config, ConfigKey.JsonKey.JSON_ENABLE_LOGGER);
         updateString(config, ConfigKey.JsonKey.JSON_LOGGER_LEVEL);
         updateBool(config, ConfigKey.JsonKey.JSON_VALIDATED_SINK);
+        updateBool(config, ConfigKey.JsonKey.JSON_ENABLE_QPS_RATE_LIMIT);
+        updateInt(config, ConfigKey.JsonKey.JSON_QPS_RATE_LIMIT);
+        updateInt(config, ConfigKey.JsonKey.JSON_TOKEN_BUCKET_POOL_SIZE);
         updateRequestDenyList(config);
     }
 
@@ -89,6 +98,8 @@ private void updateBool(JSONObject config, ConfigKey.JsonKey jsonKey) {
                 Boolean value = config.getBoolean(jsonKey.getKey());
                 conf.setValue(value);
             }
+        }catch (JSONException jsonException){
+            DongTaiLog.trace("configuration file resolution error {}",jsonException.getMessage());
         } catch (Throwable e) {
             DongTaiLog.warn(ErrorCode.UTIL_CONFIG_LOAD_FAILED,e.getMessage());
         }
@@ -102,6 +113,9 @@ private void updateInt(JSONObject config, ConfigKey.JsonKey jsonKey) {
                 Integer value = config.getInt(jsonKey.getKey());
                 conf.setValue(value);
             }
+        }catch (JSONException jsonException){
+            DongTaiLog.trace("configuration file resolution error {}",jsonException.getMessage());
+
         } catch (Throwable e) {
             DongTaiLog.warn(ErrorCode.UTIL_CONFIG_LOAD_FAILED,e.getMessage());
         }
@@ -117,6 +131,8 @@ private void updateString(JSONObject config, ConfigKey.JsonKey jsonKey) {
                     conf.setValue(value);
                 }
             }
+        }catch (JSONException jsonException){
+            DongTaiLog.trace("configuration file resolution error {}",jsonException.getMessage());
         } catch (Throwable e) {
             DongTaiLog.warn(ErrorCode.UTIL_CONFIG_LOAD_FAILED,e.getMessage());
         }
@@ -132,6 +148,8 @@ private void updateRequestDenyList(JSONObject config) {
                 RequestDenyList requestDenyList = RequestDenyList.parse(value);
                 conf.setValue(requestDenyList);
             }
+        }catch (JSONException jsonException){
+            DongTaiLog.trace("updateRequestDenyList configuration file resolution error {}",jsonException.getMessage());
         } catch (Throwable e) {
             DongTaiLog.warn(ErrorCode.UTIL_CONFIG_LOAD_FAILED,e.getMessage());
         }

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigKey.java b/dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigKey.java
@@ -9,6 +9,9 @@ public enum ConfigKey {
     ENABLE_LOGGER,
     LOGGER_LEVEL,
     VALIDATED_SINK,
+    ENABLE_QPS_RATE_LIMIT,
+    QPS_RATE_LIMIT,
+    TOKEN_BUCKET_POOL_SIZE,
     ;
 
     public enum JsonKey {
@@ -20,6 +23,9 @@ public enum JsonKey {
         JSON_ENABLE_LOGGER("enable_log", ENABLE_LOGGER),
         JSON_LOGGER_LEVEL("log_level", LOGGER_LEVEL),
         JSON_VALIDATED_SINK("report_validated_sink", VALIDATED_SINK),
+        JSON_ENABLE_QPS_RATE_LIMIT("enable_qps_rate_limit", ENABLE_QPS_RATE_LIMIT),
+        JSON_QPS_RATE_LIMIT("qps_rate_limit", QPS_RATE_LIMIT),
+        JSON_TOKEN_BUCKET_POOL_SIZE("token_bucket_pool_size", TOKEN_BUCKET_POOL_SIZE),
         ;
 
         private final String key;

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/config/RequestDeny.java b/dongtai-common/src/main/java/io/dongtai/iast/common/config/RequestDeny.java
@@ -1,9 +1,12 @@
 package io.dongtai.iast.common.config;
 
+import io.dongtai.log.DongTaiLog;
 import org.json.JSONException;
 import org.json.JSONObject;
 
 import java.util.*;
+import java.util.regex.Pattern;
+import java.util.regex.PatternSyntaxException;
 
 public class RequestDeny {
     private static final String KEY_TARGET_TYPE = "target_type";
@@ -12,7 +15,7 @@ public class RequestDeny {
 
     private static final Map<TargetType, List<Operator>> OPERATOR_MAP = new HashMap<TargetType, List<Operator>>() {{
         put(TargetType.URL, Arrays.asList(Operator.EQUAL, Operator.NOT_EQUAL, Operator.CONTAIN, Operator.NOT_CONTAIN));
-        put(TargetType.HEADER_KEY, Arrays.asList(Operator.EXISTS, Operator.NOT_EXISTS));
+        put(TargetType.HEADER_KEY, Arrays.asList(Operator.EXISTS, Operator.NOT_EXISTS,Operator.EXISTS_KEY_AND_VALUE));
     }};
 
     public enum TargetType {
@@ -47,6 +50,8 @@ public enum Operator {
         NOT_CONTAIN("NOT_CONTAIN"),
         EXISTS("EXISTS"),
         NOT_EXISTS("NOT_EXISTS"),
+        //此类型是为了兼容请求头的value匹配，此时 value使用：分割格式为key：value 值使用base64编码，需要解码使用，value使用正则匹配
+        EXISTS_KEY_AND_VALUE("EXISTS_KEY_AND_VALUE"),
         ;
 
         private final String key;
@@ -93,11 +98,22 @@ public static RequestDeny parse(JSONObject config) {
             if (operator == null || !OPERATOR_MAP.get(targetType).contains(operator)) {
                 return null;
             }
-
             String value = config.getString(KEY_VALUE);
             if (value == null || value.isEmpty()) {
                 return null;
             }
+            if (Operator.EXISTS_KEY_AND_VALUE.equals(operator)){
+                String[] split = value.split(":");
+                String vle = new String(Base64.getDecoder().decode(split[1]));
+                try {
+                    Pattern.compile(vle);
+                } catch (PatternSyntaxException e) {
+                    String key = new String(Base64.getDecoder().decode(split[0]));
+                    DongTaiLog.error("the regex is not valid please check！ key：{},value:{}",key,vle);
+                    return null;
+                }
+
+            }
 
             return new RequestDeny(targetType, operator, value);
         } catch (JSONException ignore) {
@@ -135,6 +151,18 @@ private boolean matchHeaderKey(Map<String, String> headers) {
         if (headers == null || headers.isEmpty()) {
             return false;
         }
+        if (Operator.EXISTS_KEY_AND_VALUE.equals(operator)){
+            String[] split = this.value.split(":");
+            String key = new String(Base64.getDecoder().decode(split[0])).toLowerCase();
+            String vle = new String(Base64.getDecoder().decode(split[1]));
+
+            for (Map.Entry<String, String> entry : headers.entrySet()) {
+                if (key.equals(entry.getKey().toLowerCase())) {
+                    Pattern pattern = Pattern.compile(vle);
+                    return pattern.matcher(entry.getValue()).find();
+                }
+            }
+        }
 
         boolean exists = false;
         String matchVal = this.value.toLowerCase();
@@ -147,10 +175,12 @@ private boolean matchHeaderKey(Map<String, String> headers) {
 
         if (Operator.EXISTS.equals(operator)) {
             return exists;
-        } else if (Operator.NOT_EXISTS.equals(operator)) {
+        }
+        if (Operator.NOT_EXISTS.equals(operator)) {
             return !exists;
         }
 
+
         return false;
     }
 

diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/config/RequestDenyList.java b/dongtai-common/src/main/java/io/dongtai/iast/common/config/RequestDenyList.java
@@ -2,10 +2,12 @@
 
 import org.json.JSONArray;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
 
 public class RequestDenyList {
-    private final List<List<RequestDeny>> denies = new ArrayList<List<RequestDeny>>();
+    private final List<RequestDeny> denies = new ArrayList<RequestDeny>();
 
     public static RequestDenyList parse(JSONArray config) {
         if (config == null || config.length() == 0) {
@@ -29,7 +31,7 @@ public static RequestDenyList parse(JSONArray config) {
                 andList.add(requestDeny);
             }
             if (!andList.isEmpty()) {
-                denyList.denies.add(andList);
+                denyList.denies.addAll(andList);
             }
         }
 
@@ -40,10 +42,24 @@ public static RequestDenyList parse(JSONArray config) {
     }
 
     public void addRule(List<RequestDeny> requestDenies) {
-        this.denies.add(requestDenies);
+        this.denies.addAll(requestDenies);
     }
 
+    /**
+     *
+     * @param url 请求URL
+     * @param headers 请求头集合
+     * @return true 成功匹配 false 未匹配
+     */
     public boolean match(String url, Map<String, String> headers) {
+        for (RequestDeny requestDeny : denies) {
+            if (requestDeny.match(url, headers)) {
+                return true; // 匹配到条件，提前终止循环
+            }
+        }
+        return false;
+    }
+/*    public boolean match(String url, Map<String, String> headers) {
         boolean matched = false;
         for (List<RequestDeny> denyList : this.denies) {
             boolean subHasNoMatch = false;
@@ -60,7 +76,7 @@ public boolean match(String url, Map<String, String> headers) {
             }
         }
         return matched;
-    }
+    }*/
 
     @Override
     public String toString() {

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
@@ -37,4 +37,6 @@ public class PropertyConstant {
     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";
+    //  接口速率上限，可配可不配，不配置为未开启
+    public static final String THE_UPPER_LIMIT_OF_THE_INTERFACE_RATE = "rate.caps";
 }