重构部分有风险的内存操作。

3 files changed, 33 insertions, 42 deletions

diff --git a/kni_comm.c b/kni_comm.c
index b3dd201..509b880 100644
--- a/kni_comm.c
+++ b/kni_comm.c
@@ -602,6 +602,7 @@ int kni_get_service_defined(int new_action,struct Maat_rule_t* maat_result,struc
 
 		pmeinfo->cfg_id = maat_result->config_id;
 		pmeinfo->ser_def_len = maat_result->serv_def_len;
+		assert(sizeof(pmeinfo->service_defined) > maat_result->serv_def_len);
 		memcpy(pmeinfo->service_defined,maat_result->service_defined,maat_result->serv_def_len);
 	}
 	
diff --git a/kni_entry.h b/kni_entry.h
index 1e5b51f..5ea1026 100644
--- a/kni_entry.h
+++ b/kni_entry.h
@@ -13,7 +13,8 @@
 #include "kni_intercept.h"
 #include "kni_ratelimit.h"
 
-
+#define ALLOC(type, number) ((type *)calloc(sizeof(type), number))
+#define	FREE(p)	{free(*p);*p=NULL;}
 
 #ifndef TH_FIN
 #define TH_FIN	0x01
diff --git a/kni_intercept.c b/kni_intercept.c
index 09fe08f..20059ee 100644
--- a/kni_intercept.c
+++ b/kni_intercept.c
@@ -23,50 +23,44 @@ extern "C" int sendpacket_build_ethernet(unsigned char* dst,unsigned char* src,u
 extern "C" unsigned char MESA_dir_reverse(unsigned char route_dir);
 
 
-
-int kni_set_tlvinfo(char* buf,int buflen,struct kni_repaired_fds datainfo)
+size_t add_option(char* buff, size_t size, uint16_t opt_type, uint16_t opt_len, char* opt_cont)
+{
+	if(size<opt_len+sizeof(uint16_t)*2)
+	{
+		return 0;
+	}
+	*((uint16_t*)buff)=opt_type;
+	*((uint16_t*)(buff+sizeof(uint16_t))=opt_len;
+	memcpy(buff+sizeof(uint16_t)*2, opt_cont, opt_len);
+	return opt_len+sizeof(uint16_t)*2;
+}
+int kni_set_tlvinfo(char* buf, int buflen, struct kni_lqueue_datainfo datainfo)
 {
 	int tlv_len = 0;
 
-	struct kni_tlv_header header_info;
-	struct kni_tlv_info protocol_info;
-	struct kni_tlv_info keyring_info;
-
-	header_info.magic = 0x4d5a;
-	header_info.counts = 2;
-
-	protocol_info.type = KNI_TLV_TYPE_PROTOCOL;
-	protocol_info.len = sizeof(int);
-
-	keyring_info.type =  KNI_TLV_TYPE_KEYRING_ID;
-	keyring_info.len = sizeof(int);
+	struct kni_tlv_header *header_info=(struct kni_tlv_header *)buf;
+	
+	header_info->magic = 0x4d5a;
+	header_info->counts = 2;
 
-	memcpy(buf+tlv_len,&header_info,sizeof(struct kni_tlv_header));
 	tlv_len += sizeof(struct kni_tlv_header);
 
-	memcpy(buf+tlv_len,&protocol_info,sizeof(struct kni_tlv_info));
-	tlv_len += sizeof(struct kni_tlv_info);
-	memcpy(buf+tlv_len,&(datainfo.protocol),protocol_info.len);
-	tlv_len += protocol_info.len;
+	tlv_len+=add_option(buf+tlv_len,buflen-tlv_len, KNI_TLV_TYPE_PROTOCOL, (uint16_t)sizeof(int), &(datainfo.protocol));
+	tlv_len+=add_option(buf+tlv_len,buflen-tlv_len, KNI_TLV_TYPE_KEYRING_ID, (uint16_t)sizeof(int), &(datainfo.keyring));
 
-	memcpy(buf+tlv_len,&keyring_info,sizeof(struct kni_tlv_info));
-	tlv_len += sizeof(struct kni_tlv_info);
-	memcpy(buf+tlv_len,&(datainfo.keyring),keyring_info.len);
-	tlv_len += keyring_info.len;
-	
+	assert(tlv_len=<buflen);
 
 	return tlv_len;
 
 }
 
-
 int kni_send_fds(int socket,struct kni_repaired_fds to_send_fds)  
 {
 	int flags=MSG_NOSIGNAL;
     struct msghdr msg = {0};
     struct cmsghdr *cmsg;
-    char buf[CMSG_SPACE(KNI_SENDFD_NUM * sizeof(int))], dup[256]={0};
-    memset(buf, 0, sizeof(buf));
+    char buf[CMSG_SPACE(KNI_SENDFD_NUM * sizeof(int))]={0}, dup[256]={0};
+	
     struct iovec io = { .iov_base = &dup, .iov_len = sizeof(dup) };
 
 	int dup_len = 256;
@@ -448,6 +442,7 @@ int kni_sendpkt_sockraw(int thread_seq,int iplen,char* ip,struct stream_tuple4_v
 	if(ifname_len<sizeof(ifr.ifr_name))
 	{
 		memset(ifr.ifr_name,0,IFNAMSIZ);
+		assert(ifname_len<=IFNAMSIZ);
 		memcpy(ifr.ifr_name,if_name,ifname_len);
 	}
 	else
@@ -525,7 +520,7 @@ int kni_keepalive_replay_v6(struct stream_tuple4_v6* ipv6_addr,int iprever_flag,
 
 	struct kni_ipv6_hdr* snd_iphdr=NULL;
 	struct kni_tcp_hdr* snd_tcphdr=NULL;
-	char* sendbuf=(char*)malloc(iplen);
+	char* sendbuf= ALLOC(char, iplen);
 	memcpy(sendbuf,a_packet,iplen);
 
 	snd_iphdr=(struct kni_ipv6_hdr*)sendbuf;
@@ -588,7 +583,7 @@ int kni_keepalive_replay(struct stream_tuple4_v4* ipv4_addr,int iprever_flag,str
 
 	struct ip* snd_iphdr=NULL;
 	struct tcphdr* snd_tcphdr=NULL;
-	char* sendbuf=(char*)malloc(iplen);
+	char* sendbuf=ALLOC(char, iplen);
 	memcpy(sendbuf,a_packet,iplen);
 
 	snd_iphdr=(struct ip*)sendbuf;
@@ -1054,8 +1049,7 @@ void* pthread_process_tun(void* arg)
 //read from tun
 		clock_gettime(CLOCK_MONOTONIC, &start);
 		recv_len=0;
-		
-		recv_len=tun_read_data(g_kni_comminfo.fd_tun[thread_seq],recv_buf, KNI_MAX_BUFLEN);
+		recv_len=tun_read_data(g_kni_comminfo.fd_tun[thread_seq], recv_buf, sizeof(recv_buf));
 		if(recv_len <0)
 		{
 //			MESA_handle_runtime_log(g_kni_comminfo.logger, RLOG_LV_FATAL,KNI_MODULE_READTUN,"tun_read_data()error,recv_len:%d",recv_len);
@@ -1102,14 +1096,12 @@ int tcprepair_get_addr(void** client_addr,void** server_addr,const struct layer_
 		ipv4_hdr = (struct ip*)a_packet;
 		tcphdr=(struct kni_tcp_hdr*)((char*)ipv4_hdr+4*(ipv4_hdr->ip_hl));
 	
-		client_addr_v4 = (struct sockaddr_in*)malloc(sizeof(struct sockaddr_in));
-		memset(client_addr_v4,0,sizeof(struct sockaddr_in));
+		client_addr_v4 = ALLOC(struct sockaddr_in, 1);
 		client_addr_v4->sin_family = AF_INET;
 		client_addr_v4->sin_port = tcphdr->th_sport;
 		client_addr_v4->sin_addr.s_addr = (ipv4_hdr->ip_src).s_addr;
 		
-		server_addr_v4 = (struct sockaddr_in*)malloc(sizeof(struct sockaddr_in));
-		memset(server_addr_v4,0,sizeof(struct sockaddr_in));
+		server_addr_v4 = ALLOC(struct sockaddr_in, 1);
 		server_addr_v4->sin_family = AF_INET;
 		server_addr_v4->sin_port = tcphdr->th_dport;
 		server_addr_v4->sin_addr.s_addr = (ipv4_hdr->ip_dst).s_addr;
@@ -1122,17 +1114,14 @@ int tcprepair_get_addr(void** client_addr,void** server_addr,const struct layer_
 		ipv6_hdr = (struct kni_ipv6_hdr*)a_packet;
 		tcphdr = (struct kni_tcp_hdr*)((unsigned char*)a_packet + sizeof(struct kni_ipv6_hdr));
 	
-		client_addr_v6 = (struct sockaddr_in6*)malloc(sizeof(struct sockaddr_in6));
-		memset(client_addr_v6,0,sizeof(struct sockaddr_in6));
+		client_addr_v6 = ALLOC(struct sockaddr_in6, 1);
 		client_addr_v6->sin6_family = AF_INET6;
 		client_addr_v6->sin6_port = tcphdr->th_sport;
-		memcpy(&(client_addr_v6->sin6_addr),&(ipv6_hdr->ip6_src),sizeof(struct in6_addr));
 		
-		server_addr_v6 = (struct sockaddr_in6*)malloc(sizeof(struct sockaddr_in6));
-		memset(server_addr_v6,0,sizeof(struct sockaddr_in));
+		server_addr_v6 = ALLOC(struct sockaddr_in6, 1);
 		server_addr_v6->sin6_family = AF_INET6;
 		server_addr_v6->sin6_port = tcphdr->th_dport;
-		memcpy(&(server_addr_v6->sin6_addr),&(ipv6_hdr->ip6_dst),sizeof(struct in6_addr));
+		server_addr_v6->sin6_addr=ipv6_hdr->ip6_dst;
 
 		*client_addr = client_addr_v6;
 		*server_addr = server_addr_v6;