改为同步收集结果

Author: wardlin

icmp.go

@@ -3,20 +3,20 @@ package ztrace
 import (
 	"encoding/binary"
 	"net"
-	"sync/atomic"
 	"time"
 
 	"github.com/sirupsen/logrus"
 	"golang.org/x/net/ipv4"
 )
 
 func (t *TraceRoute) SendIPv4ICMP() error {
-	key := GetHash(t.netSrcAddr.To4(), t.netDstAddr.To4(), 65535, 65535, 1)
+	key := GetHash(t.NetSrcAddr.To4(), t.NetDstAddr.To4(), 65535, 65535, 1)
 	db := NewStatsDB(key)
 
 	t.DB.Store(key, db)
 	go db.Cache.Run()
-	conn, err := net.ListenPacket("ip4:icmp", t.netSrcAddr.String())
+
+	conn, err := net.ListenPacket("ip4:icmp", t.NetSrcAddr.String())
 	if err != nil {
 		logrus.Error(err)
 		return err
@@ -36,27 +36,23 @@ func (t *TraceRoute) SendIPv4ICMP() error {
 	for ttl := 1; ttl <= int(t.MaxTTL); ttl++ {
 		hdr, payload := t.BuildIPv4ICMP(uint8(ttl), id, id, 0)
 		rSocket.WriteTo(hdr, payload, nil)
-		report := &SendMetric{
+
+		m := &SendMetric{
 			FlowKey:   key,
 			ID:        uint32(hdr.ID),
 			TTL:       uint8(ttl),
 			TimeStamp: time.Now(),
 		}
-		t.SendChan <- report
 		id = (id + 1) % mod
 
-		//logrus.Info("send icmp ttl:", ttl)
-
-		//atomic.AddUint64(db.SendCnt, 1)
+		t.RecordSend(m)
 	}
 
 	return nil
 }
 
 func (t *TraceRoute) ListenIPv4ICMP() error {
-	defer t.Stop()
-
-	laddr := &net.IPAddr{IP: t.netSrcAddr}
+	laddr := &net.IPAddr{IP: t.NetSrcAddr}
 
 	var err error
 	t.recvICMPConn, err = net.ListenIP("ip4:icmp", laddr)
@@ -67,32 +63,24 @@ func (t *TraceRoute) ListenIPv4ICMP() error {
 	}
 	defer t.recvICMPConn.Close()
 
-	//设置超时时间
 	t.recvICMPConn.SetReadDeadline(time.Now().Add(t.Timeout))
 
 	for {
-		if atomic.LoadInt32(t.stopSignal) == 1 {
-			break
-		}
-
 		buf := make([]byte, 1500)
 		n, raddr, err := t.recvICMPConn.ReadFrom(buf)
 		if err != nil {
-			//logrus.Error("recvICMPConn.ReadFrom failed:", err)
 			break
 		}
 
 		icmpType := buf[0]
-		//logrus.Info(raddr, "|", icmpType, "|", n)
 
-		//TTL Exceeded or Port Unreachable
 		if (icmpType == 11 || (icmpType == 3 && buf[1] == 3)) && (n >= 36) {
 			id := binary.BigEndian.Uint16(buf[32:34])
 
 			dstip := net.IP(buf[24:28])
 			srcip := net.IP(buf[20:24])
 
-			if dstip.Equal(t.netDstAddr) {
+			if dstip.Equal(t.NetDstAddr) {
 				key := GetHash(srcip, dstip, 65535, 65535, 1)
 
 				m := &RecvMetric{
@@ -101,12 +89,15 @@ func (t *TraceRoute) ListenIPv4ICMP() error {
 					RespAddr:  raddr.String(),
 					TimeStamp: time.Now(),
 				}
-				//logrus.Info("recv icmp ttl:", id)
 
-				t.RecvChan <- m
+				if t.RecordRecv(m) {
+					break
+				}
 			}
 		}
 	}
 
+	t.Statistics()
+
 	return nil
 }

ipv6.go

@@ -19,8 +19,6 @@ func (t *TraceRoute) TraceIpv6ICMP() error {
 		return err
 	}
 
-	//TODO：获取本地ipv6地址
-
 	icmp6Sock, err := net.ListenPacket("ip6:ipv6-icmp", t.SrcAddr)
 	if err != nil {
 		logrus.Error("Could not set a listening ICMP6 socket: ", err)
@@ -79,12 +77,8 @@ func (t *TraceRoute) TraceIpv6ICMP() error {
 
 			n, _, node, err := ipv6Sock.ReadFrom(buf)
 
-			//收集一个响应包的数据
 			hop := map[string]interface{}{}
-
-			if err != nil {
-				//fmt.Printf("%d %40s\n", i, "*")
-			} else {
+			if err == nil {
 				answer, err := icmp.ParseMessage(58, buf[:n])
 
 				if err != nil {
@@ -95,7 +89,6 @@ func (t *TraceRoute) TraceIpv6ICMP() error {
 				timeCost := time.Since(timeNow)
 
 				if answer.Type == ipv6.ICMPTypeTimeExceeded {
-					//fmt.Printf("%d   %40s   %40s\n", i, node.String(), timeCost)
 
 					hop = map[string]interface{}{
 						"rtt":   fmt.Sprintf("%s", timeCost),
@@ -104,7 +97,6 @@ func (t *TraceRoute) TraceIpv6ICMP() error {
 
 					hopData.Details = append(hopData.Details, hop)
 				} else if answer.Type == ipv6.ICMPTypeEchoReply {
-					//fmt.Printf("%d   %40s   %40s\n", i, node.String(), timeCost)
 					hop = map[string]interface{}{
 						"rtt":   fmt.Sprintf("%s", timeCost),
 						"saddr": node.String(),
@@ -113,15 +105,11 @@ func (t *TraceRoute) TraceIpv6ICMP() error {
 
 					isDest = true
 					break
-				} else {
-					//fmt.Printf("%d %40s\n", i, "*")
 				}
-
 			}
-
 		}
 
-		t.hops = append(t.hops, hopData)
+		t.Hops = append(t.Hops, hopData)
 
 		if isDest {
 			break

packet.go

@@ -116,7 +116,6 @@ func (t *TCPHeader) checksum(ip *ipv4.Header, payload []byte) {
 	t.Checksum = checkSum(b.Bytes())
 }
 
-// BuildIPv4UDPkt 构造UDP探测包，可设置TTL
 func (t *TraceRoute) BuildIPv4UDPkt(srcPort uint16, dstPort uint16, ttl uint8, id uint16, tos int) (*ipv4.Header, []byte) {
 	iph := &ipv4.Header{
 		Version:  ipv4.Version,
@@ -129,8 +128,8 @@ func (t *TraceRoute) BuildIPv4UDPkt(srcPort uint16, dstPort uint16, ttl uint8, i
 		TTL:      int(ttl),
 		Protocol: 17,
 		Checksum: 0,
-		Src:      t.netSrcAddr,
-		Dst:      t.netDstAddr,
+		Src:      t.NetSrcAddr,
+		Dst:      t.NetDstAddr,
 	}
 
 	h, err := iph.Marshal()
@@ -158,7 +157,6 @@ func (t *TraceRoute) BuildIPv4UDPkt(srcPort uint16, dstPort uint16, ttl uint8, i
 	return iph, buf.Bytes()
 }
 
-// BuildIPv4TCPSYN 构造TCP探测包，可设置TTL
 func (t *TraceRoute) BuildIPv4TCPSYN(srcPort uint16, dstPort uint16, ttl uint8, seq uint32, tos int) (*ipv4.Header, []byte) {
 	iph := &ipv4.Header{
 		Version:  ipv4.Version,
@@ -171,8 +169,8 @@ func (t *TraceRoute) BuildIPv4TCPSYN(srcPort uint16, dstPort uint16, ttl uint8,
 		TTL:      int(ttl),
 		Protocol: 6,
 		Checksum: 0,
-		Src:      t.netSrcAddr,
-		Dst:      t.netDstAddr,
+		Src:      t.NetSrcAddr,
+		Dst:      t.NetDstAddr,
 	}
 
 	h, err := iph.Marshal()
@@ -193,7 +191,6 @@ func (t *TraceRoute) BuildIPv4TCPSYN(srcPort uint16, dstPort uint16, ttl uint8,
 		Urgent:     0,
 	}
 
-	//payload is TCP Optionheader
 	payload := []byte{0x02, 0x04, 0x05, 0xb4, 0x04, 0x02, 0x08, 0x0a, 0x7f, 0x73, 0xf9, 0x3a, 0x00, 0x00, 0x00, 0x00, 0x01, 0x03, 0x03, 0x07}
 	tcp.checksum(iph, payload)
 
@@ -215,8 +212,8 @@ func (t *TraceRoute) BuildIPv4TCPPRST(srcPort uint16, dstPort uint16, ttl uint8,
 		TTL:      int(ttl),
 		Protocol: 6,
 		Checksum: 0,
-		Src:      t.netSrcAddr,
-		Dst:      t.netDstAddr,
+		Src:      t.NetSrcAddr,
+		Dst:      t.NetDstAddr,
 	}
 
 	h, err := iph.Marshal()
@@ -264,8 +261,8 @@ func (t *TraceRoute) BuildIPv4ICMP(ttl uint8, id, seq uint16, tos int) (*ipv4.He
 		TTL:      int(ttl),
 		Protocol: 1,
 		Checksum: 0,
-		Src:      t.netSrcAddr,
-		Dst:      t.netDstAddr,
+		Src:      t.NetSrcAddr,
+		Dst:      t.NetDstAddr,
 	}
 
 	h, err := iph.Marshal()
@@ -276,7 +273,7 @@ func (t *TraceRoute) BuildIPv4ICMP(ttl uint8, id, seq uint16, tos int) (*ipv4.He
 	iph.Checksum = int(checkSum(h))
 
 	icmp := ICMPHeader{
-		IType:    8, //Echo
+		IType:    8,
 		ICode:    0,
 		Checksum: 0,
 		ID:       id,
@@ -301,16 +298,14 @@ func (t *TraceRoute) BuildIPv4ICMP(ttl uint8, id, seq uint16, tos int) (*ipv4.He
 
 func (t *TraceRoute) dnsResolve(hostName string, dst *net.IPAddr) (net.IP, error) {
 	isIPv6 := false
-	if t.af == "ip6"{
+	if t.Af == "ip6"{
 		isIPv6 = true
 	}
 
 	ipAddr := net.ParseIP(hostName)
 	if isIPv6 && ipAddr.To16() != nil {
-		//logrus.Info("Using the provided ipv6 address ", ipAddr," for tracing")
 		dst.IP = ipAddr
 	} else if !isIPv6 && ipAddr.To4() != nil {
-		//logrus.Info("Using the provided ipv4 address ", ipAddr, " for tracing")
 		dst.IP = ipAddr
 	} else {
 		ips, err := net.LookupIP(hostName)
@@ -323,11 +318,9 @@ func (t *TraceRoute) dnsResolve(hostName string, dst *net.IPAddr) (net.IP, error
 		for _, ip := range ips {
 			if isIPv6 && ip.To16() != nil {
 				dst.IP = ip
-				//logrus.Info(hostName, " resolved to ", ip, ", using this ipv6 address for tracing")
 				break
 			} else if !isIPv6 && ip.To4() != nil {
 				dst.IP = ip
-				//logrus.Info(hostName, " resolved to ", ip, ", using this ipv4 address for tracing")
 				break
 			}
 		}