Peer.go | 9 +++++---- Peers.go | 7 ++++--- client_test.go | 13 ------------- connection.go | 7 +++++-- peer_protocol/pex.go | 26 ++++++++++++++++++++++++++ pex.go | 32 -------------------------------- pex_test.go => peer_protocol/pex_test.go | 4 ++-- tracker_scraper.go | 2 +- diff --git a/Peer.go b/Peer.go index 9efed032a8320d70c0c6c6a0cfdceb6dd5ffba21..e4a7a2fb4cfcd6695b79d9567cf921a16b1680ef 100644 --- a/Peer.go +++ b/Peer.go @@ -4,6 +4,7 @@ import ( "net" "github.com/anacrolix/dht/krpc" + "github.com/anacrolix/torrent/peer_protocol" ) type Peer struct { @@ -13,18 +14,18 @@ Port int Source peerSource // Peer is known to support encryption. SupportsEncryption bool - pexPeerFlags + peer_protocol.PexPeerFlags } -func (me *Peer) FromPex(na krpc.NodeAddr, fs pexPeerFlags) { +func (me *Peer) FromPex(na krpc.NodeAddr, fs peer_protocol.PexPeerFlags) { me.IP = append([]byte(nil), na.IP...) me.Port = na.Port me.Source = peerSourcePEX // If they prefer encryption, they must support it. - if fs.Get(pexPrefersEncryption) { + if fs.Get(peer_protocol.PexPrefersEncryption) { me.SupportsEncryption = true } - me.pexPeerFlags = fs + me.PexPeerFlags = fs } func (me Peer) addr() ipPort { diff --git a/Peers.go b/Peers.go index 12dda6bccbfc351cebb9cc8e5298b359c7319a3b..9c3e062cfa0ccde78e10897d02d759bb3a8ad99c 100644 --- a/Peers.go +++ b/Peers.go @@ -3,15 +3,16 @@ import ( "github.com/anacrolix/dht/krpc" + "github.com/anacrolix/torrent/peer_protocol" "github.com/anacrolix/torrent/tracker" ) type Peers []Peer -func (me *Peers) FromPex(nas []krpc.NodeAddr, fs []pexPeerFlags) { +func (me *Peers) AppendFromPex(nas []krpc.NodeAddr, fs []peer_protocol.PexPeerFlags) { for i, na := range nas { var p Peer - var f pexPeerFlags + var f peer_protocol.PexPeerFlags if i < len(fs) { f = fs[i] } @@ -20,7 +21,7 @@ *me = append(*me, p) } } -func (ret Peers) FromTracker(ps []tracker.Peer) Peers { +func (ret Peers) AppendFromTracker(ps []tracker.Peer) Peers { for _, p := range ps { _p := Peer{ IP: p.IP, diff --git a/client_test.go b/client_test.go index 76893d7303402f78fdd3d376a086f3c477606388..d8ada22e007b6103eadefd93f62c41ae9a5c537f 100644 --- a/client_test.go +++ b/client_test.go @@ -124,19 +124,6 @@ tor.cl.mu.Unlock() assert.EqualValues(t, chunkSpec{4, 1}, chunkIndexSpec(2, tor.pieceLength(0), tor.chunkSize)) } -func TestUnmarshalPEXMsg(t *testing.T) { - var m peerExchangeMessage - if err := bencode.Unmarshal([]byte("d5:added12:\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0ce"), &m); err != nil { - t.Fatal(err) - } - if len(m.Added) != 2 { - t.FailNow() - } - if m.Added[0].Port != 0x506 { - t.FailNow() - } -} - func TestReducedDialTimeout(t *testing.T) { cfg := NewDefaultClientConfig() for _, _case := range []struct { diff --git a/connection.go b/connection.go index 8239f84541bda4a9fd4a4a83930dfd2c96344081..5f550ae2a4cfb6d078450bcd913ca78043386474 100644 --- a/connection.go +++ b/connection.go @@ -1248,13 +1248,16 @@ // TODO: Maybe close the connection. Check that we're not // advertising that we support PEX if it's disabled. return nil } - var pexMsg peerExchangeMessage + var pexMsg pp.PexMsg err := bencode.Unmarshal(payload, &pexMsg) if err != nil { return fmt.Errorf("error unmarshalling PEX message: %s", err) } torrent.Add("pex added6 peers received", int64(len(pexMsg.Added6))) - t.addPeers(pexMsg.AddedPeers()) + var peers Peers + peers.AppendFromPex(pexMsg.Added6, pexMsg.Added6Flags) + peers.AppendFromPex(pexMsg.Added, pexMsg.AddedFlags) + t.addPeers(peers) return nil default: return fmt.Errorf("unexpected extended message ID: %v", id) diff --git a/peer_protocol/pex.go b/peer_protocol/pex.go new file mode 100644 index 0000000000000000000000000000000000000000..aeefd3feadac63e9cf71637cdbde3776cbcf81c7 --- /dev/null +++ b/peer_protocol/pex.go @@ -0,0 +1,26 @@ +package peer_protocol + +import "github.com/anacrolix/dht/krpc" + +type PexMsg struct { + Added krpc.CompactIPv4NodeAddrs `bencode:"added"` + AddedFlags []PexPeerFlags `bencode:"added.f"` + Added6 krpc.CompactIPv6NodeAddrs `bencode:"added6"` + Added6Flags []PexPeerFlags `bencode:"added6.f"` + Dropped krpc.CompactIPv4NodeAddrs `bencode:"dropped"` + Dropped6 krpc.CompactIPv6NodeAddrs `bencode:"dropped6"` +} + +type PexPeerFlags byte + +func (me PexPeerFlags) Get(f PexPeerFlags) bool { + return me&f == f +} + +const ( + PexPrefersEncryption = 0x01 + PexSeedUploadOnly = 0x02 + PexSupportsUtp = 0x04 + PexHolepunchSupport = 0x08 + PexOutgoingConn = 0x10 +) diff --git a/pex.go b/pex.go deleted file mode 100644 index 7b28fe6f62950b35479c4156a561de2f01bf739b..0000000000000000000000000000000000000000 --- a/pex.go +++ /dev/null @@ -1,32 +0,0 @@ -package torrent - -import "github.com/anacrolix/dht/krpc" - -type peerExchangeMessage struct { - Added krpc.CompactIPv4NodeAddrs `bencode:"added"` - AddedFlags []pexPeerFlags `bencode:"added.f"` - Added6 krpc.CompactIPv6NodeAddrs `bencode:"added6"` - Added6Flags []pexPeerFlags `bencode:"added6.f"` - Dropped krpc.CompactIPv4NodeAddrs `bencode:"dropped"` - Dropped6 krpc.CompactIPv6NodeAddrs `bencode:"dropped6"` -} - -type pexPeerFlags byte - -func (me pexPeerFlags) Get(f pexPeerFlags) bool { - return me&f == f -} - -const ( - pexPrefersEncryption = 0x01 - pexSeedUploadOnly = 0x02 - pexSupportsUtp = 0x04 - pexHolepunchSupport = 0x08 - pexOutgoingConn = 0x10 -) - -func (me *peerExchangeMessage) AddedPeers() (ret Peers) { - ret.FromPex(me.Added, me.AddedFlags) - ret.FromPex(me.Added6, me.Added6Flags) - return -} diff --git a/pex_test.go b/peer_protocol/pex_test.go rename from pex_test.go rename to peer_protocol/pex_test.go index f73425fcea9c02a5ddb2b1d020b41edc011bd3b9..054c32beb9f6e79ff2d5ad04760ffeefff53925f 100644 --- a/pex_test.go +++ b/peer_protocol/pex_test.go @@ -1,4 +1,4 @@ -package torrent +package peer_protocol import ( "testing" @@ -9,7 +9,7 @@ "github.com/anacrolix/torrent/bencode" ) func TestUnmarshalPex(t *testing.T) { - var pem peerExchangeMessage + var pem PexMsg err := bencode.Unmarshal([]byte("d5:added12:\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0ce"), &pem) require.NoError(t, err) require.EqualValues(t, 2, len(pem.Added)) diff --git a/tracker_scraper.go b/tracker_scraper.go index ec94ee9eaef657fa7232f1871f72e893fd14ebc1..39359dbe464e80b0ec747c519b891a08438e1fc0 100644 --- a/tracker_scraper.go +++ b/tracker_scraper.go @@ -124,7 +124,7 @@ if err != nil { ret.Err = fmt.Errorf("error announcing: %s", err) return } - me.t.AddPeers(Peers(nil).FromTracker(res.Peers)) + me.t.AddPeers(Peers(nil).AppendFromTracker(res.Peers)) ret.NumPeers = len(res.Peers) ret.Interval = time.Duration(res.Interval) * time.Second return