dht: Support IPv6 compact peers

cmd/dht-get-peers/main.go

@@ -12,7 +12,6 @@ import (
 	_ "github.com/anacrolix/envpprof"
 
 	"github.com/anacrolix/torrent/dht"
-	"github.com/anacrolix/torrent/util"
 	_ "github.com/anacrolix/torrent/util/profile"
 )
 
@@ -126,7 +125,7 @@ func setupSignals() {
 }
 
 func main() {
-	seen := make(map[util.CompactPeer]struct{})
+	seen := make(map[string]struct{})
 getPeers:
 	for {
 		ps, err := s.Announce(*infoHash, 0, false)
@@ -142,10 +141,10 @@ getPeers:
 				}
 				log.Printf("received %d peers from %x", len(v.Peers), v.NodeInfo.ID)
 				for _, p := range v.Peers {
-					if _, ok := seen[p]; ok {
+					if _, ok := seen[p.String()]; ok {
 						continue
 					}
-					seen[p] = struct{}{}
+					seen[p.String()] = struct{}{}
 					fmt.Println((&net.UDPAddr{
 						IP:   p.IP[:],
 						Port: int(p.Port),

tracker/udp_test.go

@@ -31,7 +31,7 @@ func TestNetIPv4Bytes(t *testing.T) {
 
 func TestMarshalAnnounceResponse(t *testing.T) {
 	w := bytes.Buffer{}
-	peers := util.CompactPeers{{[4]byte{127, 0, 0, 1}, 2}, {[4]byte{255, 0, 0, 3}, 4}}
+	peers := util.CompactPeers{{[]byte{127, 0, 0, 1}, 2}, {[]byte{255, 0, 0, 3}, 4}}
 	err := peers.WriteBinary(&w)
 	if err != nil {
 		t.Fatalf("error writing udp announce response addrs: %s", err)

util/types.go

@@ -1,11 +1,12 @@
 package util
 
 import (
-	"bytes"
 	"encoding"
 	"encoding/binary"
 	"fmt"
 	"io"
+	"net"
+	"strconv"
 
 	"github.com/anacrolix/torrent/bencode"
 )
@@ -34,25 +35,53 @@ func (me *CompactPeers) UnmarshalBinary(b []byte) (err error) {
 	return
 }
 
-func (me CompactPeers) WriteBinary(w io.Writer) error {
-	return binary.Write(w, binary.BigEndian, me)
+func (me CompactPeers) WriteBinary(w io.Writer) (err error) {
+	for _, cp := range me {
+		cp.Write(w)
+		if err != nil {
+			return
+		}
+	}
+	return
 }
 
 type CompactPeer struct {
-	IP   [4]byte
+	IP   net.IP
 	Port uint16
 }
 
 var _ encoding.BinaryUnmarshaler = &CompactPeer{}
 
 func (cp *CompactPeer) UnmarshalBinary(b []byte) (err error) {
-	r := bytes.NewReader(b)
-	err = binary.Read(r, binary.BigEndian, cp)
+	switch len(b) {
+	case 18:
+		cp.IP = make([]byte, 16)
+	case 6:
+		cp.IP = make([]byte, 4)
+	default:
+		err = fmt.Errorf("bad length: %d", len(b))
+		return
+	}
+	if n := copy(cp.IP, b); n != len(cp.IP) {
+		panic(n)
+	}
+	b = b[len(cp.IP):]
+	if len(b) != 2 {
+		panic(len(b))
+	}
+	cp.Port = binary.BigEndian.Uint16(b)
+	return
+}
+
+func (cp *CompactPeer) Write(w io.Writer) (err error) {
+	_, err = w.Write(cp.IP)
 	if err != nil {
 		return
 	}
-	if r.Len() != 0 {
-		err = fmt.Errorf("%d bytes unused", r.Len())
-	}
+	err = binary.Write(w, binary.BigEndian, cp.Port)
 	return
 }
+
+func (cp *CompactPeer) String() string {
+	return net.JoinHostPort(cp.IP.String(), strconv.FormatUint(uint64(cp.Port), 10))
+}