From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 12 Feb 2018 13:50:32 +0000 (+1100)
Subject: Support IPv6 peers over PEX
X-Git-Tag: v1.0.0~178
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=70010ce6917857752a2808c06fa01d00891fc6a5;p=btrtrc.git

Support IPv6 peers over PEX
---

diff --git a/Peers.go b/Peers.go
new file mode 100644
index 00000000..72f7e8e1
--- /dev/null
+++ b/Peers.go
@@ -0,0 +1,17 @@
+package torrent
+
+import "github.com/anacrolix/dht/krpc"
+
+type Peers []Peer
+
+func (me *Peers) FromPex(nas []krpc.NodeAddr, fs []pexPeerFlags) {
+	for i, na := range nas {
+		var p Peer
+		var f pexPeerFlags
+		if i < len(fs) {
+			f = fs[i]
+		}
+		p.FromPex(na, f)
+		*me = append(*me, p)
+	}
+}
diff --git a/connection.go b/connection.go
index 379b45f9..4e919c5f 100644
--- a/connection.go
+++ b/connection.go
@@ -1142,6 +1142,8 @@ func (c *connection) onReadExtendedMsg(id byte, payload []byte) (err error) {
 		return nil
 	case pexExtendedId:
 		if cl.config.DisablePEX {
+			// TODO: Maybe close the connection. Check that we're not
+			// advertising that we support PEX if it's disabled.
 			return nil
 		}
 		var pexMsg peerExchangeMessage
@@ -1150,21 +1152,9 @@ func (c *connection) onReadExtendedMsg(id byte, payload []byte) (err error) {
 			return fmt.Errorf("error unmarshalling PEX message: %s", err)
 		}
 		go func() {
+			ps := pexMsg.AddedPeers()
 			cl.mu.Lock()
-			t.addPeers(func() (ret []Peer) {
-				for i, cp := range pexMsg.Added {
-					p := Peer{
-						IP:     append(make(net.IP, 0, 4), cp.IP...),
-						Port:   cp.Port,
-						Source: peerSourcePEX,
-					}
-					if i < len(pexMsg.AddedFlags) && pexMsg.AddedFlags[i]&0x01 != 0 {
-						p.SupportsEncryption = true
-					}
-					ret = append(ret, p)
-				}
-				return
-			}())
+			t.addPeers(ps)
 			cl.mu.Unlock()
 		}()
 		return nil
diff --git a/pex.go b/pex.go
index b3e3e166..7b28fe6f 100644
--- a/pex.go
+++ b/pex.go
@@ -6,13 +6,17 @@ type peerExchangeMessage struct {
 	Added       krpc.CompactIPv4NodeAddrs `bencode:"added"`
 	AddedFlags  []pexPeerFlags            `bencode:"added.f"`
 	Added6      krpc.CompactIPv6NodeAddrs `bencode:"added6"`
-	AddedFlags6 []pexPeerFlags            `bencode:"added6.f"`
+	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
@@ -20,3 +24,9 @@ const (
 	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/torrent.go b/torrent.go
index f058a72f..95e7aca1 100644
--- a/torrent.go
+++ b/torrent.go
@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/anacrolix/dht"
+	"github.com/anacrolix/dht/krpc"
 	"github.com/anacrolix/log"
 	"github.com/anacrolix/missinggo"
 	"github.com/anacrolix/missinggo/bitmap"
@@ -738,6 +739,18 @@ type Peer struct {
 	Source peerSource
 	// Peer is known to support encryption.
 	SupportsEncryption bool
+	pexPeerFlags
+}
+
+func (me *Peer) FromPex(na krpc.NodeAddr, fs 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) {
+		me.SupportsEncryption = true
+	}
+	me.pexPeerFlags = fs
 }
 
 func (t *Torrent) pieceLength(piece int) pp.Integer {