TSG-1818 insert和hijack增加注入流量控制

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

plugin/business/pangu-http/src/pangu_http.cpp

@@ -25,6 +25,7 @@
 #include <string.h>
 #include <sys/types.h>
 
+#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

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