From 4f02de3a445f3a9eb46c49a8964c9660bdf6e5d7 Mon Sep 17 00:00:00 2001
From: Paul-Louis Ageneau <paul-louis@ageneau.org>
Date: Mon, 17 Jan 2022 22:03:39 +0100
Subject: feat: add BEP6 Fast Extension support (#2243)

* chore: bump bittorrent-protocol to ^3.5.0

* feat: implement BEP6 Fast Extension
---
 lib/peer.js    |  3 ++-
 lib/torrent.js | 80 +++++++++++++++++++++++++++++++++++++++++++++++++++++++---
 package.json   |  2 +-
 3 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/lib/peer.js b/lib/peer.js
index 19cb5ff..4c63269 100644
--- a/lib/peer.js
+++ b/lib/peer.js
@@ -327,7 +327,8 @@ class Peer extends EventEmitter {
 
   handshake () {
     const opts = {
-      dht: this.swarm.private ? false : !!this.swarm.client.dht
+      dht: this.swarm.private ? false : !!this.swarm.client.dht,
+      fast: true
     }
     this.wire.handshake(this.swarm.infoHash, this.swarm.client.peerId, opts)
     this.sentHandshake = true
diff --git a/lib/torrent.js b/lib/torrent.js
index ac7cb97..1bb0062 100644
--- a/lib/torrent.js
+++ b/lib/torrent.js
@@ -650,6 +650,28 @@ class Torrent extends EventEmitter {
     this.bitfield.set(index, true)
   }
 
+  _hasAllPieces () {
+    for (let index = 0; index < this.pieces.length; index++) {
+      if (!this.bitfield.get(index)) return false
+    }
+    return true
+  }
+
+  _hasNoPieces () {
+    return !this._hasMorePieces(0)
+  }
+
+  _hasMorePieces (threshold) {
+    let count = 0
+    for (let index = 0; index < this.pieces.length; index++) {
+      if (this.bitfield.get(index)) {
+        count += 1
+        if (count > threshold) return true
+      }
+    }
+    return false
+  }
+
   /**
    * Called when the metadata, listening server, and underlying chunk store is initialized.
    */
@@ -1139,6 +1161,26 @@ class Torrent extends EventEmitter {
       this._updateWireInterest(wire)
     })
 
+    // fast extension (BEP6)
+    wire.on('have-all', () => {
+      wire.isSeeder = true
+      wire.choke() // always choke seeders
+      this._update()
+      this._updateWireInterest(wire)
+    })
+
+    // fast extension (BEP6)
+    wire.on('have-none', () => {
+      wire.isSeeder = false
+      this._update()
+      this._updateWireInterest(wire)
+    })
+
+    // fast extension (BEP6)
+    wire.on('allowed-fast', (index) => {
+      this._update()
+    })
+
     wire.once('interested', () => {
       wire.unchoke()
     })
@@ -1167,7 +1209,10 @@ class Torrent extends EventEmitter {
       this.store.get(index, { offset, length }, cb)
     })
 
-    wire.bitfield(this.bitfield) // always send bitfield (required)
+    // always send bitfield or equivalent fast extension message (required)
+    if (wire.hasFast && this._hasAllPieces()) wire.haveAll()
+    else if (wire.hasFast && this._hasNoPieces()) wire.haveNone()
+    else wire.bitfield(this.bitfield)
 
     // initialize interest in case bitfield message was already received before above handler was registered
     this._updateWireInterest(wire)
@@ -1284,15 +1329,42 @@ class Torrent extends EventEmitter {
     // to allow function hoisting
     const self = this
 
-    if (wire.peerChoking) return
-    if (!wire.downloaded) return validateWire()
-
     const minOutstandingRequests = getBlockPipelineLength(wire, PIPELINE_MIN_DURATION)
     if (wire.requests.length >= minOutstandingRequests) return
     const maxOutstandingRequests = getBlockPipelineLength(wire, PIPELINE_MAX_DURATION)
 
+    if (wire.peerChoking) {
+      if (wire.hasFast && wire.peerAllowedFastSet.length > 0 &&
+        !this._hasMorePieces(wire.peerAllowedFastSet.length - 1)) {
+        requestAllowedFastSet()
+      }
+      return
+    }
+
+    if (!wire.downloaded) return validateWire()
+
     trySelectWire(false) || trySelectWire(true)
 
+    function requestAllowedFastSet () {
+      if (wire.requests.length >= maxOutstandingRequests) return false
+
+      for (const piece of wire.peerAllowedFastSet) {
+        if (wire.peerPieces.get(piece) && !self.bitfield.get(piece)) {
+          while (self._request(wire, piece, false) &&
+            wire.requests.length < maxOutstandingRequests) {
+            // body intentionally empty
+            // request all non-reserved blocks in this piece
+          }
+        }
+
+        if (wire.requests.length < maxOutstandingRequests) continue
+
+        return true
+      }
+
+      return false
+    }
+
     function genPieceFilterFunc (start, end, tried, rank) {
       return i => i >= start && i <= end && !(i in tried) && wire.peerPieces.get(i) && (!rank || rank(i))
     }
diff --git a/package.json b/package.json
index eaa657d..0b186b1 100644
--- a/package.json
+++ b/package.json
@@ -40,7 +40,7 @@
     "addr-to-ip-port": "^1.5.4",
     "bitfield": "^4.0.0",
     "bittorrent-dht": "^10.0.2",
-    "bittorrent-protocol": "^3.4.3",
+    "bittorrent-protocol": "^3.5.0",
     "cache-chunk-store": "^3.2.2",
     "chrome-net": "^3.3.4",
     "chunk-store-stream": "^4.3.0",
-- 
cgit v1.2.3