practice

ping.go

package main

import (
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"strconv"
	"time"
)

type PingPacket struct {
	Identifier     uint16
	SequenceNumber uint16
	Data           []byte
}

func checkSum(b []byte) uint16 {
	if len(b)%2 == 1 {
		b = append(b, 0)
	}
	// FIXME: Is it correct to use int?
	sum := 0
	for i := 0; i < len(b); i += 2 {
		sum += 0xffff - ((int(b[i]) << 8) | int(b[i+1]))
		sum = sum&0xffff + ((sum & 0x10000) >> 16)
	}
	return uint16(sum)
}

// See the comment in the middle of this function for information of
// `enableTimestamp`.
func (p *PingPacket) Encode(enableTimestamp bool) []byte {
	var headerSize int
	if enableTimestamp {
		headerSize = 16
	} else {
		headerSize = 8
	}
	payload := make([]byte, headerSize+len(p.Data))
	payload[0] = 8 // 8 means Echo Message
	payload[4] = uint8(p.Identifier >> 8)
	payload[5] = uint8(p.Identifier & 0xff)
	payload[6] = uint8(p.SequenceNumber >> 8)
	payload[7] = uint8(p.SequenceNumber & 0xff)

	// When I see a ping packet via Wireshark, there exits 'Timestamp from icmp
	// data' field in the tail of ICMP header, although it is not an timestamp
	// message.  This is an easy implementation of embed the current time as
	// 8 bytes data.  FIXME: It seems that this code contains some bugs.
	if enableTimestamp {
		now := time.Now().Unix()
		payload[8] = byte((now >> 56) & 0xff)
		payload[9] = byte((now >> 48) & 0xff)
		payload[10] = byte((now >> 40) & 0xff)
		payload[11] = byte((now >> 32) & 0xff)
		payload[12] = byte((now >> 24) & 0xff)
		payload[13] = byte((now >> 16) & 0xff)
		payload[14] = byte((now >> 8) & 0xff)
		payload[15] = byte(now & 0xff)
	}

	copy(payload[headerSize:], p.Data)
	c := checkSum(payload)
	copy(payload[2:4], []byte{uint8(c >> 8), uint8(c & 0xff)})

	return payload
}

func Validate(b []byte) bool {
	if ^checkSum(b) != 0 {
		return false
	}
	if len(b) < 8 {
		return false
	}
	return true
}

func Decode(b []byte) (*PingPacket, error) {
	if !Validate(b) {
		return nil, errors.New("A malformed ping packet.")
	}

	return &PingPacket{
		Identifier:     uint16(b[4])<<8 | uint16(b[5]),
		SequenceNumber: uint16(b[6])<<8 | uint16(b[7]),
		Data:           b[8:],
	}, nil
}

func (p *PingPacket) Send(laddr, raddr *net.IPAddr) (int, error) {
	conn, err := net.DialIP("ip4:icmp", laddr, raddr)
	if err != nil {
		return 0, err
	}

	n, err := conn.Write(p.Encode(false))
	if err != nil {
		return 0, err
	}

	return n, nil
}

func (p *PingPacket) SendAndReceiveResponse(laddr, raddr *net.IPAddr, timeout time.Duration) (time.Duration, error) {
	sendTime := time.Now()

	sendBytes, err := p.Send(laddr, raddr)
	if err != nil {
		return 0, err
	}

	var done = make(chan struct{})
	var fail = make(chan error)

	// Receive the response.
	go func(bufSize int) {
		// FIXME: This implementation that listens and closes the connection
		// each time the ping packets send may degrades performance.  We should
		// listen the connection just once in one PingRepeatedly invocation.
		ln, err := net.ListenIP("ip4", laddr)
		if err != nil {
			fail <- err
			return
		}
		defer ln.Close()

		buf := make([]byte, bufSize)

		// Read packets until the corresponding pong found.
		for {
			// FIXME: WHY Read returns irrelevant bytes???
			readBytes, err := ln.Read(buf)
			if err != nil {
				fail <- err
				return
			}
			packet, err := Decode(buf[readBytes-sendBytes : readBytes])
			if err != nil {
				fail <- err
				return
			}
			if packet.Identifier == p.Identifier && packet.SequenceNumber == p.SequenceNumber {
				done <- struct{}{}
				break
			} else {
				continue
			}
		}
	}(sendBytes + 100) // FIXME: I don't know the exact size of the incoming packet.

	select {
	case <-done:
		return time.Since(sendTime), err
	case err := <-fail:
		return 0, err
	case <-time.After(timeout):
		return 0, errors.New("Timeout")
	}
}

func PingRepeatedly(times int, laddr, raddr *net.IPAddr, timeout time.Duration) {
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	packet := &PingPacket{
		Identifier:     uint16(rand.Int() & 0xffff),
		SequenceNumber: 0,
		Data:           []byte("DROP YOU VIVID COLORS"),
	}

	total, success := 0, 0

	done := make(chan struct{})
	go func() {
		for i := 0; times < 0 || i < times; i++ {
			t, err := packet.SendAndReceiveResponse(laddr, raddr, timeout)
			if err != nil {
				fmt.Printf("%2d: error: %s\n", i+1, err)
			} else {
				fmt.Printf("%2d: time = %s\n", i+1, t)
				success++
			}
			total++
			packet.SequenceNumber++
			<-time.After(500 * time.Millisecond)
		}
		done <- struct{}{}
	}()

	select {
	case <-done:
		fmt.Printf("received / send = %d / %d\n", success, total)
	case <-interrupt:
		fmt.Printf("\nreceived / send = %d / %d\n", success, total)
		os.Exit(0)
	}
}

func Help() {
	fmt.Fprintf(os.Stderr, "usage: %s host [times]\n", os.Args[0])
	os.Exit(1)
}

func main() {
	if len(os.Args) < 2 {
		Help()
	}
	host := os.Args[1]
	times := -1
	if len(os.Args) >= 3 {
		var err error
		times, err = strconv.Atoi(os.Args[2])
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: [%s] is not an integer.", os.Args[2])
			Help()
		}
	}

	laddr, err := net.ResolveIPAddr("ip4", "0.0.0.0")
	if err != nil {
		fmt.Fprintln(os.Stderr, "ResolveIPAddr failed:", err)
		os.Exit(1)
	}
	raddr, err := net.ResolveIPAddr("ip4", host)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ResolveIPAddr for %s failed: %s\n", host, err)
		os.Exit(1)
	}

	PingRepeatedly(times, laddr, raddr, time.Second)
}