ping.go

package main

import (
	"crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"net"
	"os"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

const ttl = time.Second

type PingSession struct {
	ttl        time.Duration
	conn       *icmp.PacketConn
	proto      int
	targetAddr *net.IPAddr
	id         int
	seq        int
	total      float64
	lost       float64
	typ        icmp.Type
}

func main() {
	flag.Parse()

	target := flag.Arg(0)
	if target == "" {
		fmt.Println("usage: ping [host]")
		os.Exit(1)
	}

	targetAddr, err := net.ResolveIPAddr("ip", target)
	if err != nil {
		fmt.Println("failed to resolve target address:", err)
		os.Exit(1)
	}

	var conn *icmp.PacketConn
	var proto int
	var typ icmp.Type

	if targetAddr.IP.To4() == nil {
		// IPv6
		conn, err = icmp.ListenPacket("ip6:ipv6-icmp", "::")
		proto = 58
		typ = ipv6.ICMPTypeEchoRequest
	} else {
		// IPv4
		conn, err = icmp.ListenPacket("ip4:icmp", "0.0.0.0")
		proto = 1
		typ = ipv4.ICMPTypeEcho
	}
	if err != nil {
		fmt.Println("failed to open ICMP connection:", err)
		os.Exit(2)
	}

	var id uint16
	binary.Read(rand.Reader, binary.LittleEndian, &id)

	p := &PingSession{
		ttl:        time.Second,
		conn:       conn,
		proto:      proto,
		targetAddr: targetAddr,
		id:         int(id),
		seq:        0,
		total:      0,
		lost:       0,
		typ:        typ,
	}

	p.Run()
}

func (p *PingSession) Run() {
	readBuffer := make([]byte, 1500)
	t := time.NewTicker(ttl)
	for range t.C {
		data, err := (&icmp.Message{
			Type: p.typ, Code: 0,
			Body: &icmp.Echo{
				ID:   p.id,
				Seq:  p.seq,
				Data: []byte("hello world"),
			},
		}).Marshal(nil)
		if err != nil {
			panic("i suck at programming: " + err.Error())
		}

		_, err = p.conn.WriteTo(data, p.targetAddr)
		if err != nil {
			fmt.Println("warn: failed to send ICMP packet:", err)
		}

		start := time.Now()
		deadline := start.Add(ttl)

		for {
			p.conn.SetReadDeadline(deadline)
			n, from, err := p.conn.ReadFrom(readBuffer)
			if err != nil {
				if err, ok := err.(net.Error); ok && err.Timeout() {
					p.printResponse(false, "timeout", nil, time.Second)
					break
				}

				fmt.Println("failed to receive ICMP packets:", err)
				os.Exit(2)
			}

			latency := time.Since(start)

			pkt, ok := p.isMatchingPacket(readBuffer[:n])
			if !ok {
				continue
			}

			switch pkt.(type) {
			case *icmp.Echo:
				p.printResponse(true, "reply", from, latency)
			case *icmp.TimeExceeded:
				p.printResponse(false, "time limit exceeded", from, latency)
			case *icmp.DstUnreach:
				p.printResponse(false, "destination unreachable", from, latency)
			default:
				fmt.Println("lol no generics")
			}

			break
		}
	}
}

func (p *PingSession) printResponse(success bool, message string, from net.Addr, latency time.Duration) {
	p.total++
	if !success {
		p.lost++
	}

	if from == nil {
		fmt.Printf("%s - [ latency: %.1fms | loss: %.1f%% ]\n", message, float64(latency)/float64(time.Millisecond), 100.0*(p.lost/p.total))
	} else {
		fmt.Printf("%s from %s - [ latency: %.1fms | loss: %.1f%% ]\n", message, from.String(), float64(latency)/float64(time.Millisecond), 100.0*(p.lost/p.total))
	}
}

func (p *PingSession) isMatchingPacket(data []byte) (interface{}, bool) {
	msg, err := icmp.ParseMessage(p.proto, data)
	if err != nil {
		fmt.Println("warn: got malformed ICMP packet:", err)
		return nil, false
	}

	switch parsedMsg := msg.Body.(type) {
	case *icmp.Echo:
		if parsedMsg.ID == p.id && parsedMsg.Seq == p.seq {
			return parsedMsg, true
		}
	case *icmp.TimeExceeded:
		subMsg, err := icmp.ParseMessage(p.proto, parsedMsg.Data)
		if err != nil {
			break
		}

		parsedSubMsg, ok := subMsg.Body.(*icmp.Echo)
		if ok && parsedSubMsg.ID == p.id && parsedSubMsg.Seq == p.seq {
			return parsedMsg, true
		}
	case *icmp.DstUnreach:
		subMsg, err := icmp.ParseMessage(p.proto, parsedMsg.Data)
		if err != nil {
			break
		}

		parsedSubMsg, ok := subMsg.Body.(*icmp.Echo)
		if ok && parsedSubMsg.ID == p.id && parsedSubMsg.Seq == p.seq {
			return parsedMsg, true
		}
	}

	return nil, false
}