diff --git a/conf/pangu/pangu_pxy.conf b/conf/pangu/pangu_pxy.conf index 6f6d8a3..49f2c50 100644 --- a/conf/pangu/pangu_pxy.conf +++ b/conf/pangu/pangu_pxy.conf @@ -50,6 +50,12 @@ log_fsstat_trig=1 log_fsstat_dst_ip=10.4.20.202 log_fsstat_dst_port=8125 +[ratelimit] +enable=0 +token_name=ratelimit +redis_server=192.168.40.137 +redis_port=6379 + [maat] # 0:json 1: redis 2: iris maat_input_mode=1 diff --git a/plugin/business/pangu-http/src/pangu_http.cpp b/plugin/business/pangu-http/src/pangu_http.cpp index c75346a..4197fb7 100644 --- a/plugin/business/pangu-http/src/pangu_http.cpp +++ b/plugin/business/pangu-http/src/pangu_http.cpp @@ -25,6 +25,7 @@ #include #include +#include "ratelimiter.h" #define MAX_SCAN_RESULT 16 #define MAX_EDIT_ZONE_NUM 64 @@ -146,6 +147,9 @@ struct pangu_rt struct event_base* gc_evbase; struct event* gcev; + Ratelimiter_handle_t ratelimiter; + int enable_rate; + int ctrl_compile_idx; int ca_store_reseting; @@ -981,6 +985,40 @@ error_out: return ret; } +Ratelimiter_handle_t ratelimit_handle_create(const char* profile_path, const char* static_section) +{ + int ret=0; + int redis_port = 0, interval_sec = 0; + char redis_server[TFE_STRING_MAX] = {0}; + char token_name[TFE_STRING_MAX] = {0}; + + Ratelimiter_handle_t ratelimit = NULL; + + MESA_load_profile_int_def(profile_path, static_section, "enable", &(g_pangu_rt->enable_rate), 0); + MESA_load_profile_int_def(profile_path, static_section, "redis_port", &(redis_port), 6379); + MESA_load_profile_string_def(profile_path, static_section, "redis_server", redis_server, sizeof(redis_server), ""); + MESA_load_profile_string_def(profile_path, static_section, "token_name", token_name, sizeof(token_name), ""); + + MESA_load_profile_int_def(profile_path, static_section, "interval_sec", &(interval_sec), 1); + + ratelimit=Ratelimiter_create(token_name, g_pangu_rt->local_logger); + + Ratelimiter_set_opt(ratelimit, RATELIMITER_OPT_INTERVAL_SEC, &interval_sec, sizeof(interval_sec)); + Ratelimiter_set_opt(ratelimit, RATELIMITER_OPT_REDIS_IP, redis_server, strlen(redis_server) + 1); + Ratelimiter_set_opt(ratelimit, RATELIMITER_OPT_REDIS_PORT, &redis_port, sizeof(redis_port)); + + ret = Ratelimiter_start(ratelimit); + if (ret < 0) + { + TFE_LOG_ERROR(g_pangu_rt->local_logger, "%s Ratelimiter init failed.", __FUNCTION__); + goto error_out; + } + return ratelimit; +error_out: + Ratelimiter_stop(ratelimit); + return NULL; +} + int pangu_http_init(struct tfe_proxy * proxy) { const char * profile_path = "./conf/pangu/pangu_pxy.conf"; @@ -1002,8 +1040,9 @@ int pangu_http_init(struct tfe_proxy * proxy) } g_pangu_rt->fs_handle = tfe_proxy_get_fs_handle(); - pangu_http_stat_init(g_pangu_rt); + g_pangu_rt->ratelimiter=ratelimit_handle_create(profile_path, "ratelimit"); + pangu_http_stat_init(g_pangu_rt); if(pangu_policy_init(profile_path, "MAAT", "DYNAMIC_MAAT")<0) { goto error_out; @@ -1230,6 +1269,15 @@ static struct manipulate_profile* get_profile_by_id(int profile_table, int profi return result; } +static unsigned long long try_send_by_token(int inject_sz) +{ + if (g_pangu_rt->enable_rate != 1) + { + return 1; + } + return Ratelimiter_customer_factory(g_pangu_rt->ratelimiter, inject_sz); +} + static int pangu_action_weight[__PG_ACTION_MAX] = {0}; void __pangu_action_weight_init() __attribute__((constructor, used)); void __pangu_action_weight_init() @@ -1770,6 +1818,12 @@ static void http_hijack(const struct tfe_http_session * session, enum tfe_http_e ctx->action = PG_ACTION_NONE; return; } + if (try_send_by_token(hijack_size) <= 0) + { + TFE_LOG_ERROR(g_pangu_rt->local_logger, "No token is available to send data, profile_id = %d", param->profile_id); + ctx->action = PG_ACTION_NONE; + return; + } ctx->inject_sz = hijack_size; char cont_len_str[16]; @@ -1869,6 +1923,12 @@ static void http_insert(const struct tfe_stream * stream, const struct tfe_http_ ctx->action = PG_ACTION_NONE; return; } + if (try_send_by_token(ins_ctx->rule->inject_sz) <= 0) + { + TFE_LOG_ERROR(g_pangu_rt->local_logger, "No token is available to send data, profile_id = %d", param->profile_id); + ctx->action = PG_ACTION_NONE; + return; + } ctx->inject_sz = ins_ctx->rule->inject_sz; } else diff --git a/vendor/CMakeLists.txt b/vendor/CMakeLists.txt index 6cc452c..dd2c413 100644 --- a/vendor/CMakeLists.txt +++ b/vendor/CMakeLists.txt @@ -94,6 +94,22 @@ add_dependencies(nghttp2-static nghttp2) set_property(TARGET nghttp2-static PROPERTY IMPORTED_LOCATION ${INSTALL_DIR}/lib/libnghttp2.a) set_property(TARGET nghttp2-static APPEND PROPERTY INTERFACE_INCLUDE_DIRECTORIES ${INSTALL_DIR}/include;${INSTALL_DIR}/src/nghttp2/lib) +### ratelimiter +ExternalProject_Add(ratelimiter PREFIX ratelimiter + URL ${CMAKE_CURRENT_SOURCE_DIR}/ratelimiter-1.1.0-x86_64.tar.gz + URL_MD5 54cc36713c2aaa8ff378ca342778e57d + CONFIGURE_COMMAND "" + BUILD_COMMAND make + BUILD_IN_SOURCE 1) + +ExternalProject_Get_Property(ratelimiter INSTALL_DIR) +file(MAKE_DIRECTORY ${INSTALL_DIR}/include) + +add_library(ratelimiter-static STATIC IMPORTED GLOBAL) +add_dependencies(ratelimiter-static ratelimiter) +set_property(TARGET ratelimiter-static PROPERTY IMPORTED_LOCATION ${INSTALL_DIR}/lib/libratelimiter.a) +set_property(TARGET ratelimiter-static PROPERTY INTERFACE_INCLUDE_DIRECTORIES ${INSTALL_DIR}/include) + #### GoogleTest ExternalProject_Add(googletest PREFIX googletest URL ${CMAKE_CURRENT_SOURCE_DIR}/googletest-release-1.8.0.tar.gz diff --git a/vendor/ratelimiter-1.1.0-x86_64.tar.gz b/vendor/ratelimiter-1.1.0-x86_64.tar.gz new file mode 100644 index 0000000..1435650 Binary files /dev/null and b/vendor/ratelimiter-1.1.0-x86_64.tar.gz differ