From 386910032e1a4944b888acc00774eead85450f69 Mon Sep 17 00:00:00 2001
From: Roman Shtylman <shtylman@gmail.com>
Date: Sat, 10 Jan 2015 20:58:00 -0800
Subject: [PATCH] fix websocket support

---
 bin/server        |  1 +
 bookrc.js         |  1 -
 devops/run.sh     |  2 +-
 package.json      | 27 ++++++++++----------
 server.js         |  4 +++
 test/websocket.js | 63 +++++++++++++++++++++++++++++++++++++++++++++++
 6 files changed, 83 insertions(+), 15 deletions(-)
 create mode 100644 test/websocket.js

diff --git a/bin/server b/bin/server
index 60c20d5..d67e511 100755
--- a/bin/server
+++ b/bin/server
@@ -1,5 +1,6 @@
 #!/usr/bin/env node
 
+require('stackup');
 var log = require('bookrc');
 var localenv = require('localenv');
 var optimist = require('optimist');
diff --git a/bookrc.js b/bookrc.js
index db833b7..2586819 100644
--- a/bookrc.js
+++ b/bookrc.js
@@ -1,6 +1,5 @@
 /// bookrc logging setup
 var log = require('book').default();
-require('stackup');
 
 log.use(require('book-git')(__dirname));
 log.use(require('book-raven')(process.env.SENTRY_DSN));
diff --git a/devops/run.sh b/devops/run.sh
index 3a671d4..5aab5e5 100644
--- a/devops/run.sh
+++ b/devops/run.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
 
-docker run --restart always --name localtunnel -d --net host defunctzombie/localtunnel-server:0.0.5 bin/server --secure --port 2000
+docker run --restart always --name localtunnel -d --net host defunctzombie/localtunnel-server:0.0.5 bin/server --secure --port 2000 --max-sockets 5
 
 docker run --restart always --name nginx -d --net host -v /home/core/nginx/nginx.conf:/etc/nginx/nginx.conf -v /home/core/nginx/sites:/etc/nginx/sites -v /home/core/nginx/ssl:/etc/nginx/ssl nginx:1.7.8
diff --git a/package.json b/package.json
index 4d925d5..90ba990 100644
--- a/package.json
+++ b/package.json
@@ -9,25 +9,26 @@
   },
   "dependencies": {
     "book": "1.3.1",
-    "optimist": "0.6.1",
-    "debug": "2.1.0",
-    "bookrc": "0.0.1",
     "book-git": "0.0.2",
     "book-raven": "1.0.1",
-    "browserkthx": "0.0.2",
-    "hbs": "2.7.0",
-    "taters": "1.0.0",
-    "express": "4.10.5",
-    "makeup": "0.0.1",
-    "enchilada": "0.11.0",
+    "bookrc": "0.0.1",
     "bouncy": "3.2.2",
+    "browserkthx": "0.0.2",
+    "debug": "2.1.0",
+    "enchilada": "0.11.0",
+    "express": "4.10.5",
+    "hbs": "2.7.0",
+    "localenv": "0.2.2",
+    "makeover": "0.0.1",
+    "makeup": "0.0.1",
+    "on-finished": "2.2.0",
+    "optimist": "0.6.1",
+    "serve-favicon": "2.1.7",
     "stackup": "0.0.5",
     "stylish": "0.5.0",
-    "makeover": "0.0.1",
+    "taters": "1.0.0",
     "tldjs": "1.5.1",
-    "on-finished": "2.2.0",
-    "localenv": "0.2.2",
-    "serve-favicon": "2.1.7"
+    "ws": "0.6.5"
   },
   "devDependencies": {
     "mocha": "2.0.1",
diff --git a/server.js b/server.js
index 4dfabb4..80f5ab8 100644
--- a/server.js
+++ b/server.js
@@ -53,6 +53,10 @@ function maybe_bounce(req, res, bounce) {
     // we can't respond to these requests
     var finished = false;
     on_finished(res, function(err) {
+        if (req.headers['upgrade'] == 'websocket') {
+            return;
+        }
+
         finished = true;
         req.connection.destroy();
     });
diff --git a/test/websocket.js b/test/websocket.js
new file mode 100644
index 0000000..5ee9522
--- /dev/null
+++ b/test/websocket.js
@@ -0,0 +1,63 @@
+var http = require('http');
+var url = require('url');
+var assert = require('assert');
+var localtunnel = require('localtunnel');
+var WebSocket = require('ws');
+var WebSocketServer = require('ws').Server;
+
+var localtunnel_server = require('../server')();
+
+var lt_server_port
+
+test('setup localtunnel server', function(done) {
+    var server = localtunnel_server.listen(function() {
+        lt_server_port = server.address().port;
+        console.log('lt server on:', lt_server_port);
+        done();
+    });
+});
+
+test('setup local websocket server', function(done) {
+
+    var wss = new WebSocketServer({ port: 0 }, function() {
+        test._fake_port = wss._server.address().port;
+        done();
+    });
+
+    wss.on('connection', function connection(ws) {
+        ws.on('message', function incoming(message) {
+            ws.send(message);
+        });
+    });
+});
+
+test('setup localtunnel client', function(done) {
+    var opt = {
+        host: 'http://localhost:' + lt_server_port,
+    };
+
+    localtunnel(test._fake_port, opt, function(err, tunnel) {
+        assert.ifError(err);
+        var url = tunnel.url;
+        assert.ok(new RegExp('^http:\/\/.*localhost:' + lt_server_port + '$').test(url));
+        test._fake_url = url;
+        done(err);
+    });
+});
+
+test('test websocket server request', function(done) {
+    var hostname = url.parse(test._fake_url).hostname;
+    var ws = new WebSocket('http://localhost:' + lt_server_port, {
+        headers: {
+            host: hostname + '.tld'
+        }
+    });
+    ws.on('message', function(msg) {
+        assert.equal(msg, 'something');
+        done();
+    });
+
+    ws.on('open', function open() {
+        ws.send('something');
+    });
+});