TSG-7392: 限速dns、wechat协议时，有限速效果没有但是安全日志

3 files changed, 66 insertions, 38 deletions

diff --git a/src/tsg_action.cpp b/src/tsg_action.cpp
index 8ca7eda..b70d23d 100644
--- a/src/tsg_action.cpp
+++ b/src/tsg_action.cpp
@@ -35,6 +35,17 @@ static int set_drop_stream(const struct streaminfo *a_stream)
 	return STATE_DROPME|STATE_DROPPKT;				
 }
 
+static int set_ratelimit_flag(const struct streaminfo *a_stream)
+{	
+	struct master_context *_context=(struct master_context *)get_struct_project(a_stream, g_tsg_para.context_project_id);
+	if(_context!=NULL)
+	{
+		_context->is_ratelimit=1;
+	}
+
+	return 0;
+}
+
 static int get_http_header(char *buff, int len, int code, char *user_define)
 {
 	int used_len=0;
@@ -516,15 +527,22 @@ static unsigned char do_action_drop(const struct streaminfo *a_stream, Maat_rule
 	return STATE_DROPME|STATE_DROPPKT;
 }
 
-static unsigned char do_action_ratelimit(const struct streaminfo *a_stream, Maat_rule_t *p_result, struct compile_user_region *user_region)
+static unsigned char do_action_ratelimit(const struct streaminfo *a_stream, Maat_rule_t *p_result, struct compile_user_region *user_region, enum ACTION_RETURN_TYPE type)
 {
 	struct tcpall_context *context=NULL;
 	struct leaky_bucket *bucket=create_bucket((double)((user_region->deny->bps)/1000000), user_region->deny->bps, a_stream->threadnum);
 	tsg_set_bucket_to_tcpall(a_stream, &context, bucket, a_stream->threadnum);
 
+	set_ratelimit_flag(a_stream);
+
 	context=NULL;
-	
-	return STATE_GIVEME;
+
+	if(type==ACTION_RETURN_TYPE_PROT)
+	{		
+		return STATE_DROPME;
+	}
+
+	return STATE_GIVEME|STATE_KILL_OTHER;
 }
 
 static unsigned char do_action_block_mail(const struct streaminfo *a_stream, Maat_rule_t *p_result, struct compile_user_region *user_region)
@@ -681,7 +699,7 @@ unsigned char tsg_deal_deny_action(const struct streaminfo *a_stream, Maat_rule_
 			local_state=do_action_redirect_xxx( a_stream, p_result, user_region, protocol, user_data);
 			break;
 		case	TSG_METHOD_TYPE_RATE_LINIT:
-			local_state=do_action_ratelimit(a_stream, p_result, user_region);
+			local_state=do_action_ratelimit(a_stream, p_result, user_region, type);
 			break;
 		default:
 			break;
@@ -693,7 +711,8 @@ unsigned char tsg_deal_deny_action(const struct streaminfo *a_stream, Maat_rule_
 	state|=((type==ACTION_RETURN_TYPE_PROT) ? (local_state&STATE_DROPME ? PROT_STATE_DROPME : 0) : (local_state&STATE_DROPME ? APP_STATE_DROPME : 0));	
 	state|=((type==ACTION_RETURN_TYPE_PROT) ? (local_state&STATE_GIVEME ? PROT_STATE_GIVEME: 0) : (local_state&STATE_GIVEME ? APP_STATE_GIVEME: 0));
 	state|=((type==ACTION_RETURN_TYPE_PROT) ? (local_state&STATE_DROPPKT ? PROT_STATE_DROPPKT : 0) : (local_state&STATE_DROPPKT ? APP_STATE_DROPPKT: 0));
-
+	state|=((type==ACTION_RETURN_TYPE_PROT) ? (0) : (local_state&STATE_KILL_OTHER ? APP_STATE_KILL_OTHER : 0));
+	
 	return state;
 }
 
diff --git a/src/tsg_entry.cpp b/src/tsg_entry.cpp
index ee946ed..9fb167c 100644
--- a/src/tsg_entry.cpp
+++ b/src/tsg_entry.cpp
@@ -1243,7 +1243,7 @@ static unsigned char master_deal_scan_result(const struct streaminfo *a_stream,
 		{
 			case	TSG_ACTION_DENY:
 				state=tsg_deal_deny_action(a_stream, p_result, context->proto, ACTION_RETURN_TYPE_APP, a_packet);
-				if((state&APP_STATE_DROPPKT)==APP_STATE_DROPPKT)
+				if((state&APP_STATE_DROPPKT)==APP_STATE_DROPPKT || (state&APP_STATE_KILL_OTHER))
 				{
 					context->hit_cnt=0;
 					master_send_log(a_stream, p_result, 1, context->domain, context->proto, a_stream->threadnum);					
@@ -1486,6 +1486,11 @@ static unsigned char tsg_master_data_entry(const struct streaminfo *a_stream, vo
 		default:
 			break;
 	}
+
+	if(context->is_ratelimit==1 && a_stream->type==STREAM_TYPE_TCP)
+	{
+		state=APP_STATE_KILL_OTHER|APP_STATE_DROPME;
+	}
 	
 	if((a_stream->opstate==OP_STATE_CLOSE) || (state&APP_STATE_DROPME)==APP_STATE_DROPME)
 	{
@@ -1511,12 +1516,6 @@ static unsigned char tsg_master_all_entry(const struct streaminfo *a_stream, uns
 	struct Maat_rule_t result[MAX_RESULT_NUM]={0};	
 	struct mirrored_vlan vlan[MAX_RESULT_NUM]={0};
 	struct tcpall_context *context=(struct tcpall_context *)(*pme);
-
-	if(context==NULL)
-	{
-		context=(struct tcpall_context *)get_struct_project(a_stream, g_tsg_para.tcpall_project_id);
-		*pme=(void *)context;
-	}
 		
 	switch(stream_state)
 	{
@@ -1536,7 +1535,7 @@ static unsigned char tsg_master_all_entry(const struct streaminfo *a_stream, uns
 						ret=tsg_set_vlan_id_to_tcpall(a_stream, &context, vlan, vlan_num, thread_seq);
 						if(ret<=0)
 						{
-							break;
+							return state;
 						}
 						
 						*pme=(void *)(context);
@@ -1551,35 +1550,41 @@ static unsigned char tsg_master_all_entry(const struct streaminfo *a_stream, uns
 			scan_mid=NULL;
 		case	OP_STATE_DATA:
 		case	OP_STATE_CLOSE:
-				if(context==NULL || context->para==NULL)
+
+			break;
+		default:
+			break;
+	}
+
+	if(context==NULL)
+	{
+		context=(struct tcpall_context *)get_struct_project(a_stream, g_tsg_para.tcpall_project_id);
+		*pme=(void *)context;
+	}
+
+	if(context!=NULL && context->para!=NULL)
+	{
+		switch(context->method_type)
+		{
+			case	TSG_METHOD_TYPE_MIRRORED:
+				tsg_send_raw_packet(a_stream, context->vlan, context->vlan_num, thread_seq);
+				break;
+			case	TSG_METHOD_TYPE_RATE_LINIT:
+				eth_rawpkt_len=get_raw_packet_len(a_stream);
+				if(eth_rawpkt_len<=0)
 				{
 					break;
 				}
-				
-				switch(context->method_type)
+		
+				ret=is_permit_pass(eth_rawpkt_len*8, context->bucket, thread_seq);
+				if(ret==0)
 				{
-					case	TSG_METHOD_TYPE_MIRRORED:
-						tsg_send_raw_packet(a_stream, context->vlan, context->vlan_num, thread_seq);
-						break;
-					case	TSG_METHOD_TYPE_RATE_LINIT:
-						eth_rawpkt_len=get_raw_packet_len(a_stream);
-						if(eth_rawpkt_len<=0)
-						{
-							break;
-						}
-
-						ret=is_permit_pass(eth_rawpkt_len*8, context->bucket, thread_seq);
-						if(ret==0)
-						{
-							state=APP_STATE_GIVEME|APP_STATE_DROPPKT;
-						}
-						break;
-					default:
-						break;
+					state|=APP_STATE_GIVEME|APP_STATE_DROPPKT;
 				}
-			break;
-		default:
-			break;
+				break;
+			default:
+				break;
+		}
 	}
 
 	return state;
@@ -1610,7 +1615,10 @@ extern "C" unsigned char TSG_MASTER_UDP_ENTRY(const struct streaminfo *a_udp, vo
 	}
 	
 	state1=tsg_master_all_entry(a_udp, a_udp->opstate, (void **)&(context->all_entry), thread_seq, a_packet);
-	state2=tsg_master_data_entry(a_udp, (void **)&(context->data_entry), thread_seq, a_packet);
+	if(context->all_entry==NULL || context->all_entry->method_type!=TSG_METHOD_TYPE_RATE_LINIT)
+	{
+		state2=tsg_master_data_entry(a_udp, (void **)&(context->data_entry), thread_seq, a_packet);		
+	}
 
 	if(state1&APP_STATE_DROPME || state2&APP_STATE_DROPME || a_udp->opstate==OP_STATE_CLOSE)
 	{
diff --git a/src/tsg_entry.h b/src/tsg_entry.h
index ac6bb34..525e910 100644
--- a/src/tsg_entry.h
+++ b/src/tsg_entry.h
@@ -174,6 +174,7 @@ struct master_context
 	int hit_cnt;
 	int is_esni;
 	int is_log;
+	int is_ratelimit;
 	char *domain;
 	scan_status_t mid;	
 	struct Maat_rule_t *result;