Change to use golang.org/x/net/icmp instead of my own icmp.go

fastping.go

@@ -45,6 +45,11 @@ import (
 	"sync"
 	"syscall"
 	"time"
+
+	"golang.org/x/net/icmp"
+	"golang.org/x/net/internal/iana"
+	"golang.org/x/net/ipv4"
+	"golang.org/x/net/ipv6"
 )
 
 const TimeSliceLength = 8
@@ -83,6 +88,14 @@ func isIPv6(ip net.IP) bool {
 	return len(ip) == net.IPv6len
 }
 
+func ipv4Payload(b []byte) []byte {
+	if len(b) < ipv4.HeaderLen {
+		return b
+	}
+	hdrlen := int(b[0]&0x0f) << 2
+	return b[hdrlen:]
+}
+
 type packet struct {
 	bytes []byte
 	addr  *net.IPAddr
@@ -389,13 +402,13 @@ func (p *Pinger) sendICMP(conn, conn6 *net.IPConn) (map[string]*net.IPAddr, erro
 	queue := make(map[string]*net.IPAddr)
 	wg := new(sync.WaitGroup)
 	for key, addr := range p.addrs {
-		var typ int
+		var typ icmp.Type
 		var cn *net.IPConn
 		if isIPv4(addr.IP) {
-			typ = icmpv4EchoRequest
+			typ = ipv4.ICMPTypeEcho
 			cn = conn
 		} else if isIPv6(addr.IP) {
-			typ = icmpv6EchoRequest
+			typ = ipv6.ICMPTypeEchoRequest
 			cn = conn6
 		} else {
 			continue
@@ -411,13 +424,13 @@ func (p *Pinger) sendICMP(conn, conn6 *net.IPConn) (map[string]*net.IPAddr, erro
 		}
 
 		p.mu.Lock()
-		bytes, err := (&icmpMessage{
+		bytes, err := (&icmp.Message{
 			Type: typ, Code: 0,
-			Body: &icmpEcho{
+			Body: &icmp.Echo{
 				ID: p.id, Seq: p.seq,
 				Data: t,
 			},
-		}).Marshal()
+		}).Marshal(nil)
 		p.mu.Unlock()
 		if err != nil {
 			wg.Wait()
@@ -506,27 +519,30 @@ func (p *Pinger) procRecv(recv *packet, queue map[string]*net.IPAddr) {
 	p.mu.Unlock()
 
 	var bytes []byte
+	var proto int
 	if isIPv4(recv.addr.IP) {
 		bytes = ipv4Payload(recv.bytes)
+		proto = iana.ProtocolICMP
 	} else if isIPv6(recv.addr.IP) {
 		bytes = recv.bytes
+		proto = iana.ProtocolIPv6ICMP
 	} else {
 		return
 	}
 
-	var m *icmpMessage
+	var m *icmp.Message
 	var err error
-	if m, err = parseICMPMessage(bytes); err != nil {
+	if m, err = icmp.ParseMessage(proto, bytes); err != nil {
 		return
 	}
 
-	if m.Type != icmpv4EchoReply && m.Type != icmpv6EchoReply {
+	if m.Type != ipv4.ICMPTypeEchoReply && m.Type != ipv6.ICMPTypeEchoReply {
 		return
 	}
 
 	var rtt time.Duration
 	switch pkt := m.Body.(type) {
-	case *icmpEcho:
+	case *icmp.Echo:
 		p.mu.Lock()
 		if pkt.ID == p.id && pkt.Seq == p.seq {
 			rtt = time.Since(bytesToTime(pkt.Data[:TimeSliceLength]))

icmp.go

@@ -1,158 +0,0 @@
-// Almost all this code is taken from go net ipraw_test.go implementation.
-// It's under a BSD-style license.
-
-// Copyright (c) 2012 The Go Authors. All rights reserved.
-//
-// Redistribution and use in source and binary forms, with or without
-// modification, are permitted provided that the following conditions are
-// met:
-//
-//    * Redistributions of source code must retain the above copyright
-// notice, this list of conditions and the following disclaimer.
-//    * Redistributions in binary form must reproduce the above
-// copyright notice, this list of conditions and the following disclaimer
-// in the documentation and/or other materials provided with the
-// distribution.
-//    * Neither the name of Google Inc. nor the names of its
-// contributors may be used to endorse or promote products derived from
-// this software without specific prior written permission.
-//
-// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
-// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
-// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
-// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
-// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
-// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
-// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-// Copyright 2009 The Go Authors.  All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-package fastping
-
-import (
-	"errors"
-)
-
-func ipv4Payload(b []byte) []byte {
-	if len(b) < 20 {
-		return b
-	}
-	hdrlen := int(b[0]&0x0f) << 2
-	return b[hdrlen:]
-}
-
-const (
-	icmpv4EchoRequest = 8
-	icmpv4EchoReply   = 0
-	icmpv6EchoRequest = 128
-	icmpv6EchoReply   = 129
-)
-
-// icmpMessage represents an ICMP message.
-type icmpMessage struct {
-	Type     int             // type
-	Code     int             // code
-	Checksum int             // checksum
-	Body     icmpMessageBody // body
-}
-
-// icmpMessageBody represents an ICMP message body.
-type icmpMessageBody interface {
-	Len() int
-	Marshal() ([]byte, error)
-}
-
-// Marshal returns the binary enconding of the ICMP echo request or
-// reply message m.
-func (m *icmpMessage) Marshal() ([]byte, error) {
-	b := []byte{byte(m.Type), byte(m.Code), 0, 0}
-	if m.Body != nil && m.Body.Len() != 0 {
-		mb, err := m.Body.Marshal()
-		if err != nil {
-			return nil, err
-		}
-		b = append(b, mb...)
-	}
-	switch m.Type {
-	case icmpv6EchoRequest, icmpv6EchoReply:
-		return b, nil
-	}
-	csumcv := len(b) - 1 // checksum coverage
-	s := uint32(0)
-	for i := 0; i < csumcv; i += 2 {
-		s += uint32(b[i+1])<<8 | uint32(b[i])
-	}
-	if csumcv&1 == 0 {
-		s += uint32(b[csumcv])
-	}
-	s = s>>16 + s&0xffff
-	s = s + s>>16
-	// Place checksum back in header; using ^= avoids the
-	// assumption the checksum bytes are zero.
-	b[2] ^= byte(^s)
-	b[3] ^= byte(^s >> 8)
-	return b, nil
-}
-
-// parseICMPMessage parses b as an ICMP message.
-func parseICMPMessage(b []byte) (*icmpMessage, error) {
-	msglen := len(b)
-	if msglen < 4 {
-		return nil, errors.New("message too short")
-	}
-	m := &icmpMessage{Type: int(b[0]), Code: int(b[1]), Checksum: int(b[2])<<8 | int(b[3])}
-	if msglen > 4 {
-		var err error
-		switch m.Type {
-		case icmpv4EchoRequest, icmpv4EchoReply, icmpv6EchoRequest, icmpv6EchoReply:
-			m.Body, err = parseICMPEcho(b[4:])
-			if err != nil {
-				return nil, err
-			}
-		}
-	}
-	return m, nil
-}
-
-// imcpEcho represenets an ICMP echo request or reply message body.
-type icmpEcho struct {
-	ID   int    // identifier
-	Seq  int    // sequence number
-	Data []byte // data
-}
-
-// Len returns message length in bytes
-func (p *icmpEcho) Len() int {
-	if p == nil {
-		return 0
-	}
-	return 4 + len(p.Data)
-}
-
-// Marshal returns the binary enconding of the ICMP echo request or
-// reply message body p.
-func (p *icmpEcho) Marshal() ([]byte, error) {
-	b := make([]byte, 4+len(p.Data))
-	b[0], b[1] = byte(p.ID>>8), byte(p.ID)
-	b[2], b[3] = byte(p.Seq>>8), byte(p.Seq)
-	copy(b[4:], p.Data)
-	return b, nil
-}
-
-// parseICMPEcho parses b as an ICMP echo request or reply message
-// body.
-func parseICMPEcho(b []byte) (*icmpEcho, error) {
-	bodylen := len(b)
-	p := &icmpEcho{ID: int(b[0])<<8 | int(b[1]), Seq: int(b[2])<<8 | int(b[3])}
-	if bodylen > 4 {
-		p.Data = make([]byte, bodylen-4)
-		copy(p.Data, b[4:])
-	}
-	return p, nil
-}