dht: Support IPv6 compact peers

diff --git a/cmd/dht-get-peers/main.go b/cmd/dht-get-peers/main.go
index d1f3a116790617026d30a82236879f6b38a5d919..646104d9b4f3a65d7a850f5ce32bd5e4ab2468ba 100644 (file)
--- 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 2f0c364a421e2df5d0a6287c21e4294aa902a942..fd845f1b5ad2f4fec74373555d194e4235dd5f2c 100644 (file)
--- 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 1ffff796c667dbbe405bc8c077a900b6188caa13..68135f8573091d6d874507e55016752eac813f0f 100644 (file)
--- 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))
+}