Merge pull request #4 from meyskens/concurrency

Make ping after discovery concurrent

hop/hop.go

@@ -25,6 +25,7 @@ type HopStatistic struct {
 	Lost           int
 	Packets        *ring.Ring
 	RingBufferSize int
+	pingSeq        int
 }
 
 type packet struct {
@@ -33,18 +34,24 @@ type packet struct {
 }
 
 func (s *HopStatistic) Next() {
-	r, _ := imcp.SendIMCP("0.0.0.0", s.Dest, s.TTL, s.PID, s.Timeout)
-	s.Packets = s.Packets.Prev()
-	s.Packets.Value = r
 	if s.Target == "" {
-		s.Target = r.Addr
+		return
 	}
+	s.pingSeq++
+	r, _ := imcp.SendIMCP("0.0.0.0", s.Dest, s.Target, s.TTL, s.PID, s.Timeout, s.pingSeq)
+	s.Packets = s.Packets.Prev()
+	s.Packets.Value = r
+
 	s.Sent++
-	s.SumElapsed = r.Elapsed + s.SumElapsed
+
+	s.Last = r
 	if !r.Success {
 		s.Lost++
+		return // do not count failed into statistics
 	}
-	s.Last = r
+
+	s.SumElapsed = r.Elapsed + s.SumElapsed
+
 	if s.Best.Elapsed > r.Elapsed {
 		s.Best = r
 	}

imcp/icmp.go

@@ -1,6 +1,8 @@
 package imcp
 
 import (
+	"bytes"
+	"fmt"
 	"net"
 	"time"
 
@@ -15,21 +17,29 @@ type ICMPReturn struct {
 	Elapsed time.Duration
 }
 
-// SendIMCP sends a IMCP to a given destination
-func SendIMCP(localAddr string, dst net.Addr, ttl, pid int, timeout time.Duration) (hop ICMPReturn, err error) {
+// SendDiscoverIMCP sends a IMCP to a given destination with a TTL to discover hops
+func SendDiscoverIMCP(localAddr string, dst net.Addr, ttl, pid int, timeout time.Duration, seq int) (hop ICMPReturn, err error) {
 	hop.Success = false
 	start := time.Now()
 	c, err := icmp.ListenPacket("ip4:icmp", localAddr)
 	if err != nil {
 		return hop, err
 	}
 	defer c.Close()
-	c.IPv4PacketConn().SetTTL(ttl)
-	c.SetDeadline(time.Now().Add(timeout))
+
+	err = c.IPv4PacketConn().SetTTL(ttl)
+	if err != nil {
+		return hop, err
+	}
+	err = c.SetDeadline(time.Now().Add(timeout))
+	if err != nil {
+		return hop, err
+	}
+
 	wm := icmp.Message{
 		Type: ipv4.ICMPTypeEcho, Code: 0,
 		Body: &icmp.Echo{
-			ID: pid, Seq: 1,
+			ID: pid, Seq: seq,
 			Data: []byte(""),
 		},
 	}
@@ -53,3 +63,101 @@ func SendIMCP(localAddr string, dst net.Addr, ttl, pid int, timeout time.Duratio
 	hop.Success = true
 	return hop, err
 }
+
+// SendIMCP sends a IMCP to a given destination which requires a  reply from that specific destination
+func SendIMCP(localAddr string, dst net.Addr, target string, ttl, pid int, timeout time.Duration, seq int) (hop ICMPReturn, err error) {
+	hop.Success = false
+	start := time.Now()
+	c, err := icmp.ListenPacket("ip4:icmp", localAddr)
+	if err != nil {
+		return hop, err
+	}
+	defer c.Close()
+	err = c.IPv4PacketConn().SetTTL(ttl)
+	if err != nil {
+		return hop, err
+	}
+	err = c.SetDeadline(time.Now().Add(timeout))
+	if err != nil {
+		return hop, err
+	}
+
+	body := fmt.Sprintf("ping%d", seq)
+	wm := icmp.Message{
+		Type: ipv4.ICMPTypeEcho, Code: 0,
+		Body: &icmp.Echo{
+			ID: pid, Seq: seq,
+			Data: []byte(body),
+		},
+	}
+	wb, err := wm.Marshal(nil)
+	if err != nil {
+		return hop, err
+	}
+
+	if _, err := c.WriteTo(wb, dst); err != nil {
+		return hop, err
+	}
+
+	peer, _, err := listenForSpecific(c, time.Now().Add(timeout), target, body, seq, wb)
+	if err != nil {
+		return hop, err
+	}
+
+	elapsed := time.Since(start)
+	hop.Elapsed = elapsed
+	hop.Addr = peer
+	hop.Success = true
+	return hop, err
+}
+
+// listenForSpecific listens for a reply from a specific destination with a specifi body and returns the body if returned
+func listenForSpecific(conn *icmp.PacketConn, deadline time.Time, neededPeer, neededBody string, needSeq int, sent []byte) (string, string, error) {
+	for {
+		b := make([]byte, 1500)
+		n, peer, err := conn.ReadFrom(b)
+		if err != nil {
+			if neterr, ok := err.(*net.OpError); ok {
+				return "", "", neterr
+			}
+		}
+		if n == 0 {
+			continue
+		}
+
+		if peer.String() != neededPeer {
+			continue
+		}
+
+		x, err := icmp.ParseMessage(1, b[:n])
+		if err != nil {
+			continue
+		}
+
+		if x.Type.(ipv4.ICMPType).String() == "time exceeded" {
+			body := x.Body.(*icmp.TimeExceeded).Data
+
+			index := bytes.Index(body, sent[:4])
+			if index > 0 {
+				x, _ := icmp.ParseMessage(1, body[index:])
+				switch x.Body.(type) {
+				case *icmp.Echo:
+					seq := x.Body.(*icmp.Echo).Seq
+					if seq == needSeq {
+						return peer.String(), "", nil
+					}
+				default:
+					// ignore
+				}
+			}
+		}
+
+		if x.Type.(ipv4.ICMPType).String() == "echo reply" {
+			b, _ := x.Body.Marshal(1)
+			if string(b[4:]) != neededBody {
+				continue
+			}
+			return peer.String(), string(b[4:]), nil
+		}
+	}
+}

mtr/mtr.go

@@ -97,7 +97,7 @@ func (m *MTR) discover(ch chan struct{}) {
 	unknownHopsCount := 0
 	for ttl := 1; ttl < m.maxHops; ttl++ {
 		time.Sleep(m.hopsleep)
-		hopReturn, err := imcp.SendIMCP("0.0.0.0", &ipAddr, ttl, pid, m.timeout)
+		hopReturn, err := imcp.SendDiscoverIMCP("0.0.0.0", &ipAddr, ttl, pid, m.timeout, 1)
 
 		m.mutex.Lock()
 		s := m.registerStatistic(ttl, hopReturn)