OMPUB-740 Lua插件支持获取五元组和重写HTTP头信息feature-OMPUB-740

3 files changed, 90 insertions, 12 deletions

diff --git a/plugin/business/tsg-http/include/http_lua.h b/plugin/business/tsg-http/include/http_lua.h
index 5f506cb..060fcec 100644
--- a/plugin/business/tsg-http/include/http_lua.h
+++ b/plugin/business/tsg-http/include/http_lua.h
@@ -19,11 +19,13 @@ struct tsg_script_ctx
     int config_id;
 	int profile_id;
     int http_req_uri;
+    int rewrite_header;
     char *rewrite_uri;
     int execut_lua_sucess;
     struct elua_context *elua_ctx;
     enum tfe_http_event events;
     void * local_logger;
+    struct tfe_stream_addr * addr;
     const struct tfe_http_session *session;
 	struct tfe_http_half *replacing;
 	struct evbuffer *http_body;
diff --git a/plugin/business/tsg-http/src/http_lua.cpp b/plugin/business/tsg-http/src/http_lua.cpp
index b7bfb82..05cbb1f 100644
--- a/plugin/business/tsg-http/src/http_lua.cpp
+++ b/plugin/business/tsg-http/src/http_lua.cpp
@@ -449,6 +449,35 @@ static int http_lua_get_current_stage(struct elua_vm *vm)
 	return 1;
 }
 
+static int http_lua_get_5tuple(struct elua_vm *vm)
+{
+	struct tsg_script_ctx *tsg_ctx = (struct tsg_script_ctx *)elua_get_execute_userdata(vm);
+	if(tsg_ctx == NULL)
+	{
+		return 0;
+	}
+    struct tfe_stream_addr * addr = tsg_ctx->addr;
+	if(addr == NULL)
+	{
+		return 0;
+	}
+
+	char ip_addr[64]={0};
+	unsigned int source=0,dest=0,protocol;
+	char src_ip_str[MAX(INET6_ADDRSTRLEN,INET_ADDRSTRLEN)] = {0};
+	char dst_ip_str[MAX(INET6_ADDRSTRLEN,INET_ADDRSTRLEN)] = {0};
+
+	protocol = addr->addrtype;
+	source = ntohs(addr->tuple4_v4->source);
+	dest = ntohs(addr->tuple4_v4->dest);
+	inet_ntop(AF_INET, &addr->tuple4_v4->saddr, src_ip_str, sizeof(src_ip_str));
+	inet_ntop(AF_INET, &addr->tuple4_v4->daddr, dst_ip_str, sizeof(dst_ip_str));
+	snprintf(ip_addr, sizeof(ip_addr), "%d %s %d %s %d", protocol, src_ip_str, source, dst_ip_str, dest);
+	http_set_string_to_lua(vm, ip_addr, strlen(ip_addr));
+
+	return 1;
+}
+
 static int http_lua_get_headers(struct elua_vm *vm)
 {
 	struct tsg_script_ctx *tsg_ctx = (struct tsg_script_ctx *)elua_get_execute_userdata(vm);
@@ -594,6 +623,48 @@ static int http_lua_set_headers(struct elua_vm *vm)
 	return 0;
 }
 
+static int http_lua_rewrite_header(struct elua_vm *vm)
+{
+	struct tsg_script_ctx *tsg_ctx = (struct tsg_script_ctx *)elua_get_execute_userdata(vm);
+	if(tsg_ctx == NULL)
+	{
+		return 0;
+	}
+	const struct tfe_http_session * session = tsg_ctx->session;
+	if(session == NULL)
+	{
+		return 0;
+	}
+	struct tfe_http_half *replacing = tsg_ctx->replacing;
+	if(replacing == NULL)
+	{
+		return 0;
+	}
+
+	int out_lua_argc = 0;
+	struct elua_data *out_lua_argv = NULL;
+
+	out_lua_argc = http_get_param_from_lua(vm, &out_lua_argv);
+	if(out_lua_argc != 2 || out_lua_argv == NULL)
+	{
+		return 0;
+	}
+
+	char *field_name=out_lua_argv[0].buff, *field_value=out_lua_argv[1].buff;
+
+	if(field_name == NULL || field_value == NULL)
+	{
+		return 0;
+	}
+
+	tfe_http_nonstd_field_write(tsg_ctx->replacing, field_name, field_value);
+	tsg_ctx->execut_lua_sucess=1;
+	tsg_ctx->rewrite_header=1;
+
+	http_free_params(out_lua_argv);
+	return 0;
+}
+
 static int http_lua_get_body(struct elua_vm *vm)
 {
 	struct tsg_script_ctx *tsg_ctx = (struct tsg_script_ctx *)elua_get_execute_userdata(vm);
@@ -669,6 +740,7 @@ void http_lua_ctx_free(struct tsg_lua_script *lua_script, unsigned int thread_id
 void http_lua_inject_http_consts(struct elua_vm *vm)
 {
 	elua_register_cbinding(vm, NULL, "get_current_stage", http_lua_get_current_stage);
+	elua_register_cbinding(vm, NULL, "get_5tuple",        http_lua_get_5tuple);
 	elua_register_cbinding(vm, NULL, "log_debug", 	      http_lua_log_debug);
 	elua_register_cbinding(vm, NULL, "log_info", 		  http_lua_log_info);
 	elua_register_cbinding(vm, NULL, "log_error", 	      http_lua_log_error);
@@ -682,6 +754,7 @@ void http_lua_inject_req_header_api(struct elua_vm *vm)
 	elua_register_cbinding(vm, "req",  "set_uri",       http_lua_set_uri);
 	elua_register_cbinding(vm, "req",  "get_headers",   http_lua_get_headers);
 	elua_register_cbinding(vm, "req",  "set_header",    http_lua_set_headers);
+	elua_register_cbinding(vm, "req",  "rewrite_header", http_lua_rewrite_header);
 }
 
 void http_lua_inject_req_body_api(struct elua_vm *vm)
@@ -695,8 +768,9 @@ void http_lua_inject_resp_header_api(struct elua_vm *vm)
 	elua_register_cbinding(vm, "resp", "get_status_code",   http_lua_get_status_code);
 	elua_register_cbinding(vm, "resp", "set_status_code",   http_lua_set_status_code);
 
-	elua_register_cbinding(vm, "resp", "get_headers",   http_lua_get_headers);
-	elua_register_cbinding(vm, "resp", "set_header",    http_lua_set_headers);
+	elua_register_cbinding(vm, "resp", "get_headers",    http_lua_get_headers);
+	elua_register_cbinding(vm, "resp", "set_header",     http_lua_set_headers);
+	elua_register_cbinding(vm, "resp", "rewrite_header", http_lua_rewrite_header);
 }
 
 void http_lua_inject_resp_body_api(struct elua_vm *vm)
diff --git a/plugin/business/tsg-http/src/tsg_http.cpp b/plugin/business/tsg-http/src/tsg_http.cpp
index 8c20465..d88a6f4 100644
--- a/plugin/business/tsg-http/src/tsg_http.cpp
+++ b/plugin/business/tsg-http/src/tsg_http.cpp
@@ -1644,6 +1644,7 @@ void http_lua(const struct tfe_stream * stream, const struct tfe_http_session *
 			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->addr = stream->addr;
 			tsg_ctx->elua_ctx = http_lua_ctx_new(lua_script, ctx->thread_id);
 		}
 		else
@@ -1665,16 +1666,17 @@ void http_lua(const struct tfe_stream * stream, const struct tfe_http_session *
 
 	if ((events & EV_HTTP_REQ_HDR) || (events & EV_HTTP_RESP_HDR))
 	{
+		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->elua_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->execut_lua_sucess=0;
+
 		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->elua_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->execut_lua_sucess=0;
 			tsg_ctx->replacing = tfe_http_session_request_create(to_write_sess, in_req_spec->method, tsg_ctx->rewrite_uri !=NULL ? tsg_ctx->rewrite_uri : in_req_spec->uri);
 			tfe_http_session_request_set(to_write_sess, tsg_ctx->replacing);
 		}
@@ -1694,14 +1696,13 @@ void http_lua(const struct tfe_stream * stream, const struct tfe_http_session *
 		{
 			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)
+		while (true && tsg_ctx->rewrite_header!=1)
 		{
 			if ((in_header_value = tfe_http_field_iterate(in_half, &iterator, &in_header_field)) == NULL)
 			{
@@ -1709,6 +1710,7 @@ void http_lua(const struct tfe_stream * stream, const struct tfe_http_session *
 			}
 			tfe_http_field_write(tsg_ctx->replacing, &in_header_field, in_header_value);
 		}
+		tsg_ctx->rewrite_header=0;
 	}
 
 	if ((events & EV_HTTP_REQ_BODY_BEGIN) || (events & EV_HTTP_RESP_BODY_BEGIN))