TSG-11399 tfe支持TSG Script

2022-07-29 15:39:25 +08:00
parent b00782dff3
commit c4c03bfc4e
9 changed files with 1621 additions and 10 deletions
--- a/plugin/business/pangu-http/src/pangu_http.cpp
+++ b/plugin/business/pangu-http/src/pangu_http.cpp
@@ -1,6 +1,7 @@
 #include "pangu_logger.h"
 #include "edit_element.h"
 #include "pattern_replace.h"
+#include "http_lua.h"
 #include "pangu_web_cache.h"

 #include <tfe_proxy.h>
@@ -55,6 +56,7 @@ enum manipulate_action
 	MA_ACTION_HIJACK,
 	MA_ACTION_INSERT,
 	MA_ACTION_ELEMENT,
+	MA_ACTION_LUA_SCRIPT,
 	__MA_ACTION_MAX
 };

@@ -84,6 +86,7 @@ enum pangu_http_stat
 	STAT_ACTION_INSERT,
 	STAT_ACTION_INSERT_SZ,
 	STAT_ACTION_EDIT_ELEMENT,
+	STAT_ACTION_RUN_SCRIPT,
 	STAT_ACTION_WHITELSIT,
 	STAT_SUSPENDING,
 	__PG_STAT_MAX
@@ -94,6 +97,7 @@ enum manipulate_profile_table
 	POLICY_PROFLIE_TABLE_REJECT,
 	POLICY_PROFILE_TABLE_INSERT,
 	POLICY_PROFILE_TABLE_HIJACK,
+	POLICY_PROFILE_TABLE_LUA,
 	POLICY_PROFILE_TABLE_MAX
 };

@@ -106,6 +110,8 @@ struct manipulate_profile
 	char *profile_msg;
 	char *profile_type;
 	char *profile_position;
+	int *script_id;
+	int timeout;
 	ctemplate::Template * tpl;
 	pthread_mutex_t lock;
 };
@@ -153,6 +159,7 @@ struct pangu_rt
 	struct event_base* gc_evbase;
 	struct event* gcev;

+	struct tsg_lua_script lua_script;
 	Ratelimiter_handle_t ratelimiter;
 	int enable_rate;

@@ -189,6 +196,7 @@ static void pangu_http_stat_init(struct pangu_rt * pangu_runtime)
 	spec[STAT_ACTION_INSERT]="intcp_ins_num";
 	spec[STAT_ACTION_INSERT_SZ]="ins_bytes";
 	spec[STAT_ACTION_EDIT_ELEMENT]="intcp_edit_elem_num";
+	spec[STAT_ACTION_RUN_SCRIPT]="intcp_rus_num";
 	spec[STAT_ACTION_WHITELSIT]="intcp_allow_num";
 	spec[STAT_SUSPENDING]="suspending";

@@ -317,6 +325,7 @@ static enum manipulate_action manipulate_action_str2idx(const char *action_str)
 	clue_action_map[MA_ACTION_HIJACK]= "hijack";
 	clue_action_map[MA_ACTION_INSERT]= "insert";
 	clue_action_map[MA_ACTION_ELEMENT] = "edit_element";
+	clue_action_map[MA_ACTION_LUA_SCRIPT] = "run_script";

 	size_t i = 0;

@@ -517,6 +526,21 @@ void policy_action_param_new(int idx, const struct Maat_rule_t* rule, const char
 			}
 			param->e_rule = rule_id;
 			break;
+		case MA_ACTION_LUA_SCRIPT:
+			item=cJSON_GetObjectItem(json,"run_script_profile");
+			if(item && item->type==cJSON_Number){
+				param->profile_id = item->valueint;
+			}
+			item=cJSON_GetObjectItem(json,"enforcement_ratio");
+			if(item && item->type==cJSON_Number)
+			{
+				param->enforcement_ratio = item->valuedouble;
+			}
+			else
+			{
+				param->enforcement_ratio = 1;
+			}
+			break;
 		default: assert(0);
 			break;
 	}
@@ -699,6 +723,52 @@ void ma_hijack_profile_table_new_cb(int table_id, const char* key, const char* t
 	return;
 }

+void ma_lua_profile_table_new_cb(int table_id, const char* key, const char* table_line, MAAT_PLUGIN_EX_DATA* ad, long argl, void* argp)
+{
+	int timeout=0;
+	int ret=0, profile_id=0, is_valid=0;
+	char profile_name[128]={0}, profile_path[TFE_PATH_MAX]={0};
+
+	ret=sscanf(table_line, "%d\t%s\t%d\t%d", &profile_id, profile_path, &timeout, &is_valid);
+	if(ret!=4)
+	{
+		TFE_LOG_ERROR(g_pangu_rt->local_logger, "Policy table parse config failed: %s", table_line);
+		return;
+	}
+	struct manipulate_profile* ply_profile=ALLOC(struct manipulate_profile, 1);
+	ply_profile->ref_cnt=1;
+	pthread_mutex_init(&(ply_profile->lock), NULL);
+
+	ply_profile->profile_id=profile_id;
+	ply_profile->profile_name=tfe_strdup(profile_path);
+	ply_profile->profile_msg = tfe_read_file(profile_path, &ply_profile->msg_len);
+	ply_profile->timeout=timeout;
+	if (ply_profile->profile_msg == NULL)
+	{
+		TFE_LOG_ERROR(g_pangu_rt->local_logger, "Read file failed %d:%s:%s", profile_id, profile_name, profile_path);
+		*ad = ply_profile;
+		return;
+	}
+
+	int i=0, thread_num = g_pangu_rt->thread_num;
+	struct tsg_lua_script *lua_script = &(g_pangu_rt->lua_script);
+
+	if(lua_script->lua_is_cache==1)
+	{
+		ply_profile->script_id = ALLOC(int, thread_num);
+
+		for(i=0; i<thread_num; i++)
+		{
+			ply_profile->script_id[i]=http_lua_map_cache_script(lua_script->http_lua_handle[i], ply_profile->profile_msg, ply_profile->msg_len);
+		}
+	}
+
+	TFE_LOG_INFO(g_pangu_rt->local_logger, "Policy table add success %d", profile_id);
+
+	*ad = ply_profile;
+	return;
+}
+
 void ma_profile_table_free_cb(int table_id, MAAT_PLUGIN_EX_DATA* ad, long argl, void *argp)
 {
 	if(*ad==NULL)
@@ -716,8 +786,15 @@ void ma_profile_table_free_cb(int table_id, MAAT_PLUGIN_EX_DATA* ad, long argl,
 	pthread_mutex_unlock(&(ply_obj->lock));
 	pthread_mutex_destroy(&(ply_obj->lock));

-	FREE(&ply_obj->profile_type);
-	FREE(&ply_obj->profile_msg);
+	if(ply_obj->profile_type)
+		FREE(&ply_obj->profile_type);
+
+	if(ply_obj->profile_msg)
+		FREE(&ply_obj->profile_msg);
+
+	if(ply_obj->script_id)
+		FREE(&ply_obj->script_id);
+
 	FREE(&ply_obj->profile_name);
 	if (ply_obj->profile_position)
 		FREE(&ply_obj->profile_position);
@@ -759,7 +836,8 @@ const char* table_name_idx2str(int profile_idx)
 {
 	const char *table_name_map[] = {"TSG_PROFILE_RESPONSE_PAGES",
 									"PXY_PROFILE_INSERT_SCRIPTS",
-									"PXY_PROFILE_HIJACK_FILES"};
+									"PXY_PROFILE_HIJACK_FILES",
+									"PXY_PROFILE_RUN_SCRIPTS"};

 	return table_name_map[profile_idx];
 }
@@ -775,6 +853,7 @@ int maat_table_ex_init(int profile_idx,
 			[POLICY_PROFLIE_TABLE_REJECT] = ma_profile_table_new_cb,
 			[POLICY_PROFILE_TABLE_INSERT] = ma_insert_profile_table_new_cb,
 			[POLICY_PROFILE_TABLE_HIJACK] = ma_hijack_profile_table_new_cb,
+			[POLICY_PROFILE_TABLE_LUA]    = ma_lua_profile_table_new_cb,
 	};

 	table_id=g_pangu_rt->plolicy_table_id[profile_idx]=Maat_table_register(g_pangu_rt->maat, table_name);
@@ -843,7 +922,7 @@ int pangu_policy_init(const char* profile_path, const char* static_section, cons
 		goto error_out;
 	}

-	for(int i = 0; i <= POLICY_PROFILE_TABLE_HIJACK; i++)
+	for(int i = 0; i <= POLICY_PROFILE_TABLE_LUA; i++)
 	{
 		ret = maat_table_ex_init(i, ma_profile_table_free_cb, ma_profile_table_dup_cb);
 		if(ret<0)
@@ -934,6 +1013,13 @@ int pangu_http_init(struct tfe_proxy * proxy)
 	g_pangu_rt->ratelimiter=ratelimit_handle_create(profile_path, "ratelimit");

 	pangu_http_stat_init(g_pangu_rt);
+
+	g_pangu_rt->lua_script.http_lua_handle=http_lua_handle_create(g_pangu_rt->thread_num, "tfe");
+	if(g_pangu_rt->lua_script.http_lua_handle==NULL)
+	{
+		goto error_out;
+	}
+
 	if(pangu_policy_init(profile_path, "MAAT", "DYNAMIC_MAAT")<0)
 	{
 		goto error_out;
@@ -1030,6 +1116,7 @@ struct pangu_http_ctx
 	struct replace_ctx * rep_ctx;
 	struct insert_ctx * ins_ctx;
 	struct edit_element_ctx * edit_ctx;
+	struct tsg_script_ctx *tsg_ctx;
 	struct ip_data_ctx ip_ctx;

 	int (* resumed_cb)(const struct tfe_stream * stream,
@@ -1048,7 +1135,6 @@ struct pangu_http_ctx
 	int thread_id;
 };

-
 void http_repl_ctx_free(struct replace_ctx* rep_ctx)
 {
 	if (rep_ctx->http_body)
@@ -1088,6 +1174,23 @@ void http_element_ctx_free(struct edit_element_ctx *edit_ctx)
 	return;
 }

+void http_tsg_ctx_free(struct tsg_script_ctx *tsg_ctx, int thread_id)
+{
+	if (tsg_ctx->http_body)
+	{
+		evbuffer_free(tsg_ctx->http_body);
+		tsg_ctx->http_body = NULL;
+	}
+	if (tsg_ctx->http_lua_body)
+	{
+		evbuffer_free(tsg_ctx->http_lua_body);
+		tsg_ctx->http_lua_body = NULL;
+	}
+	http_lua_ctx_free(&g_pangu_rt->lua_script, thread_id, tsg_ctx->lua_ctx);
+	FREE(&tsg_ctx);
+	return;
+}
+
 void http_ip_ctx_free(struct ip_data_ctx *ip_ctx)
 {
 	if(ip_ctx->asn_client)
@@ -1128,6 +1231,11 @@ static void pangu_http_ctx_free(struct pangu_http_ctx * ctx)
 		http_element_ctx_free(ctx->edit_ctx);
 		ctx->edit_ctx = NULL;
 	}
+	if(ctx->tsg_ctx)
+	{
+		http_tsg_ctx_free(ctx->tsg_ctx, ctx->thread_id);
+		ctx->tsg_ctx = NULL;
+	}

 	http_ip_ctx_free(&ctx->ip_ctx);
 	ctx->manipulate_replaced=0;
@@ -1403,6 +1511,172 @@ static int http_enforcement_ratio(float enforcement_ratio)
    return 0;
 }

+int http_lua_profile(int profile_id, char **profile_msg, size_t *msg_len, int **script_id, int *timeout)
+{
+	int ret = 0;
+
+	struct manipulate_profile* lua_profile=get_profile_by_id(POLICY_PROFILE_TABLE_LUA, profile_id);
+	if(lua_profile==NULL)
+	{
+		ret=-1;
+		return ret;
+	}
+
+	*profile_msg = tfe_strdup(lua_profile->profile_msg);
+	*msg_len = lua_profile->msg_len;
+	*script_id = lua_profile->script_id;
+	*timeout=lua_profile->timeout;
+
+	ma_profile_table_free(lua_profile);
+	lua_profile = NULL;
+	return ret;
+}
+
+void http_lua(const struct tfe_stream * stream, const struct tfe_http_session * session, enum tfe_http_event events,
+	const unsigned char * body_frag, size_t frag_size, struct pangu_http_ctx * ctx)
+{
+	int ret = 0;
+	struct tfe_http_session * to_write_sess = NULL;
+
+	struct tsg_lua_script *lua_script=&g_pangu_rt->lua_script;
+	lua_script->http_lua_profile=http_lua_profile;
+
+	to_write_sess = tfe_http_session_allow_write(session);
+	if (to_write_sess == NULL)    //fail to wirte, abandon.
+	{
+		TFE_STREAM_LOG_INFO(stream, "tfe_http_session_allow_write() %s failed.", session->req->req_spec.uri);
+		ctx->action = PG_ACTION_NONE;
+		tfe_http_session_detach(session); return;
+	}
+
+	struct tsg_script_ctx *tsg_ctx = ctx->tsg_ctx;
+	if (ctx->tsg_ctx == NULL)
+	{
+		if ((events & EV_HTTP_REQ_HDR) || (events & EV_HTTP_RESP_HDR))
+		{
+			struct policy_action_param *param = ctx->param;
+			ctx->tsg_ctx = tsg_ctx = ALLOC(struct tsg_script_ctx, 1);
+			tsg_ctx->profile_id = param->profile_id;
+			tsg_ctx->lua_ctx = http_lua_ctx_new(lua_script, ctx->thread_id);
+		}
+		else
+		{
+			TFE_STREAM_LOG_INFO(stream, "Can only setup replace on REQ/RESP headers, detached.");
+			ctx->action = PG_ACTION_NONE;
+			tfe_http_session_detach(session); return;
+		}
+	}
+	tsg_ctx->events = events;
+	tsg_ctx->session = session;
+	tsg_ctx->local_logger = g_pangu_rt->local_logger;
+	tsg_ctx->config_id = ctx->enforce_rules[0].config_id;
+
+	struct tfe_http_half * in_req_half = session->req;
+	struct tfe_http_half * in_resp_half = session->resp;
+	struct tfe_http_req_spec * in_req_spec = &in_req_half->req_spec;
+	struct tfe_http_resp_spec * in_resp_spec = &in_resp_half->resp_spec;
+
+	if ((events & EV_HTTP_REQ_HDR) || (events & EV_HTTP_RESP_HDR))
+	{
+		if (tfe_http_in_request(events))
+		{
+			tsg_ctx->http_req_uri=1; tsg_ctx->execut_lua_sucess=0;
+			ret=execute_lua_script_rule(lua_script, tsg_ctx->profile_id, tsg_ctx->lua_ctx, ctx->thread_id, (void *)tsg_ctx);
+			if(ret>0 && tsg_ctx->execut_lua_sucess==1)
+			{
+				tsg_ctx->actually_executed =1;
+			}
+			tsg_ctx->http_req_uri=0;
+			tsg_ctx->replacing = tfe_http_session_request_create(to_write_sess, in_req_spec->method, in_req_spec->uri);
+			tfe_http_session_request_set(to_write_sess, tsg_ctx->replacing);
+		}
+		else
+		{
+			tsg_ctx->replacing = tfe_http_session_response_create(to_write_sess, in_resp_spec->resp_code);
+			tfe_http_session_response_set(to_write_sess, tsg_ctx->replacing);
+		}
+
+		tsg_ctx->execut_lua_sucess=0;
+		ret=execute_lua_script_rule(lua_script, tsg_ctx->profile_id, tsg_ctx->lua_ctx, ctx->thread_id, (void *)tsg_ctx);
+		if(ret>0 && tsg_ctx->actually_executed==1)
+		{
+			tsg_ctx->actually_executed =1;
+		}
+
+		struct tfe_http_half * in_half = tfe_http_in_request(events) ? in_req_half : in_resp_half;
+
+		struct http_field_name in_header_field{};
+		const char * in_header_value = NULL;
+		void * iterator = NULL;
+
+		while (true)
+		{
+			if ((in_header_value = tfe_http_field_iterate(in_half, &iterator, &in_header_field)) == NULL)
+			{
+				break;
+			}
+			tfe_http_field_write(tsg_ctx->replacing, &in_header_field, in_header_value);
+		}
+	}
+
+	if ((events & EV_HTTP_REQ_BODY_BEGIN) || (events & EV_HTTP_RESP_BODY_BEGIN))
+	{
+		assert(tsg_ctx->http_body == NULL);
+		tsg_ctx->http_body = evbuffer_new();
+	}
+
+	if ((events & EV_HTTP_REQ_BODY_CONT) || (events & EV_HTTP_RESP_BODY_CONT))
+	{
+		evbuffer_add(tsg_ctx->http_body, body_frag, frag_size);
+	}
+
+	if ((events & EV_HTTP_REQ_BODY_END) || (events & EV_HTTP_RESP_BODY_END))
+	{
+		char * __http_body = (char *) evbuffer_pullup(tsg_ctx->http_body, -1);
+		size_t __http_body_len = evbuffer_get_length(tsg_ctx->http_body);
+
+		ret=execute_lua_script_rule(lua_script, tsg_ctx->profile_id, tsg_ctx->lua_ctx, ctx->thread_id, (void *)tsg_ctx);
+
+		char * __http_lua_body = NULL; size_t __http_body_lua_len = 0;
+		if(ret > 0 && tsg_ctx->http_lua_body != NULL)
+		{
+			tsg_ctx->actually_executed =1;
+			__http_lua_body = (char *) evbuffer_pullup(tsg_ctx->http_lua_body, -1);
+			__http_body_lua_len = evbuffer_get_length(tsg_ctx->http_lua_body);
+		}
+
+		if (__http_body_lua_len >0)
+		{
+			tfe_http_half_append_body(tsg_ctx->replacing, __http_lua_body, __http_body_lua_len, 0);
+			tsg_ctx->actually_executed=1;
+		}
+		else
+		{
+			tfe_http_half_append_body(tsg_ctx->replacing, __http_body, __http_body_len, 0);
+		}
+
+		if (tsg_ctx->http_lua_body != NULL)
+		{
+			evbuffer_free(tsg_ctx->http_lua_body);
+			tsg_ctx->http_lua_body = NULL;
+		}
+
+		if (tsg_ctx->http_body != NULL)
+		{
+			evbuffer_free(tsg_ctx->http_body);
+			tsg_ctx->http_body = NULL;
+		}
+	}
+
+	if ((events & EV_HTTP_REQ_END) || (events & EV_HTTP_RESP_END))
+	{
+		tfe_http_half_append_body(tsg_ctx->replacing, NULL, 0, 0);
+		tsg_ctx->replacing = NULL;
+	}
+
+	return;
+}
+
 void http_replace(const struct tfe_stream * stream, const struct tfe_http_session * session, enum tfe_http_event events,
 	const unsigned char * body_frag, size_t frag_size, struct pangu_http_ctx * ctx)
 {
@@ -2185,6 +2459,9 @@ static void http_manipulate(const struct tfe_stream * stream, const struct tfe_h
 		case MA_ACTION_ELEMENT:
 			http_element(stream, session, events, body_frag, frag_size, ctx);
 			break;
+		case MA_ACTION_LUA_SCRIPT:
+			http_lua(stream, session, events, body_frag, frag_size, ctx);
+			break;
 		default: assert(0);
 			break;
 	}
@@ -2702,6 +2979,21 @@ static inline int ctx_actually_replaced(struct pangu_http_ctx * ctx)
 	}
 }

+static inline int ctx_actually_ran_script(struct pangu_http_ctx * ctx)
+{
+
+	if(ctx->action == PG_ACTION_MANIPULATE &&
+	   ctx->param->action == MA_ACTION_LUA_SCRIPT &&
+	   ctx->tsg_ctx->actually_executed==1)
+	{
+		return 1;
+	}
+	else
+	{
+		return 0;
+	}
+}
+
 static inline int ctx_actually_inserted(struct pangu_http_ctx * ctx)
 {

@@ -2791,7 +3083,7 @@ void pangu_on_http_end(const struct tfe_stream * stream,

 	if(ctx->action != PG_ACTION_NONE &&
 	   (((ctx_actually_replaced(ctx)) || (ctx_actually_inserted(ctx)) || (ctx_actually_edited(ctx)) ||
-	   (ctx_actually_manipulate(ctx))) || (ctx->action == PG_ACTION_MONIT ||
+	   (ctx_actually_manipulate(ctx)) || ctx_actually_ran_script(ctx)) || (ctx->action == PG_ACTION_MONIT ||
 	   	ctx->action == PG_ACTION_REJECT || ctx->action == PG_ACTION_WHITELIST)))
 	{
 		ret=pangu_send_log(g_pangu_rt->send_logger, &log_msg);
@@ -2811,6 +3103,11 @@ void pangu_on_http_end(const struct tfe_stream * stream,
 		ATOMIC_INC(&(g_pangu_rt->stat_val[STAT_ACTION_REPLACE]));
 	}

+	if(ctx->tsg_ctx && ctx->tsg_ctx->actually_executed==1)
+	{
+		ATOMIC_INC(&(g_pangu_rt->stat_val[STAT_ACTION_RUN_SCRIPT]));
+	}
+
 	TFE_LOG_DEBUG(g_pangu_rt->local_logger, "cache %s %s upload=%d",
 											session->req->req_spec.url,
 											cache_pending_result_string(ctx->pending_result),