TSG-14484 Pxoxy支持Maat4

1 files changed, 76 insertions, 2 deletions

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)