extension support

README.md

@@ -266,7 +266,8 @@ If `opts` is specified, then the default options (shown below) will be overridde
   rtcConfig: Object,     // RTCPeerConnection configuration object (default=STUN only)
   storage: Function,     // custom storage engine, or `false` to use in-memory engine
   tracker: Boolean,      // Whether or not to enable trackers (default=true)
-  wrtc: {}               // custom webrtc implementation (in node, specify the [wrtc](https://www.npmjs.com/package/wrtc) package)
+  wrtc: {},              // custom webrtc implementation (in node, specify the [wrtc](https://www.npmjs.com/package/wrtc) package)
+  onWire: Function       // function to be called on each new wire. Use this to specify [custom bittorrent extensions](https://www.npmjs.com/package/bittorrent-protocol#extension-api)
 }
 ```
 

index.js

@@ -43,6 +43,7 @@ function WebTorrent (opts) {
   self.destroyed = false
   self.torrentPort = opts.torrentPort || 0
   self.tracker = opts.tracker !== undefined ? opts.tracker : true
+  self.onWire = opts.onWire
 
   self._rtcConfig = opts.rtcConfig
   self._wrtc = opts.wrtc || global.WRTC // to support `webtorrent-hybrid` package
@@ -156,6 +157,7 @@ WebTorrent.prototype.download = function (torrentId, opts, ontorrent) {
   }
   if (!opts) opts = {}
   if (!opts.storage) opts.storage = self.storage
+  if (!opts.onWire) opts.onWire = self.onWire
   opts.client = self
 
   var torrent = self.get(torrentId)

lib/torrent.js

@@ -501,6 +501,11 @@ Torrent.prototype._onWire = function (wire, addr) {
     })
   }
 
+  // trigger opts.onWire to allow user-defined extensions
+  if (typeof this.opts.onWire === 'function') {
+    this.opts.onWire.call(this, wire)
+  }
+
   // Send KEEP-ALIVE (every 60s) so peers will not disconnect the wire
   wire.setKeepAlive(true)
 

test/extensions.js

@@ -0,0 +1,54 @@
+var fs = require('fs')
+var parseTorrent = require('parse-torrent')
+var test = require('tape')
+var WebTorrent = require('../')
+
+var leaves = fs.readFileSync(__dirname + '/torrents/leaves.torrent')
+var leavesTorrent = parseTorrent(leaves)
+
+test('onWire option', {timeout: 500}, function (t) {
+  t.plan(6)
+  var extendedHandshakes = 0
+
+  function Extension (wire) {
+    wire.extendedHandshake.test = 'Hello, World!'
+  }
+
+  Extension.prototype.name = 'wt_test'
+  Extension.prototype.onExtendedHandshake = function (handshake) {
+    t.equal(handshake.test.toString(), 'Hello, World!', 'handshake.test === Hello, World!')
+    extendedHandshakes++
+    if (extendedHandshakes === 2) {
+      client1.destroy(function () {
+        t.pass('client1 destroyed')
+      })
+      client2.destroy(function () {
+        t.pass('client1 destroyed')
+      })
+    }
+  }
+
+  var client1 = new WebTorrent({
+    dht: false,
+    tracker: false,
+    onWire: function (wire) {
+      t.pass('client1 onWire')
+      wire.use(Extension)
+    }
+  })
+  var client2 = new WebTorrent({
+    dht: false,
+    tracker: false,
+    onWire: function (wire) {
+      t.pass('client2 onWire')
+      wire.use(Extension)
+    }
+  })
+
+  client1.add(leavesTorrent, function (torrent1) {
+    client2.add(leavesTorrent.infoHash)
+    client2.on('listening', function (port, torrent2) {
+      torrent2.addPeer('127.0.0.1:' + client1.torrentPort)
+    })
+  })
+})