test/ping/probing/icmp/packetconn.go

package icmp

import (
	"encoding/binary"
	"fmt"
	"net"
	"sync"
	"syscall"
	"time"

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

const (
	protocolICMP     = 1
	protocolIPv6ICMP = 58
)

var ENOLISTENER = fmt.Errorf("no listener")

type Type icmp.Type

type PacketConn interface {
	Close() error
	ICMPRequestType() Type
	ReadFrom(b []byte) (n int, ttl int, src net.Addr, err error)
	SetFlagTTL() error
	SetReadDeadline(t time.Time) error
	WriteTo(b []byte, dst net.Addr) (int, error)
	SetTTL(ttl int)
}

var (
	ipv4Proto = map[string]string{"icmp": "ip4:icmp", "udp": "udp4"}
	ipv6Proto = map[string]string{"icmp": "ip6:ipv6-icmp", "udp": "udp6"}
)

func Listen(ipv4 bool, protocol string, source string) (conn PacketConn, err error) {
	if ipv4 {
		var c icmpv4Conn
		c.c, err = icmp.ListenPacket(ipv4Proto[protocol], source)
		conn = &c
	} else {
		var c icmpV6Conn
		c.c, err = icmp.ListenPacket(ipv6Proto[protocol], source)
		conn = &c
	}
	return
}

// Packet represents a received and processed ICMP echo packet.
type Packet struct {
	// IPAddr is the address of the host being pinged.
	IPAddr *net.IPAddr
	// ID is the ICMP identifier.
	ID int
	// Seq is the ICMP sequence number.
	Seq int
	// TTL is the Time To Live on the packet.
	TTL int
	// NBytes is the number of bytes in the message.
	Nbytes int
	// Rtt is the round-trip time it took to ping.
	Rtt time.Duration
}

type recvPkt struct {
	recvtime time.Time
	addr     net.Addr
	bytes    []byte
	nbytes   int
	ttl      int
}

type MPacketConn struct {
	IPV4         bool
	Protocol     string
	Source       string
	Backlog      int
	OnRecvPacket func(addr net.Addr, pkt *Packet)
	OnRecvError  func(error)
	mutex        sync.Mutex
	conn         PacketConn
	done         chan interface{}
	recvbuf      chan *recvPkt
}

func (mp *MPacketConn) Listen() error {
	conn, err := Listen(mp.IPV4, mp.Protocol, mp.Source)
	if err != nil {
		return err
	}

	conn.SetTTL(64)
	if err := conn.SetFlagTTL(); err != nil {
		return err
	}

	mp.done = make(chan interface{})
	mp.recvbuf = make(chan *recvPkt, mp.Backlog)
	mp.conn = conn
	go mp.recvICMP()
	go mp.processRecvPacket()

	return nil
}

func (mp *MPacketConn) Close() error {
	mp.mutex.Lock()
	defer mp.mutex.Unlock()

	open := true
	select {
	case _, open = <-mp.done:
	default:
	}

	if open {
		close(mp.done)
	}
	return mp.conn.Close()
}

func (mp *MPacketConn) recvICMP() {
	bytes := make([]byte, 65536)
	for {
		select {
		case <-mp.done:
			return
		default:
			var n, ttl int
			var addr net.Addr
			var err error
			n, ttl, addr, err = mp.conn.ReadFrom(bytes)
			if err != nil {
				if neterr, ok := err.(*net.OpError); ok {
					if neterr.Timeout() {
						// Read timeout
						continue
					}
				}
				if mp.OnRecvError != nil {
					mp.OnRecvError(err)
				} else {
					fmt.Println(err)
				}
			}
			bs := make([]byte, n)
			copy(bs, bytes[:n])
			select {
			case <-mp.done:
				return
			case mp.recvbuf <- &recvPkt{recvtime: time.Now(), addr: addr, bytes: bs, nbytes: n, ttl: ttl}:
			}
		}
	}
}

func (mp *MPacketConn) SendPacket(pkt *Packet, addr *net.IPAddr) error {
	if mp.conn == nil {
		return ENOLISTENER
	}
	msgBytes, err := BuildEchoRequestMessage(pkt.ID, pkt.Seq, pkt.Nbytes, mp.conn.ICMPRequestType())
	if err != nil {
		return err
	}
	var dst net.Addr = addr
	if mp.Protocol == "udp" {
		dst = &net.UDPAddr{IP: addr.IP, Zone: addr.Zone}
	}
	for {
		select {
		case <-mp.done:
			return nil
		default:
		}
		if _, err := mp.conn.WriteTo(msgBytes, dst); err != nil {
			if neterr, ok := err.(*net.OpError); ok {
				if neterr.Err == syscall.ENOBUFS {
					if mp.OnRecvError != nil {
						mp.OnRecvError(neterr.Err)
					} else {
						fmt.Println("缓存不够，发送失败，重发")
					}
					continue
				}
			}
			return err
		} else {
			return nil
		}
	}
}

func (mp *MPacketConn) processRecvPacket() {
	for pkt := range mp.recvbuf {
		err := mp.processPacket(pkt)
		if err != nil {
			if mp.OnRecvError != nil {
				mp.OnRecvError(err)
			} else {
				fmt.Println(err)
			}
		}
	}
}

var count = 0

func (mp *MPacketConn) processPacket(recv *recvPkt) error {
	var proto int
	if mp.IPV4 {
		proto = protocolICMP
	} else {
		proto = protocolIPv6ICMP
	}

	// fmt.Println(count, "from", recv.addr.String(), "bytes", recv.bytes)

	var m *icmp.Message
	var err error
	if m, err = icmp.ParseMessage(proto, recv.bytes); err != nil {
		return fmt.Errorf("error parsing icmp message: %w", err)
	}

	if m.Type != ipv4.ICMPTypeEchoReply && m.Type != ipv6.ICMPTypeEchoReply {
		// Not an echo reply, ignore it
		return nil
	}

	switch pkt := m.Body.(type) {
	case *icmp.Echo:
		return mp.processEchoReply(pkt, recv)
	default:
		// Very bad, not sure how this can happen
		return fmt.Errorf("invalid ICMP echo reply; type: '%T', '%v'", pkt, pkt)
	}
}

func (mp *MPacketConn) processEchoReply(pkt *icmp.Echo, recv *recvPkt) error {
	if len(pkt.Data) < 24 {
		return nil
	}

	sendtime := int64(binary.BigEndian.Uint64(pkt.Data[:8]))
	fullseq := int(binary.BigEndian.Uint64(pkt.Data[8:16]))
	fullid := int(binary.BigEndian.Uint64(pkt.Data[16:24]))

	if fullid%65536 != pkt.ID || fullseq%65536 != pkt.Seq {
		return nil
	}

	inPkt := &Packet{
		Addr:   recv.addr.String(),
		ID:     fullid,
		Seq:    fullseq,
		Nbytes: recv.nbytes,
		TTL:    recv.ttl,
		Rtt:    recv.recvtime.Sub(time.Unix(0, sendtime)),
	}

	// fmt.Printf("%s %d bytes from %s: icmp_seq=%d time=%v\n",
	// 	time.Now().Format("15:04:05.000"), inPkt.Nbytes, inPkt.IPAddr, inPkt.Seq, inPkt.Rtt)

	p.mutex.Lock()
	chpkt, inflight := pinfo.seqpkt[fullseq]
	if inflight {
		// remove it from the list of sequences we're waiting for so we don't get duplicates.
		delete(pinfo.seqpkt, fullseq)
	}
	p.mutex.Unlock()

	if chpkt != nil {
		chpkt <- inPkt
	}
	return nil
}