From 1cc5a511bd8a5e698b1f5343ad54f99d1cbe488e Mon Sep 17 00:00:00 2001
From: Feross Aboukhadijeh <feross@feross.org>
Date: Fri, 1 May 2015 17:36:07 -0700
Subject: [PATCH] udp server: support multiple info_hash scrape

Fixes #33
---
 client.js           | 35 +++++++++++++++++++++-------
 lib/http-tracker.js | 26 ++++++++++++++-------
 lib/parse_http.js   |  2 +-
 lib/parse_udp.js    | 10 ++++----
 lib/udp-tracker.js  | 39 +++++++++++++++++++------------
 server.js           | 24 +++++++++----------
 test/scrape.js      | 56 +++++++++++++++++++++++++++++++++++++--------
 7 files changed, 135 insertions(+), 57 deletions(-)

diff --git a/client.js b/client.js
index 4884696..77fca84 100644
--- a/client.js
+++ b/client.js
@@ -71,10 +71,11 @@ function Client (peerId, port, torrent, opts) {
 }
 
 /**
- * Simple convenience function to scrape a tracker for an infoHash without
- * needing to create a Client, pass it a parsed torrent, etc.
- * @param  {string}   announceUrl
- * @param  {string}   infoHash
+ * Simple convenience function to scrape a tracker for an info hash without needing to
+ * create a Client, pass it a parsed torrent, etc. Support scraping a tracker for multiple
+ * torrents at the same time.
+ * @param  {string} announceUrl
+ * @param  {string|Array.<string>} infoHash
  * @param  {function} cb
  */
 Client.scrape = function (announceUrl, infoHash, cb) {
@@ -83,15 +84,33 @@ Client.scrape = function (announceUrl, infoHash, cb) {
   var peerId = new Buffer('01234567890123456789') // dummy value
   var port = 6881 // dummy value
   var torrent = {
-    infoHash: infoHash,
+    infoHash: Array.isArray(infoHash) ? infoHash[0] : infoHash,
     announce: [ announceUrl ]
   }
   var client = new Client(peerId, port, torrent)
   client.once('error', cb)
-  client.once('scrape', function (data) {
-    cb(null, data)
+
+  var len = Array.isArray(infoHash) ? infoHash.length : 1
+  var results = {}
+  client.on('scrape', function (data) {
+    len -= 1
+    results[data.infoHash] = data
+    if (len === 0) {
+      client.destroy()
+      var keys = Object.keys(results)
+      if (keys.length === 1) {
+        cb(null, results[keys[0]])
+      } else {
+        cb(null, results)
+      }
+    }
   })
-  client.scrape()
+
+  infoHash = Array.isArray(infoHash)
+    ? infoHash.map(function (infoHash) { return new Buffer(infoHash, 'hex') })
+    : new Buffer(infoHash, 'hex')
+  client.scrape({ infoHash: infoHash })
+
 }
 
 /**
diff --git a/lib/http-tracker.js b/lib/http-tracker.js
index 6095e77..45e1d9b 100644
--- a/lib/http-tracker.js
+++ b/lib/http-tracker.js
@@ -63,7 +63,12 @@ HTTPTracker.prototype.scrape = function (opts) {
     return
   }
 
-  opts.info_hash = self.client._infoHash.toString('binary')
+  opts.info_hash = (Array.isArray(opts.infoHash) && opts.infoHash.length > 0)
+    ? opts.infoHash.map(function (infoHash) { return infoHash.toString('binary') })
+    : (opts.infoHash || self.client._infoHash).toString('binary')
+
+  if (opts.infoHash) delete opts.infoHash
+
   self._request(self._scrapeUrl, opts, self._onScrapeResponse.bind(self))
 }
 
@@ -186,18 +191,23 @@ HTTPTracker.prototype._onScrapeResponse = function (data) {
   // NOTE: the unofficial spec says to use the 'files' key, 'host' has been
   // seen in practice
   data = data.files || data.host || {}
-  data = data[self.client._infoHash.toString('binary')]
 
-  if (!data) {
+  var keys = Object.keys(data)
+  if (keys.length === 0) {
     self.client.emit('warning', new Error('invalid scrape response'))
-  } else {
+    return
+  }
+
+  keys.forEach(function (infoHash) {
+    var response = data[infoHash]
     // TODO: optionally handle data.flags.min_request_interval
     // (separate from announce interval)
     self.client.emit('scrape', {
       announce: self._announceUrl,
-      complete: data.complete,
-      incomplete: data.incomplete,
-      downloaded: data.downloaded
+      infoHash: common.binaryToHex(infoHash),
+      complete: response.complete,
+      incomplete: response.incomplete,
+      downloaded: response.downloaded
     })
-  }
+  })
 }
diff --git a/lib/parse_http.js b/lib/parse_http.js
index 6e575f8..fa044b4 100644
--- a/lib/parse_http.js
+++ b/lib/parse_http.js
@@ -38,8 +38,8 @@ function parseHttpRequest (req, opts) {
     params.addr = (common.IPV6_RE.test(params.ip) ? '[' + params.ip + ']' : params.ip) + ':' + params.port
   } else if (opts.action === 'scrape' || s[0] === '/scrape') {
     params.action = common.ACTIONS.SCRAPE
-    if (typeof params.info_hash === 'string') params.info_hash = [ params.info_hash ]
 
+    if (typeof params.info_hash === 'string') params.info_hash = [ params.info_hash ]
     if (Array.isArray(params.info_hash)) {
       params.info_hash = params.info_hash.map(function (binaryInfoHash) {
         if (typeof binaryInfoHash !== 'string' || binaryInfoHash.length !== 20) {
diff --git a/lib/parse_udp.js b/lib/parse_udp.js
index e1819fd..0701d94 100644
--- a/lib/parse_udp.js
+++ b/lib/parse_udp.js
@@ -50,10 +50,12 @@ function parseUdpRequest (msg, rinfo) {
     params.addr = params.ip + ':' + params.port // TODO: ipv6 brackets
     params.compact = 1 // udp is always compact
   } else if (params.action === common.ACTIONS.SCRAPE) { // scrape message
-    // TODO: support multiple info_hash scrape
-    if (msg.length > 36) throw new Error('multiple info_hash scrape not supported')
-
-    params.info_hash = [ msg.slice(16, 36).toString('hex') ] // 20 bytes
+    if ((msg.length - 16) % 20 !== 0) throw new Error('invalid scrape message')
+    params.info_hash = []
+    for (var i = 0, len = (msg.length - 16) / 20; i < len; i += 1) {
+      var infoHash = msg.slice(16 + (i * 20), 36 + (i * 20)).toString('hex') // 20 bytes
+      params.info_hash.push(infoHash)
+    }
   } else {
     throw new Error('Invalid action in UDP packet: ' + params.action)
   }
diff --git a/lib/udp-tracker.js b/lib/udp-tracker.js
index 102bdd7..279931e 100644
--- a/lib/udp-tracker.js
+++ b/lib/udp-tracker.js
@@ -35,19 +35,19 @@ function UDPTracker (client, announceUrl, opts) {
 
 UDPTracker.prototype.announce = function (opts) {
   var self = this
-  self._request(self._announceUrl, opts)
+  self._request(opts)
 }
 
 UDPTracker.prototype.scrape = function (opts) {
   var self = this
   opts._scrape = true
-  self._request(self._announceUrl, opts) // udp scrape uses same announce url
+  self._request(opts) // udp scrape uses same announce url
 }
 
-UDPTracker.prototype._request = function (requestUrl, opts) {
+UDPTracker.prototype._request = function (opts) {
   var self = this
   if (!opts) opts = {}
-  var parsedUrl = url.parse(requestUrl)
+  var parsedUrl = url.parse(self._announceUrl)
   var socket = dgram.createSocket('udp4')
   var transactionId = genTransactionId()
 
@@ -78,7 +78,7 @@ UDPTracker.prototype._request = function (requestUrl, opts) {
     }
 
     var action = msg.readUInt32BE(0)
-    debug(requestUrl + ' UDP response, action ' + action)
+    debug(self._announceUrl + ' UDP response, action ' + action)
     switch (action) {
       case 0: // handshake
         if (msg.length < 16) {
@@ -125,15 +125,22 @@ UDPTracker.prototype._request = function (requestUrl, opts) {
 
       case 2: // scrape
         cleanup()
-        if (msg.length < 20) {
+        if (msg.length < 20 || (msg.length - 8) % 12 !== 0) {
           return error('invalid scrape message')
         }
-        self.client.emit('scrape', {
-          announce: self._announceUrl,
-          complete: msg.readUInt32BE(8),
-          downloaded: msg.readUInt32BE(12),
-          incomplete: msg.readUInt32BE(16)
-        })
+        var infoHashes = (Array.isArray(opts.infoHash) && opts.infoHash.length > 0)
+          ? opts.infoHash.map(function (infoHash) { return infoHash.toString('hex') })
+          : (opts.infoHash || self.client._infoHash).toString('hex')
+
+        for (var i = 0, len = (msg.length - 8) / 12; i < len; i += 1) {
+          self.client.emit('scrape', {
+            announce: self._announceUrl,
+            infoHash: infoHashes[i],
+            complete: msg.readUInt32BE(8 + (i * 12)),
+            downloaded: msg.readUInt32BE(12 + (i * 12)),
+            incomplete: msg.readUInt32BE(16 + (i * 12))
+          })
+        }
         break
 
       case 3: // error
@@ -159,7 +166,7 @@ UDPTracker.prototype._request = function (requestUrl, opts) {
 
   function error (message) {
     // errors will often happen if a tracker is offline, so don't treat it as fatal
-    self.client.emit('warning', new Error(message + ' (' + requestUrl + ')'))
+    self.client.emit('warning', new Error(message + ' (' + self._announceUrl + ')'))
     cleanup()
   }
 
@@ -195,11 +202,15 @@ UDPTracker.prototype._request = function (requestUrl, opts) {
   function scrape (connectionId) {
     transactionId = genTransactionId()
 
+    var infoHash = (Array.isArray(opts.infoHash) && opts.infoHash.length > 0)
+      ? Buffer.concat(opts.infoHash)
+      : (opts.infoHash || self.client._infoHash)
+
     send(Buffer.concat([
       connectionId,
       common.toUInt32(common.ACTIONS.SCRAPE),
       transactionId,
-      self.client._infoHash
+      infoHash
     ]))
   }
 }
diff --git a/server.js b/server.js
index 07e41e8..f39dd49 100644
--- a/server.js
+++ b/server.js
@@ -437,19 +437,19 @@ function makeUdpPacket (params) {
       ])
       break
     case common.ACTIONS.SCRAPE:
-      var firstInfoHash = Object.keys(params.files)[0]
-      var scrapeInfo = firstInfoHash ? {
-        complete: params.files[firstInfoHash].complete,
-        incomplete: params.files[firstInfoHash].incomplete,
-        completed: params.files[firstInfoHash].complete // TODO: this only provides a lower-bound
-      } : {}
-      packet = Buffer.concat([
+      var scrapeResponse = [
         common.toUInt32(common.ACTIONS.SCRAPE),
-        common.toUInt32(params.transactionId),
-        common.toUInt32(scrapeInfo.complete),
-        common.toUInt32(scrapeInfo.completed),
-        common.toUInt32(scrapeInfo.incomplete)
-      ])
+        common.toUInt32(params.transactionId)
+      ]
+      for (var infoHash in params.files) {
+        var file = params.files[infoHash]
+        scrapeResponse.push(
+          common.toUInt32(file.complete),
+          common.toUInt32(file.downloaded), // TODO: this only provides a lower-bound
+          common.toUInt32(file.incomplete)
+        )
+      }
+      packet = Buffer.concat(scrapeResponse)
       break
     case common.ACTIONS.ERROR:
       packet = Buffer.concat([
diff --git a/test/scrape.js b/test/scrape.js
index daf45cd..05c9341 100644
--- a/test/scrape.js
+++ b/test/scrape.js
@@ -8,25 +8,33 @@ var parseTorrent = require('parse-torrent')
 var Server = require('../').Server
 var test = require('tape')
 
-function hexToBinary (str) {
-  return new Buffer(str, 'hex').toString('binary')
-}
-
 var infoHash1 = 'aaa67059ed6bd08362da625b3ae77f6f4a075aaa'
-var binaryInfoHash1 = hexToBinary(infoHash1)
+var binaryInfoHash1 = commonLib.hexToBinary(infoHash1)
 var infoHash2 = 'bbb67059ed6bd08362da625b3ae77f6f4a075bbb'
-var binaryInfoHash2 = hexToBinary(infoHash2)
+var binaryInfoHash2 = commonLib.hexToBinary(infoHash2)
 
 var bitlove = fs.readFileSync(__dirname + '/torrents/bitlove-intro.torrent')
 var parsedBitlove = parseTorrent(bitlove)
-var binaryBitlove = hexToBinary(parsedBitlove.infoHash)
+var binaryBitlove = commonLib.hexToBinary(parsedBitlove.infoHash)
 
 var peerId = new Buffer('01234567890123456789')
 
 function testSingle (t, serverType) {
   commonTest.createServer(t, serverType, function (server, announceUrl) {
-    Client.scrape(announceUrl, infoHash1, function (err, data) {
+    parsedBitlove.announce = [ announceUrl ]
+    var client = new Client(peerId, 6881, parsedBitlove)
+
+    client.on('error', function (err) {
       t.error(err)
+    })
+
+    client.on('warning', function (err) {
+      t.error(err)
+    })
+
+    client.scrape()
+
+    client.on('scrape', function (data) {
       t.equal(data.announce, announceUrl)
       t.equal(typeof data.complete, 'number')
       t.equal(typeof data.incomplete, 'number')
@@ -69,9 +77,37 @@ test('udp: scrape using Client.scrape static method', function (t) {
   clientScrapeStatic(t, 'udp')
 })
 
-// TODO: test client for multiple scrape for UDP trackers
+function clientScrapeMulti (t, serverType) {
+  commonTest.createServer(t, serverType, function (server, announceUrl) {
+    Client.scrape(announceUrl, [ infoHash1, infoHash2 ], function (err, results) {
+      t.error(err)
 
-test('server: multiple info_hash scrape', function (t) {
+      t.equal(results[infoHash1].announce, announceUrl)
+      t.equal(typeof results[infoHash1].complete, 'number')
+      t.equal(typeof results[infoHash1].incomplete, 'number')
+      t.equal(typeof results[infoHash1].downloaded, 'number')
+
+      t.equal(results[infoHash2].announce, announceUrl)
+      t.equal(typeof results[infoHash2].complete, 'number')
+      t.equal(typeof results[infoHash2].incomplete, 'number')
+      t.equal(typeof results[infoHash2].downloaded, 'number')
+
+      server.close(function () {
+        t.end()
+      })
+    })
+  })
+}
+
+test('http: MULTI scrape using Client.scrape static method', function (t) {
+  clientScrapeMulti(t, 'http')
+})
+
+test('udp: MULTI scrape using Client.scrape static method', function (t) {
+  clientScrapeMulti(t, 'udp')
+})
+
+test('server: multiple info_hash scrape (manual http request)', function (t) {
   var server = new Server({ udp: false })
   server.on('error', function (err) {
     t.error(err)