icmp: Unifiy ICMP send, warp data into binary

Fixes #5

icmp/icmp.go

@@ -2,7 +2,7 @@ package icmp
 
 import (
 	"bytes"
-	"fmt"
+	"encoding/binary"
 	"log"
 	"net"
 	"time"
@@ -20,51 +20,7 @@ type ICMPReturn struct {
 
 // SendDiscoverICMP sends a ICMP to a given destination with a TTL to discover hops
 func SendDiscoverICMP(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 {
-		log.Panicf("Failed to listen to address %v. Msg: %v.", localAddr, err.Error())
-		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
-	}
-
-	wm := icmp.Message{
-		Type: ipv4.ICMPTypeEcho, Code: 0,
-		Body: &icmp.Echo{
-			ID: pid, Seq: seq,
-			Data: []byte(""),
-		},
-	}
-	wb, err := wm.Marshal(nil)
-	if err != nil {
-		return hop, err
-	}
-
-	if _, err := c.WriteTo(wb, dst); err != nil {
-		return hop, err
-	}
-
-	// rb := make([]byte, 1500)
-	rb := []byte{}
-	_, peer, err := c.ReadFrom(rb)
-	if err != nil {
-		return hop, err
-	}
-	elapsed := time.Since(start)
-	hop.Elapsed = elapsed
-	hop.Addr = peer.String()
-	hop.Success = true
-	return hop, err
+	return SendICMP(localAddr, dst, "", ttl, pid, timeout, seq)
 }
 
 // SendICMP sends a ICMP to a given destination which requires a  reply from that specific destination
@@ -86,12 +42,13 @@ func SendICMP(localAddr string, dst net.Addr, target string, ttl, pid int, timeo
 		return hop, err
 	}
 
-	body := fmt.Sprintf("ping%d", seq)
+	bs := make([]byte, 4)
+	binary.LittleEndian.PutUint32(bs, uint32(seq))
 	wm := icmp.Message{
 		Type: ipv4.ICMPTypeEcho, Code: 0,
 		Body: &icmp.Echo{
 			ID: pid, Seq: seq,
-			Data: []byte(body),
+			Data: append(bs, 'x'),
 		},
 	}
 	wb, err := wm.Marshal(nil)
@@ -103,7 +60,7 @@ func SendICMP(localAddr string, dst net.Addr, target string, ttl, pid int, timeo
 		return hop, err
 	}
 
-	peer, _, err := listenForSpecific(c, time.Now().Add(timeout), target, body, seq, wb)
+	peer, _, err := listenForSpecific(c, time.Now().Add(timeout), target, append(bs, 'x'), seq, wb)
 	if err != nil {
 		return hop, err
 	}
@@ -116,20 +73,20 @@ func SendICMP(localAddr string, dst net.Addr, target string, ttl, pid int, timeo
 }
 
 // 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) {
+func listenForSpecific(conn *icmp.PacketConn, deadline time.Time, neededPeer string, neededBody []byte, needSeq int, sent []byte) (string, []byte, error) {
 	for {
 		b := make([]byte, 1500)
 		n, peer, err := conn.ReadFrom(b)
 		if err != nil {
 			if neterr, ok := err.(*net.OpError); ok {
-				return "", "", neterr
+				return "", []byte{}, neterr
 			}
 		}
 		if n == 0 {
 			continue
 		}
 
-		if peer.String() != neededPeer {
+		if neededPeer != "" && peer.String() != neededPeer {
 			continue
 		}
 
@@ -148,7 +105,7 @@ func listenForSpecific(conn *icmp.PacketConn, deadline time.Time, neededPeer, ne
 				case *icmp.Echo:
 					seq := x.Body.(*icmp.Echo).Seq
 					if seq == needSeq {
-						return peer.String(), "", nil
+						return peer.String(), []byte{}, nil
 					}
 				default:
 					// ignore
@@ -158,10 +115,10 @@ func listenForSpecific(conn *icmp.PacketConn, deadline time.Time, neededPeer, ne
 
 		if x.Type.(ipv4.ICMPType).String() == "echo reply" {
 			b, _ := x.Body.Marshal(1)
-			if string(b[4:]) != neededBody {
+			if string(b[4:]) != string(neededBody) {
 				continue
 			}
-			return peer.String(), string(b[4:]), nil
+			return peer.String(), b[4:], nil
 		}
 	}
 }