From 2f37e6c476089d14f0adb2bf170cddb872ee67dd Mon Sep 17 00:00:00 2001
From: Yoann Ciabaud <yoann@sonora.io>
Date: Fri, 11 Mar 2016 10:09:09 +0100
Subject: [PATCH] Adds a type param to identify the source of the peer

---
 lib/common-node.js            |  5 +++++
 lib/server/parse-http.js      |  1 +
 lib/server/parse-udp.js       |  3 ++-
 lib/server/parse-websocket.js |  3 ++-
 lib/server/swarm.js           | 10 +++++++---
 5 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/lib/common-node.js b/lib/common-node.js
index 7573367..b017bc0 100644
--- a/lib/common-node.js
+++ b/lib/common-node.js
@@ -24,6 +24,11 @@ exports.EVENT_NAMES = {
   started: 'start',
   stopped: 'stop'
 }
+exports.PEER_TYPES = {
+  http: 'http',
+  udp: 'udp',
+  websocket: 'ws'
+}
 
 function toUInt32 (n) {
   var buf = new Buffer(4)
diff --git a/lib/server/parse-http.js b/lib/server/parse-http.js
index 9f21005..4681e43 100644
--- a/lib/server/parse-http.js
+++ b/lib/server/parse-http.js
@@ -9,6 +9,7 @@ function parseHttpRequest (req, opts) {
 
   if (opts.action === 'announce' || s[0] === '/announce') {
     params.action = common.ACTIONS.ANNOUNCE
+    params.type = common.PEER_TYPES.http
 
     if (typeof params.info_hash !== 'string' || params.info_hash.length !== 20) {
       throw new Error('invalid info_hash')
diff --git a/lib/server/parse-udp.js b/lib/server/parse-udp.js
index 0eda11d..6524ce2 100644
--- a/lib/server/parse-udp.js
+++ b/lib/server/parse-udp.js
@@ -10,7 +10,8 @@ function parseUdpRequest (msg, rinfo) {
   var params = {
     connectionId: msg.slice(0, 8), // 64-bit
     action: msg.readUInt32BE(8),
-    transactionId: msg.readUInt32BE(12)
+    transactionId: msg.readUInt32BE(12),
+    type: common.PEER_TYPES.udp
   }
 
   if (!bufferEqual(params.connectionId, common.CONNECTION_ID)) {
diff --git a/lib/server/parse-websocket.js b/lib/server/parse-websocket.js
index 8fb4819..a28f5aa 100644
--- a/lib/server/parse-websocket.js
+++ b/lib/server/parse-websocket.js
@@ -7,6 +7,7 @@ function parseWebSocketRequest (socket, opts, params) {
   params = JSON.parse(params) // may throw
 
   params.action = common.ACTIONS.ANNOUNCE
+  params.type = common.PEER_TYPES.websocket
   params.socket = socket
 
   if (typeof params.info_hash !== 'string' || params.info_hash.length !== 20) {
@@ -35,7 +36,7 @@ function parseWebSocketRequest (socket, opts, params) {
 
   params.ip = opts.trustProxy
       ? socket.upgradeReq.headers['x-forwarded-for'] || socket.upgradeReq.connection.remoteAddress
-      : (socket.upgradeReq.connection.remoteAddress && socket.upgradeReq.connection.remoteAddress.replace(common.REMOVE_IPV4_MAPPED_IPV6_RE, '')) // force ipv4
+      : socket.upgradeReq.connection.remoteAddress.replace(common.REMOVE_IPV4_MAPPED_IPV6_RE, '') // force ipv4
   params.port = socket.upgradeReq.connection.remotePort
   if (params.port) {
     params.addr = (common.IPV6_RE.test(params.ip) ? '[' + params.ip + ']' : params.ip) + ':' + params.port
diff --git a/lib/server/swarm.js b/lib/server/swarm.js
index 8d2765c..d18b5b8 100644
--- a/lib/server/swarm.js
+++ b/lib/server/swarm.js
@@ -2,6 +2,7 @@ module.exports = Swarm
 
 var debug = require('debug')('bittorrent-tracker')
 var randomIterate = require('random-iterate')
+var common = require('../common')
 
 // Regard this as the default implementation of an interface that you
 // need to support when overriding Server.createSwarm() and Server.getSwarm()
@@ -13,7 +14,8 @@ function Swarm (infoHash, server) {
 
 Swarm.prototype.announce = function (params, cb) {
   var self = this
-  var peer = self.peers[params.addr || params.peer_id]
+  var id = params.type === common.PEER_TYPES.websocket ? params.peer_id : params.addr
+  var peer = self.peers[id]
 
   if (params.event === 'started') {
     self._onAnnounceStarted(params, peer)
@@ -49,7 +51,8 @@ Swarm.prototype._onAnnounceStarted = function (params, peer) {
 
   if (params.left === 0) this.complete += 1
   else this.incomplete += 1
-  peer = this.peers[params.addr || params.peer_id] = {
+  var id = params.type === common.PEER_TYPES.websocket ? params.peer_id : params.addr
+  peer = this.peers[id] = {
     complete: params.left === 0,
     peerId: params.peer_id, // as hex
     ip: params.ip, // only http, udp
@@ -66,7 +69,8 @@ Swarm.prototype._onAnnounceStopped = function (params, peer) {
 
   if (peer.complete) this.complete -= 1
   else this.incomplete -= 1
-  this.peers[params.addr || params.peer_id] = null
+  var id = params.type === common.PEER_TYPES.websocket ? params.peer_id : params.addr
+  this.peers[id] = null
 }
 
 Swarm.prototype._onAnnounceCompleted = function (params, peer) {