From abdf1f049fe959d80395c65f49647f7706d0657d Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 6 Aug 2015 09:01:38 +1000
Subject: [PATCH] dht: Support IPv6 compact peers

---
 cmd/dht-get-peers/main.go |  7 +++---
 tracker/udp_test.go       |  2 +-
 util/types.go             | 47 +++++++++++++++++++++++++++++++--------
 3 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/cmd/dht-get-peers/main.go b/cmd/dht-get-peers/main.go
index d1f3a116..646104d9 100644
--- a/cmd/dht-get-peers/main.go
+++ b/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),
diff --git a/tracker/udp_test.go b/tracker/udp_test.go
index 2f0c364a..fd845f1b 100644
--- a/tracker/udp_test.go
+++ b/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)
diff --git a/util/types.go b/util/types.go
index 1ffff796..68135f85 100644
--- a/util/types.go
+++ b/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)
-	if err != nil {
+	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 r.Len() != 0 {
-		err = fmt.Errorf("%d bytes unused", r.Len())
+	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
+	}
+	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))
+}
-- 
2.44.0