starnet/ping.go

package starnet

import (
	"bytes"
	"encoding/binary"
	"net"
	"time"
)

type ICMP struct {
	Type        uint8
	Code        uint8
	CheckSum    uint16
	Identifier  uint16
	SequenceNum uint16
}

func getICMP(seq uint16) ICMP {
	icmp := ICMP{
		Type:        8,
		Code:        0,
		CheckSum:    0,
		Identifier:  0,
		SequenceNum: seq,
	}
	var buffer bytes.Buffer
	binary.Write(&buffer, binary.BigEndian, icmp)
	icmp.CheckSum = checkSum(buffer.Bytes())
	buffer.Reset()

	return icmp
}

func sendICMPRequest(icmp ICMP, destAddr *net.IPAddr, timeout time.Duration) (PingResult, error) {
	var res PingResult
	conn, err := net.DialIP("ip:icmp", nil, destAddr)
	if err != nil {
		return res, err
	}
	defer conn.Close()
	var buffer bytes.Buffer
	binary.Write(&buffer, binary.BigEndian, icmp)

	if _, err := conn.Write(buffer.Bytes()); err != nil {
		return res, err
	}

	tStart := time.Now()

	conn.SetReadDeadline((time.Now().Add(timeout)))

	recv := make([]byte, 1024)
	res.RecvCount, err = conn.Read(recv)

	if err != nil {
		return res, err
	}

	tEnd := time.Now()
	res.Duration = tEnd.Sub(tStart)

	return res, err
}

func checkSum(data []byte) uint16 {
	var (
		sum    uint32
		length int = len(data)
		index  int
	)
	for length > 1 {
		sum += uint32(data[index])<<8 + uint32(data[index+1])
		index += 2
		length -= 2
	}
	if length > 0 {
		sum += uint32(data[index])
	}
	sum += (sum >> 16)

	return uint16(^sum)
}

type PingResult struct {
	Duration  time.Duration
	RecvCount int
}

func Ping(ip string, seq int, timeout time.Duration) (PingResult, error) {
	var res PingResult
	ipAddr, err := net.ResolveIPAddr("ip", ip)
	if err != nil {
		return res, err
	}
	icmp := getICMP(uint16(seq))
	return sendICMPRequest(icmp, ipAddr, timeout)
}