From 146d4d46dd10588a116a192335637cb243c1f07b Mon Sep 17 00:00:00 2001
From: Yoann Ciabaud <yoann@sonora.io>
Date: Fri, 24 Feb 2017 16:18:32 +0100
Subject: [PATCH] Extension point on requests and responses

---
 server.js               | 31 +++++++++++++-----
 test/request-handler.js | 70 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 8 deletions(-)
 create mode 100644 test/request-handler.js

diff --git a/server.js b/server.js
index 3723516..abb2d2d 100644
--- a/server.js
+++ b/server.js
@@ -27,14 +27,15 @@ inherits(Server, EventEmitter)
  * metrics from clients that help the tracker keep overall statistics about the torrent.
  * Responses include a peer list that helps the client participate in the torrent.
  *
- * @param {Object}  opts            options object
- * @param {Number}  opts.interval   tell clients to announce on this interval (ms)
- * @param {Number}  opts.trustProxy trust 'x-forwarded-for' header from reverse proxy
- * @param {boolean} opts.http       start an http server? (default: true)
- * @param {boolean} opts.udp        start a udp server? (default: true)
- * @param {boolean} opts.ws         start a websocket server? (default: true)
- * @param {boolean} opts.stats      enable web-based statistics? (default: true)
- * @param {function} opts.filter    black/whitelist fn for disallowing/allowing torrents
+ * @param {Object}  opts                  options object
+ * @param {Number}  opts.interval         tell clients to announce on this interval (ms)
+ * @param {Number}  opts.trustProxy       trust 'x-forwarded-for' header from reverse proxy
+ * @param {boolean} opts.http             start an http server? (default: true)
+ * @param {boolean} opts.udp              start a udp server? (default: true)
+ * @param {boolean} opts.ws               start a websocket server? (default: true)
+ * @param {boolean} opts.stats            enable web-based statistics? (default: true)
+ * @param {function} opts.filter          black/whitelist fn for disallowing/allowing torrents
+ * @param {function} opts.requestHandler  functions to handle params / response
  */
 function Server (opts) {
   var self = this
@@ -64,6 +65,18 @@ function Server (opts) {
   self.udp6 = null
   self.ws = null
 
+  self._reqHandler = opts.requestHandler || {}
+  if (!self._reqHandler.getParams) {
+    self._reqHandler.getParams = function (params) {
+      return params
+    }
+  }
+  if (!self._reqHandler.getResponse) {
+    self._reqHandler.getResponse = function (params, cb) {
+      return cb
+    }
+  }
+
   // start an http tracker unless the user explictly says no
   if (opts.http !== false) {
     self.http = http.createServer()
@@ -636,6 +649,8 @@ Server.prototype._onWebSocketError = function (socket, err) {
 
 Server.prototype._onRequest = function (params, cb) {
   var self = this
+  params = self._reqHandler.getParams(params)
+  cb = self._reqHandler.getResponse(params, cb)
   if (params && params.action === common.ACTIONS.CONNECT) {
     cb(null, { action: common.ACTIONS.CONNECT })
   } else if (params && params.action === common.ACTIONS.ANNOUNCE) {
diff --git a/test/request-handler.js b/test/request-handler.js
new file mode 100644
index 0000000..f186230
--- /dev/null
+++ b/test/request-handler.js
@@ -0,0 +1,70 @@
+var Buffer = require('safe-buffer').Buffer
+var Client = require('../')
+var common = require('./common')
+var fixtures = require('webtorrent-fixtures')
+var test = require('tape')
+
+var peerId = Buffer.from('01234567890123456789')
+
+function testRequestHandler (t, serverType) {
+  t.plan(4)
+
+  var opts = { serverType: serverType } // this is test-suite-only option
+  opts.requestHandler = {
+    getParams: function (params) {
+      params.extra = 123
+      return params
+    },
+    getResponse: function (params, cb) {
+      return function (err, response) {
+        response.complete = params.extra * 2
+        cb(err, response)
+      }
+    }
+  }
+
+  common.createServer(t, opts, function (server, announceUrl) {
+    var client1 = new Client({
+      infoHash: fixtures.alice.parsedTorrent.infoHash,
+      announce: announceUrl,
+      peerId: peerId,
+      port: 6881,
+      wrtc: {}
+    })
+
+    client1.on('error', function (err) { t.error(err) })
+    if (serverType === 'ws') common.mockWebsocketTracker(client1)
+
+    server.once('start', function () {
+      t.pass('got start message from client1')
+    })
+
+    client1.once('update', function (data) {
+      console.log(data)
+
+      t.equal(data.complete, 246)
+
+      client1.destroy(function () {
+        t.pass('client1 destroyed')
+      })
+
+      server.close(function () {
+        t.pass('server destroyed')
+      })
+    })
+
+    client1.start()
+  })
+}
+
+test('http: request handler option intercepts announce requests and responses', function (t) {
+  testRequestHandler(t, 'http')
+})
+
+test('udp: request handler option intercepts announce requests and responses', function (t) {
+  testRequestHandler(t, 'udp')
+})
+
+test('ws: request handler option intercepts announce requests and responses', function (t) {
+  testRequestHandler(t, 'ws')
+})