Move PEX stuff into peer_protocol

This commit is contained in:
Matt Joiner 2018-07-10 11:21:24 +10:00
parent e3cb94006c
commit 50def7a688
8 changed files with 43 additions and 57 deletions

View File

@ -4,6 +4,7 @@ import (
"net"
"github.com/anacrolix/dht/krpc"
"github.com/anacrolix/torrent/peer_protocol"
)
type Peer struct {
@ -13,18 +14,18 @@ type Peer struct {
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 {

View File

@ -3,15 +3,16 @@ package torrent
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 @@ func (me *Peers) FromPex(nas []krpc.NodeAddr, fs []pexPeerFlags) {
}
}
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,

View File

@ -124,19 +124,6 @@ func TestTorrentInitialState(t *testing.T) {
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 {

View File

@ -1248,13 +1248,16 @@ func (c *connection) onReadExtendedMsg(id byte, payload []byte) (err error) {
// 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)

26
peer_protocol/pex.go Normal file
View File

@ -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
)

View File

@ -1,4 +1,4 @@
package torrent
package peer_protocol
import (
"testing"
@ -9,7 +9,7 @@ import (
)
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))

32
pex.go
View File

@ -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
}

View File

@ -124,7 +124,7 @@ func (me *trackerScraper) announce() (ret trackerAnnounceResult) {
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