From ccdf6046af2b2697ed5b2c21b32efd0111e32d72 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Sat, 27 Aug 2016 17:12:53 -0700
Subject: [PATCH] browser: rewrite wsproxy.

---
 browser/server.js  |  10 +-
 browser/wsproxy.js | 316 ++++++++++++++++++++++++++-------------------
 2 files changed, 193 insertions(+), 133 deletions(-)

diff --git a/browser/server.js b/browser/server.js
index 8d1ebfe1..dc8d0817 100644
--- a/browser/server.js
+++ b/browser/server.js
@@ -1,11 +1,17 @@
 'use strict';
 
 var HTTPBase = require('../lib/http/base');
+var WSProxy = require('./wsproxy');
 var fs = require('fs');
 
 var server = new HTTPBase();
-var proxy = require('./wsproxy')({
-  pow: process.argv.indexOf('--pow') !== -1
+var proxy = new WSProxy({
+  pow: process.argv.indexOf('--pow') !== -1,
+  ports: [8333, 18333, 18444, 28333, 28901]
+});
+
+proxy.on('error', function(err) {
+  console.error(err.stack + '');
 });
 
 var index = fs.readFileSync(__dirname + '/index.html');
diff --git a/browser/wsproxy.js b/browser/wsproxy.js
index f4cd3a01..6700a477 100644
--- a/browser/wsproxy.js
+++ b/browser/wsproxy.js
@@ -1,150 +1,204 @@
 'use strict';
 
-var http = require('http');
 var net = require('net');
 var IOServer = require('socket.io');
 var utils = require('../lib/utils/utils');
 var IP = require('../lib/utils/ip');
-var networks = require('../lib/protocol/networks');
 var BufferWriter = require('../lib/utils/writer');
-var ports = [];
-var i, type;
+var EventEmitter = require('events').EventEmitter;
 
-for (i = 0; i < networks.types.length; i++) {
-  type = networks.types[i];
-  ports.push(networks[type].port);
-}
+var TARGET = new Buffer(
+  '0000ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff',
+  'hex');
 
-module.exports = function wsproxy(options) {
-  var target, io;
+function WSProxy(options) {
+  if (!(this instanceof WSProxy))
+    return new WSProxy(options);
+
+  EventEmitter.call(this);
 
   if (!options)
     options = {};
 
-  target = new Buffer(
-    '0000ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff',
-    'hex');
+  this.options = options;
+  this.target = options.target || TARGET;
+  this.pow = options.pow === true;
+  this.ports = options.ports || [];
+  this.io = new IOServer();
+  this.sockets = new WeakMap();
 
-  io = new IOServer();
+  this._init();
+}
 
-  io.on('error', function(err) {
-    utils.error(err.stack + '');
+utils.inherits(WSProxy, EventEmitter);
+
+WSProxy.prototype._init = function _init() {
+  var self = this;
+
+  this.io.on('error', function(err) {
+    self.emit('error', err);
   });
 
-  io.on('connection', function(ws) {
-    var snonce = utils.nonce().toArrayLike(Buffer, 'be', 8);
-    var socket, pow;
-
-    ws.emit('info', {
-      pow: !!options.pow,
-      snonce: snonce.toString('hex'),
-      target: target.toString('hex')
-    });
-
-    ws.on('error', function(err) {
-      utils.error(err.stack + '');
-    });
-
-    ws.on('tcp connect', function(port, host, nonce) {
-      if (socket)
-        return;
-
-      if (!utils.isNumber(port)
-          || typeof host !== 'string'
-          || host.length === 0) {
-        utils.error('Client gave bad arguments.');
-        ws.emit('tcp close');
-        ws.disconnect();
-        return;
-      }
-
-      if (options.pow) {
-        if (!utils.isNumber(nonce)) {
-          utils.error('Client did not solve proof of work.');
-          ws.emit('tcp close');
-          ws.disconnect();
-          return;
-        }
-
-        pow = new BufferWriter();
-        pow.writeU32(nonce);
-        pow.writeBytes(snonce);
-        pow.writeU32(port);
-        pow.writeString(host, 'ascii');
-        pow = pow.render();
-
-        if (utils.cmp(utils.dsha256(pow), target) > 0) {
-          utils.error('Client did not solve proof of work.');
-          ws.emit('tcp close');
-          ws.disconnect();
-          return;
-        }
-      }
-
-      if (!/^[a-zA-Z0-9\.:\-]+$/.test(host)) {
-        utils.error('Client gave a bad host.');
-        ws.emit('tcp close');
-        ws.disconnect();
-        return;
-      }
-
-      if (IP.isPrivate(host)) {
-        utils.error('Client is trying to connect to a private ip.');
-        ws.emit('tcp close');
-        ws.disconnect();
-        return;
-      }
-
-      if (ports.indexOf(port) === -1) {
-        utils.error('Client is trying to connect to a non-bitcoin port.');
-        ws.emit('tcp close');
-        ws.disconnect();
-        return;
-      }
-
-      try {
-        socket = net.connect(port, host);
-        utils.log('Connecting to %s:%d.', host, port);
-      } catch (e) {
-        utils.error(e.message);
-        utils.error('Closing %s:%d.', host, port);
-        ws.emit('tcp close');
-        ws.disconnect();
-        return;
-      }
-
-      socket.on('connect', function() {
-        ws.emit('tcp connect');
-      });
-
-      socket.on('data', function(data) {
-        ws.emit('tcp data', data.toString('hex'));
-      });
-
-      socket.on('error', function(err) {
-        ws.emit('tcp error', {
-          message: err.message,
-          code: err.code || null
-        });
-      });
-
-      socket.on('close', function() {
-        utils.log('Closing %s:%d.', host, port);
-        ws.emit('tcp close');
-        ws.disconnect();
-      });
-
-      ws.on('tcp data', function(data) {
-        if (typeof data !== 'string')
-          return;
-        socket.write(new Buffer(data, 'hex'));
-      });
-
-      ws.on('disconnect', function() {
-        socket.destroy();
-      });
-    });
+  this.io.on('connection', function(ws) {
+    self._handleSocket(ws);
   });
-
-  return io;
 };
+
+WSProxy.prototype._handleSocket = function _handleSocket(ws) {
+  var self = this;
+  var state = new SocketState(this, ws);
+
+  // Use a weak map to avoid
+  // mutating the websocket object.
+  this.sockets.set(ws, state);
+
+  ws.emit('info', state.toInfo());
+
+  ws.on('error', function(err) {
+    self.emit('error', err);
+  });
+
+  ws.on('tcp connect', function(port, host, nonce) {
+    self._handleConnect(ws, port, host, nonce);
+  });
+};
+
+WSProxy.prototype._handleConnect = function _handleConnect(ws, port, host, nonce) {
+  var self = this;
+  var state = this.sockets.get(ws);
+  var socket, pow;
+
+  if (state.socket) {
+    this.log('Client is trying to reconnect (%s).', state.host);
+    return;
+  }
+
+  if (!utils.isNumber(port)
+      || typeof host !== 'string'
+      || host.length === 0) {
+    this.log('Client gave bad arguments (%s).', state.host);
+    ws.emit('tcp close');
+    ws.disconnect();
+    return;
+  }
+
+  if (this.pow) {
+    if (!utils.isNumber(nonce)) {
+      this.log('Client did not solve proof of work.', state.host);
+      ws.emit('tcp close');
+      ws.disconnect();
+      return;
+    }
+
+    pow = new BufferWriter();
+    pow.writeU32(nonce);
+    pow.writeBytes(state.snonce);
+    pow.writeU32(port);
+    pow.writeString(host, 'ascii');
+    pow = pow.render();
+
+    if (utils.cmp(utils.hash256(pow), this.target) > 0) {
+      this.log('Client did not solve proof of work (%s).', state.host);
+      ws.emit('tcp close');
+      ws.disconnect();
+      return;
+    }
+  }
+
+  if (!/^[a-zA-Z0-9\.:\-]+$/.test(host)) {
+    this.log('Client gave a bad host (%s).', state.host);
+    ws.emit('tcp close');
+    ws.disconnect();
+    return;
+  }
+
+  if (IP.isPrivate(host)) {
+    this.log('Client is trying to connect to a private ip (%s).', state.host);
+    ws.emit('tcp close');
+    ws.disconnect();
+    return;
+  }
+
+  if (this.ports.indexOf(port) === -1) {
+    this.log('Client is connecting to non-whitelist port (%s).', state.host);
+    ws.emit('tcp close');
+    ws.disconnect();
+    return;
+  }
+
+  try {
+    socket = state.connect(port, host);
+    this.log('Connecting to %s (%s).', state.remoteHost, state.host);
+  } catch (e) {
+    this.log(e.message);
+    this.log('Closing %s (%s).', state.remoteHost, state.host);
+    ws.emit('tcp close');
+    ws.disconnect();
+    return;
+  }
+
+  socket.on('connect', function() {
+    ws.emit('tcp connect');
+  });
+
+  socket.on('data', function(data) {
+    ws.emit('tcp data', data.toString('hex'));
+  });
+
+  socket.on('error', function(err) {
+    ws.emit('tcp error', {
+      message: err.message,
+      code: err.code || null
+    });
+  });
+
+  socket.on('close', function() {
+    self.log('Closing %s (%s).', state.remoteHost, state.host);
+    ws.emit('tcp close');
+    ws.disconnect();
+  });
+
+  ws.on('tcp data', function(data) {
+    if (typeof data !== 'string')
+      return;
+    socket.write(new Buffer(data, 'hex'));
+  });
+
+  ws.on('disconnect', function() {
+    socket.destroy();
+  });
+};
+
+WSProxy.prototype.log = function log() {
+  process.stdout.write('wsproxy: ');
+  console.log.apply(console, arguments);
+};
+
+WSProxy.prototype.attach = function attach(server) {
+  this.io.attach(server);
+};
+
+function SocketState(server, socket) {
+  this.pow = server.pow;
+  this.target = server.target;
+  this.snonce = utils.nonce(true);
+  this.socket = null;
+  this.host = socket.conn.remoteAddress;
+  this.remoteHost = null;
+}
+
+SocketState.prototype.toInfo = function toInfo() {
+  return {
+    pow: this.pow,
+    target: this.target.toString('hex'),
+    snonce: this.snonce.toString('hex')
+  };
+};
+
+SocketState.prototype.connect = function connect(port, host) {
+  this.socket = net.connect(port, host);
+  this.remoteHost = IP.hostname(host, port);
+  return this.socket;
+};
+
+module.exports = WSProxy;