From 60f03b9529e0f02bc74a3c469613a72c264135c7 Mon Sep 17 00:00:00 2001
From: Yoann Ciabaud <yoann@sonora.io>
Date: Thu, 3 Mar 2016 11:10:54 +0100
Subject: [PATCH] Provide IP and HTTP headers in both HTTP and Websocket server

---
 lib/common-node.js            |  1 +
 lib/server/parse-http.js      |  6 +++---
 lib/server/parse-websocket.js | 13 ++++++++++++-
 server.js                     | 10 ++++++----
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/lib/common-node.js b/lib/common-node.js
index 12c45b2..7573367 100644
--- a/lib/common-node.js
+++ b/lib/common-node.js
@@ -7,6 +7,7 @@ var querystring = require('querystring')
 
 exports.IPV4_RE = /^[\d\.]+$/
 exports.IPV6_RE = /^[\da-fA-F:]+$/
+exports.REMOVE_IPV4_MAPPED_IPV6_RE = /^::ffff:/
 
 exports.CONNECTION_ID = Buffer.concat([ toUInt32(0x417), toUInt32(0x27101980) ])
 exports.ACTIONS = { CONNECT: 0, ANNOUNCE: 1, SCRAPE: 2, ERROR: 3 }
diff --git a/lib/server/parse-http.js b/lib/server/parse-http.js
index cb3e481..9f21005 100644
--- a/lib/server/parse-http.js
+++ b/lib/server/parse-http.js
@@ -2,8 +2,6 @@ module.exports = parseHttpRequest
 
 var common = require('../common')
 
-var REMOVE_IPV4_MAPPED_IPV6_RE = /^::ffff:/
-
 function parseHttpRequest (req, opts) {
   if (!opts) opts = {}
   var s = req.url.split('?')
@@ -34,8 +32,10 @@ function parseHttpRequest (req, opts) {
 
     params.ip = opts.trustProxy
       ? req.headers['x-forwarded-for'] || req.connection.remoteAddress
-      : req.connection.remoteAddress.replace(REMOVE_IPV4_MAPPED_IPV6_RE, '') // force ipv4
+      : req.connection.remoteAddress.replace(common.REMOVE_IPV4_MAPPED_IPV6_RE, '') // force ipv4
     params.addr = (common.IPV6_RE.test(params.ip) ? '[' + params.ip + ']' : params.ip) + ':' + params.port
+
+    params.headers = req.headers
   } else if (opts.action === 'scrape' || s[0] === '/scrape') {
     params.action = common.ACTIONS.SCRAPE
 
diff --git a/lib/server/parse-websocket.js b/lib/server/parse-websocket.js
index 944843a..8fb4819 100644
--- a/lib/server/parse-websocket.js
+++ b/lib/server/parse-websocket.js
@@ -2,7 +2,8 @@ module.exports = parseWebSocketRequest
 
 var common = require('../common')
 
-function parseWebSocketRequest (socket, params) {
+function parseWebSocketRequest (socket, opts, params) {
+  if (!opts) opts = {}
   params = JSON.parse(params) // may throw
 
   params.action = common.ACTIONS.ANNOUNCE
@@ -32,5 +33,15 @@ function parseWebSocketRequest (socket, params) {
   )
   params.compact = -1 // return full peer objects (used for websocket responses)
 
+  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
+  params.port = socket.upgradeReq.connection.remotePort
+  if (params.port) {
+    params.addr = (common.IPV6_RE.test(params.ip) ? '[' + params.ip + ']' : params.ip) + ':' + params.port
+  }
+
+  params.headers = socket.upgradeReq.headers
+
   return params
 }
diff --git a/server.js b/server.js
index a0e107c..e60a581 100644
--- a/server.js
+++ b/server.js
@@ -294,21 +294,23 @@ Server.prototype.onUdpRequest = function (msg, rinfo) {
   })
 }
 
-Server.prototype.onWebSocketConnection = function (socket) {
+Server.prototype.onWebSocketConnection = function (socket, opts) {
   var self = this
+  if (!opts) opts = {}
+  opts.trustProxy = opts.trustProxy || self._trustProxy
   socket.peerId = null // as hex
   socket.infoHashes = []
   socket.onSend = self._onWebSocketSend.bind(self, socket)
-  socket.on('message', self._onWebSocketRequest.bind(self, socket))
+  socket.on('message', self._onWebSocketRequest.bind(self, socket, opts))
   socket.on('error', self._onWebSocketError.bind(self, socket))
   socket.on('close', self._onWebSocketClose.bind(self, socket))
 }
 
-Server.prototype._onWebSocketRequest = function (socket, params) {
+Server.prototype._onWebSocketRequest = function (socket, opts, params) {
   var self = this
 
   try {
-    params = parseWebSocketRequest(socket, params)
+    params = parseWebSocketRequest(socket, opts, params)
   } catch (err) {
     socket.send(JSON.stringify({
       'failure reason': err.message