1 files changed, 21 insertions, 11 deletions

diff --git a/enforcer/security/security_enforcer.cpp b/enforcer/security/security_enforcer.cpp
index 309a453..7dec5c9 100644
--- a/enforcer/security/security_enforcer.cpp
+++ b/enforcer/security/security_enforcer.cpp
@@ -14,10 +14,10 @@
 #include <stellar/module.h>
 #include <stellar/session.h>
 #include <stellar/packet.h>
+#include <stellar/security_enforcer.h>
 
 #include "bucket.h"
 #include "security_maat.h"
-#include "security_enforcer.h"
 
 #define	ENFORCER_RULE_UUID_NUM	128
 
@@ -353,7 +353,7 @@ void security_enforcer_enforce_drop(struct security_enforcer_env *enforcer_env,
     }
 
 	packet_set_action((struct packet *)rawpkt, PACKET_ACTION_DROP);
-	session_set_discard((struct session *)sess);
+	session_manager_discard_session(enforcer_env->sess_mgr, module_manager_get_thread_id(enforcer_env->mod_mgr), (struct session *)sess);
 
     // if(drop->send_icmp_enable)
     // {
@@ -489,7 +489,7 @@ void security_enforcer_enforce_mail_block(struct security_enforcer_env *enforcer
 	}
 
     packet_set_action((struct packet *)rawpkt, PACKET_ACTION_DROP);
-	session_set_discard((struct session *)sess);
+	session_manager_discard_session(enforcer_env->sess_mgr, module_manager_get_thread_id(enforcer_env->mod_mgr), (struct session *)sess);
 }
 
 void policy_user_define_variable_replace(ctemplate::TemplateDictionary *tpl_dict, uuid_t rule_uuid, char *client_ip, char *subscriber)
@@ -556,7 +556,7 @@ void security_enforcer_enforce_http_block(struct security_enforcer_env *enforcer
 	}
 
     packet_set_action((struct packet *)rawpkt, PACKET_ACTION_DROP);
-	session_set_discard((struct session *)sess);
+	session_manager_discard_session(enforcer_env->sess_mgr, module_manager_get_thread_id(enforcer_env->mod_mgr), (struct session *)sess);
 
 	uint16_t http_hdr_len=0;
 	char http_hdr[512]={0};
@@ -967,7 +967,7 @@ void security_enforcer_enforce_http_redirect(struct security_enforcer_env *enfor
 		struct packet *redirect_pkt=packet_manager_build_tcp_packet(enforcer_env->pkt_mgr, module_manager_get_thread_id(enforcer_env->mod_mgr), c2s_origin_pkt, th_seq, th_ack, TH_PUSH, NULL, 0, payload, payload_offset);
 		packet_manager_schedule_packet(enforcer_env->pkt_mgr, module_manager_get_thread_id(enforcer_env->mod_mgr), redirect_pkt, PACKET_STAGE_POSTROUTING);
 		packet_set_action((struct packet *)rawpkt, PACKET_ACTION_DROP);
-		session_set_discard((struct session *)sess);
+		session_manager_discard_session(enforcer_env->sess_mgr, module_manager_get_thread_id(enforcer_env->mod_mgr), (struct session *)sess);
 	}
 }
 
@@ -1076,10 +1076,15 @@ size_t maat_state_compile(struct maat_state *state __attribute__((unused)), cons
 	return 0;
 }
 
-void security_enforcer_packet_based_node_callback(const struct packet *rawpkt, void *arg)
+void packet_based_security_enforcer_node_entry(struct packet *rawpkt, struct module *mod_enforcer)
 {
-	struct security_enforcer_env *enforcer_env=(struct security_enforcer_env *)arg;
-	if(enforcer_env==NULL || rawpkt==NULL)
+	if(rawpkt==NULL || mod_enforcer==NULL)
+	{
+		return ;
+	}
+
+	struct security_enforcer_env *enforcer_env=(struct security_enforcer_env *)module_get_ctx(mod_enforcer);
+	if(enforcer_env==NULL)
 	{
 		return ;
 	}
@@ -1107,10 +1112,15 @@ void security_enforcer_packet_based_node_callback(const struct packet *rawpkt, v
 	// packet_tag_get(rawpkt, &pkt_tag_key_bits, &pkt_tag_val_bits);
 }
 
-void security_enforcer_session_based_node_callback(const struct packet *rawpkt, void *arg)
+void session_based_security_enforcer_node_entry(struct packet *rawpkt, struct module *mod_enforcer)
 {
-	struct security_enforcer_env *enforcer_env=(struct security_enforcer_env *)arg;
-	if(enforcer_env==NULL || rawpkt==NULL)
+	if(rawpkt==NULL || mod_enforcer==NULL)
+	{
+		return ;
+	}
+
+	struct security_enforcer_env *enforcer_env=(struct security_enforcer_env *)module_get_ctx(mod_enforcer);
+	if(enforcer_env==NULL)
 	{
 		return ;
 	}