调整edit_addr逻辑

1 files changed, 13 insertions, 23 deletions

diff --git a/src/tcpreplay.c b/src/tcpreplay.c
index e246c73..d457cd1 100644
--- a/src/tcpreplay.c
+++ b/src/tcpreplay.c
@@ -1209,7 +1209,7 @@ static inline int stream_addr_cmp_udp(struct ip *iphdr, struct udphdr *udphdr)
    为了不从新计算校验和，以达到较高的发送速率，
    改地址时，仅对IP地址进行修改，增加和减少的值要相等.
 */
-static int stream_edit_addr(u_char *pkt, int differ, struct ip* iphdr)
+static int stream_edit_addr(struct ip* raw_iphdr, int differ, struct ip* iphdr)
 {
 	//struct ip *iphdr;
 	struct tcphdr *tcphdr;
@@ -1219,10 +1219,10 @@ static int stream_edit_addr(u_char *pkt, int differ, struct ip* iphdr)
 
 	//iphdr = (struct ip *)(pkt + 14);
 
-	if(IPPROTO_TCP == iphdr->ip_p)
+	if(IPPROTO_TCP == raw_iphdr->ip_p || IPPROTO_UDP == raw_iphdr->ip_p)
 	{
-		tcphdr = (struct tcphdr *)(pkt + 14 + iphdr->ip_hl*4);
-		if(stream_addr_cmp_tcp(iphdr, tcphdr) < 0)
+		tcphdr = (struct tcphdr *)((char*)iphdr + iphdr->ip_hl*4);
+		if(stream_addr_cmp_tcp(raw_iphdr, tcphdr) < 0)
 		{		
 			iphdr->ip_src.s_addr -= differ;
 			iphdr->ip_dst.s_addr += differ;
@@ -1232,20 +1232,6 @@ static int stream_edit_addr(u_char *pkt, int differ, struct ip* iphdr)
 			iphdr->ip_src.s_addr += differ;
 			iphdr->ip_dst.s_addr -= differ;
 		}
-	}
-	else if(IPPROTO_UDP == iphdr->ip_p)
-	{
-		udphdr = (struct udphdr *)(pkt + 14 + iphdr->ip_hl*4);
-		if(stream_addr_cmp_tcp(iphdr, udphdr) < 0)
-		{
-			iphdr->ip_src.s_addr -= differ;
-			iphdr->ip_dst.s_addr += differ;
-		}
-		else
-		{
-			iphdr->ip_src.s_addr += differ;
-			iphdr->ip_dst.s_addr -= differ;
-		}
 	}else{ /* 增加对GRE, IPinIP, 6over4的支持 */
 		if(iphdr->ip_src.s_addr <= iphdr->ip_dst.s_addr){
 			iphdr->ip_src.s_addr += differ;
@@ -1531,6 +1517,7 @@ static const u_char *find_iphdr(const u_char *pkt, size_t pktlen, const u_char *
 	//struct ip *iphdr = (struct ip *)(pkt + 14);
 	struct ethhdr * eth = (struct ethhdr *)pkt;
 	struct ip *iphdr = NULL;
+	struct ip *t_iphdr = NULL;
 	
 	if(1 == flush) /* 原始包已读完, 强制刷新队列中剩余数据包 */
 	{
@@ -1583,13 +1570,17 @@ static const u_char *find_iphdr(const u_char *pkt, size_t pktlen, const u_char *
 		return -1;
 	}
 	
+	if(options.stream_multiple > 0)
+	{
+		memcpy(pbuf, pkt, pktlen);
+		t_iphdr = (struct ip *)(pbuf+((const u_char*)iphdr - pkt));	
+	}
 	if(options.pkt_distance > 0)
 	{
 		/* 暂存到队列 */
 		for(i = 0; i < options.stream_multiple; i++)
 		{
-			memcpy(pbuf, pkt, pktlen);
-			stream_edit_addr(pbuf, i+1, iphdr);
+			stream_edit_addr(iphdr, i+1, t_iphdr);
 			MESA_fixed_q_join_tail(&pkt_queue[i], pbuf, pktlen);
 		}
 
@@ -1610,9 +1601,8 @@ static const u_char *find_iphdr(const u_char *pkt, size_t pktlen, const u_char *
 		/* 如果没有包间距需求, 直接发送N倍数据包即可 */
 		for(i = 0; i < options.stream_multiple; i++)
 		{
-			memcpy(pbuf, pkt, pktlen);
-			stream_edit_addr(pbuf, i+1, iphdr); /* 在函数内部修改包头IP */
-			stream_burst_send_pkt(sp, pkt, pktlen);
+			stream_edit_addr(iphdr, i+1, t_iphdr); /* 在函数内部修改包头IP */
+			stream_burst_send_pkt(sp, pbuf, pktlen);
 			pkts_sent ++;
     			bytes_sent += pktlen;
 		}