From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 10 Jul 2018 01:21:24 +0000 (+1000)
Subject: Move PEX stuff into peer_protocol
X-Git-Tag: v1.0.0~95
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=50def7a6889505c1837ea8b240d4caa9f8ce1e45;p=btrtrc.git

Move PEX stuff into peer_protocol
---

diff --git a/Peer.go b/Peer.go
index 9efed032..e4a7a2fb 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 @@ 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 {
diff --git a/Peers.go b/Peers.go
index 12dda6bc..9c3e062c 100644
--- a/Peers.go
+++ b/Peers.go
@@ -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,
diff --git a/client_test.go b/client_test.go
index 76893d73..d8ada22e 100644
--- a/client_test.go
+++ b/client_test.go
@@ -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 {
diff --git a/connection.go b/connection.go
index 8239f845..5f550ae2 100644
--- a/connection.go
+++ b/connection.go
@@ -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)
diff --git a/peer_protocol/pex.go b/peer_protocol/pex.go
new file mode 100644
index 00000000..aeefd3fe
--- /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_test.go b/peer_protocol/pex_test.go
similarity index 90%
rename from pex_test.go
rename to peer_protocol/pex_test.go
index f73425fc..054c32be 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 @@ 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))
diff --git a/pex.go b/pex.go
deleted file mode 100644
index 7b28fe6f..00000000
--- 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/tracker_scraper.go b/tracker_scraper.go
index ec94ee9e..39359dbe 100644
--- a/tracker_scraper.go
+++ b/tracker_scraper.go
@@ -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