starnet/internal/pingcore/core.go

package pingcore

import (
	"encoding/binary"
	"net"
	"os"
	"sync/atomic"
	"time"
)

const icmpHeaderLen = 8

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

type Options struct {
	Count       int
	Timeout     time.Duration
	Interval    time.Duration
	Deadline    time.Time
	PreferIPv4  bool
	PreferIPv6  bool
	SourceIP    net.IP
	PayloadSize int
}

type Result struct {
	Duration  time.Duration
	RecvCount int
	RemoteIP  string
}

var identifierSeed uint32

func NextIdentifier() uint16 {
	pid := uint32(os.Getpid() & 0xffff)
	n := atomic.AddUint32(&identifierSeed, 1)
	return uint16((pid + n) & 0xffff)
}

func Payload(size int) []byte {
	if size <= 0 {
		return nil
	}
	payload := make([]byte, size)
	for index := 0; index < len(payload); index++ {
		payload[index] = byte(index)
	}
	return payload
}

func BuildICMP(seq, identifier uint16, typ uint8, payload []byte) ICMP {
	icmp := ICMP{
		Type:        typ,
		Code:        0,
		CheckSum:    0,
		Identifier:  identifier,
		SequenceNum: seq,
	}
	buf := MarshalPacket(icmp, payload)
	icmp.CheckSum = Checksum(buf)
	return icmp
}

func Checksum(data []byte) uint16 {
	var (
		sum    uint32
		length = 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]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return uint16(^sum)
}

func Marshal(icmp ICMP) []byte {
	return MarshalPacket(icmp, nil)
}

func MarshalPacket(icmp ICMP, payload []byte) []byte {
	buf := make([]byte, icmpHeaderLen+len(payload))
	buf[0] = icmp.Type
	buf[1] = icmp.Code
	binary.BigEndian.PutUint16(buf[2:], icmp.CheckSum)
	binary.BigEndian.PutUint16(buf[4:], icmp.Identifier)
	binary.BigEndian.PutUint16(buf[6:], icmp.SequenceNum)
	copy(buf[icmpHeaderLen:], payload)
	return buf
}

func IsExpectedEchoReply(packet []byte, family int, expectedType uint8, identifier, seq uint16) bool {
	for _, offset := range CandidateICMPOffsets(packet, family) {
		if offset < 0 || offset+icmpHeaderLen > len(packet) {
			continue
		}
		if packet[offset] != expectedType || packet[offset+1] != 0 {
			continue
		}
		if binary.BigEndian.Uint16(packet[offset+4:offset+6]) != identifier {
			continue
		}
		if binary.BigEndian.Uint16(packet[offset+6:offset+8]) != seq {
			continue
		}
		return true
	}
	return false
}

func CandidateICMPOffsets(packet []byte, family int) []int {
	offsets := []int{0}
	if len(packet) == 0 {
		return offsets
	}

	version := packet[0] >> 4
	if version == 4 && len(packet) >= 20 {
		ihl := int(packet[0]&0x0f) * 4
		if ihl >= 20 && ihl <= len(packet)-icmpHeaderLen {
			offsets = append(offsets, ihl)
		}
	} else if version == 6 && len(packet) >= 40+icmpHeaderLen {
		offsets = append(offsets, 40)
	}

	if family == 4 && len(packet) >= 20+icmpHeaderLen {
		offsets = append(offsets, 20)
	}
	if family == 6 && len(packet) >= 40+icmpHeaderLen {
		offsets = append(offsets, 40)
	}

	return DedupOffsets(offsets)
}

func DedupOffsets(offsets []int) []int {
	if len(offsets) <= 1 {
		return offsets
	}
	seen := make(map[int]struct{}, len(offsets))
	out := make([]int, 0, len(offsets))
	for _, offset := range offsets {
		if _, ok := seen[offset]; ok {
			continue
		}
		seen[offset] = struct{}{}
		out = append(out, offset)
	}
	return out
}

func ResolveTargets(host string, preferIPv4, preferIPv6 bool) ([]*net.IPAddr, error) {
	if parsed := net.ParseIP(host); parsed != nil {
		return []*net.IPAddr{{IP: parsed}}, nil
	}

	var targets []*net.IPAddr
	var err4 error
	var err6 error

	if ip4, err := net.ResolveIPAddr("ip4", host); err == nil && ip4 != nil && ip4.IP != nil {
		targets = append(targets, ip4)
	} else {
		err4 = err
	}

	if ip6, err := net.ResolveIPAddr("ip6", host); err == nil && ip6 != nil && ip6.IP != nil {
		targets = append(targets, ip6)
	} else {
		err6 = err
	}

	if len(targets) > 0 {
		return OrderTargets(targets, preferIPv4, preferIPv6), nil
	}
	if err4 != nil {
		return nil, err4
	}
	if err6 != nil {
		return nil, err6
	}
	return nil, nil
}

func OrderTargets(targets []*net.IPAddr, preferIPv4, preferIPv6 bool) []*net.IPAddr {
	if len(targets) <= 1 || preferIPv4 == preferIPv6 {
		return targets
	}

	ordered := make([]*net.IPAddr, 0, len(targets))
	if preferIPv4 {
		for _, target := range targets {
			if target != nil && target.IP != nil && target.IP.To4() != nil {
				ordered = append(ordered, target)
			}
		}
		for _, target := range targets {
			if target != nil && target.IP != nil && target.IP.To4() == nil {
				ordered = append(ordered, target)
			}
		}
		return ordered
	}

	for _, target := range targets {
		if target != nil && target.IP != nil && target.IP.To4() == nil {
			ordered = append(ordered, target)
		}
	}
	for _, target := range targets {
		if target != nil && target.IP != nil && target.IP.To4() != nil {
			ordered = append(ordered, target)
		}
	}
	return ordered
}