feat: Add PE/MSE support (#1820)

* reimplement #1384

* add option for secure - performance issue

* fixes after code review

* fix error

* use const in hex

* use stored hash

* Update lib/peer.js

Co-authored-by: Diego Rodríguez Baquero <github@diegorbaquero.com>

* fix const

Co-authored-by: Diego Rodríguez Baquero <github@diegorbaquero.com>

2 files changed, 88 insertions, 5 deletions

diff --git a/lib/conn-pool.js b/lib/conn-pool.js
index 9ef02a6..52fb9b5 100644
--- a/lib/conn-pool.js
+++ b/lib/conn-pool.js
@@ -136,15 +136,30 @@ class ConnPool {
       : Peer.createTCPIncomingPeer(conn, this._client.throttleGroups)
 
     const wire = peer.wire
+    wire.once('pe3', onPe3)
     wire.once('handshake', onHandshake)
 
+    function onPe3 (infoHashHash) {
+      const torrent = self._client._getByHash(infoHashHash)
+      if (torrent) {
+        peer.swarm = torrent
+        torrent._addIncomingPeer(peer)
+        peer.onPe3(infoHashHash)
+      } else {
+        peer.destroy(new Error(`Unexpected info hash hash ${infoHashHash} from incoming peer ${peer.id}`))
+      }
+    }
+
     function onHandshake (infoHash, peerId) {
       cleanupPending()
 
       const torrent = self._client.get(infoHash)
+      // only add incoming peer if didn't already do so in protocol encryption handshake
       if (torrent) {
-        peer.swarm = torrent
-        torrent._addIncomingPeer(peer)
+        if (!peer.swarm) {
+          peer.swarm = torrent
+          torrent._addIncomingPeer(peer)
+        }
         peer.onHandshake(infoHash, peerId)
       } else {
         const err = new Error(
diff --git a/lib/peer.js b/lib/peer.js
index dd048fc..1116f32 100644
--- a/lib/peer.js
+++ b/lib/peer.js
@@ -9,6 +9,12 @@ const CONNECT_TIMEOUT_UTP = 5000
 const CONNECT_TIMEOUT_WEBRTC = 25000
 const HANDSHAKE_TIMEOUT = 25000
 
+let secure = false
+
+exports.enableSecure = () => {
+  secure = true
+}
+
 /**
  * WebRTC peer connections start out connected, because WebRTC peers require an
  * "introduction" (i.e. WebRTC signaling), and there's no equivalent to an IP address
@@ -137,6 +143,10 @@ class Peer extends EventEmitter {
     this.timeout = null // handshake timeout
     this.retries = 0 // outgoing TCP connection retry count
 
+    this.sentPe1 = false
+    this.sentPe2 = false
+    this.sentPe3 = false
+    this.sentPe4 = false
     this.sentHandshake = false
   }
 
@@ -166,8 +176,8 @@ class Peer extends EventEmitter {
       this.destroy(err)
     })
 
-    const wire = this.wire = new Wire()
-    wire.type = this.type
+    const wire = this.wire = new Wire(this.type, this.retries, secure)
+
     wire.once('end', () => {
       this.destroy()
     })
@@ -181,6 +191,18 @@ class Peer extends EventEmitter {
       this.destroy(err)
     })
 
+    wire.once('pe1', () => {
+      this.onPe1()
+    })
+    wire.once('pe2', () => {
+      this.onPe2()
+    })
+    wire.once('pe3', () => {
+      this.onPe3()
+    })
+    wire.once('pe4', () => {
+      this.onPe4()
+    })
     wire.once('handshake', (infoHash, peerId) => {
       this.onHandshake(infoHash, peerId)
     })
@@ -188,7 +210,53 @@ class Peer extends EventEmitter {
 
     this.setThrottlePipes()
 
-    if (this.swarm && !this.sentHandshake) this.handshake()
+    if (this.swarm) {
+      if (this.type === 'tcpOutgoing') {
+        if (secure && this.retries === 0 && !this.sentPe1) this.sendPe1()
+        else if (!this.sentHandshake) this.handshake()
+      } else if (this.type !== 'tcpIncoming' && !this.sentHandshake) this.handshake()
+    }
+  }
+
+  sendPe1 () {
+    this.wire.sendPe1()
+    this.sentPe1 = true
+  }
+
+  onPe1 () {
+    this.sendPe2()
+  }
+
+  sendPe2 () {
+    this.wire.sendPe2()
+    this.sentPe2 = true
+  }
+
+  onPe2 () {
+    this.sendPe3()
+  }
+
+  sendPe3 () {
+    this.wire.sendPe3(this.swarm.infoHash)
+    this.sentPe3 = true
+  }
+
+  onPe3 (infoHashHash) {
+    if (this.swarm) {
+      if (this.swarm.infoHashHash !== infoHashHash) {
+        this.destroy(new Error('unexpected crypto handshake info hash for this swarm'))
+      }
+      this.sendPe4()
+    }
+  }
+
+  sendPe4 () {
+    this.wire.sendPe4(this.swarm.infoHash)
+    this.sentPe4 = true
+  }
+
+  onPe4 () {
+    if (!this.sentHandshake) this.handshake()
   }
 
   clearPipes () {