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) + 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)) +}