refactor with async/await

Trying to be more robust about error handling and failure.

server.js

@@ -1,14 +1,19 @@
-var log = require('bookrc');
-var express = require('express');
-var tldjs = require('tldjs');
-var on_finished = require('on-finished');
-var debug = require('debug')('localtunnel-server');
-var http_proxy = require('http-proxy');
-var http = require('http');
+import log from 'bookrc';
+import express from 'express';
+import tldjs from 'tldjs';
+import on_finished from 'on-finished';
+import Debug from 'debug';
+import http_proxy from 'http-proxy';
+import http from 'http';
+import Promise from 'bluebird';
 
-var BindingAgent = require('./lib/BindingAgent');
+import Proxy from './proxy';
+import rand_id from './lib/rand_id';
+import BindingAgent from './lib/BindingAgent';
 
-var proxy = http_proxy.createProxyServer({
+const debug = Debug('localtunnel:server');
+
+const proxy = http_proxy.createProxyServer({
     target: 'http://localtunnel.github.io'
 });
 
@@ -23,51 +28,54 @@ proxy.on('proxyReq', function(proxyReq, req, res, options) {
     proxyReq.setHeader('host', 'localtunnel.github.io');
 });
 
-var Proxy = require('./proxy');
-var rand_id = require('./lib/rand_id');
-
-var PRODUCTION = process.env.NODE_ENV === 'production';
+const PRODUCTION = process.env.NODE_ENV === 'production';
 
 // id -> client http server
-var clients = Object.create(null);
+const clients = Object.create(null);
 
 // proxy statistics
-var stats = {
+const stats = {
     tunnels: 0
 };
 
+// handle proxying a request to a client
+// will wait for a tunnel socket to become available
 function maybe_bounce(req, res, sock, head) {
     // without a hostname, we won't know who the request is for
-    var hostname = req.headers.host;
+    const hostname = req.headers.host;
     if (!hostname) {
         return false;
     }
 
-    var subdomain = tldjs.getSubdomain(hostname);
+    const subdomain = tldjs.getSubdomain(hostname);
     if (!subdomain) {
         return false;
     }
 
-    var client_id = subdomain;
-    var client = clients[client_id];
+    const client = clients[subdomain];
 
     // no such subdomain
     // we use 502 error to the client to signify we can't service the request
     if (!client) {
-        res.statusCode = 502;
-        res.end('localtunnel error: no active client for \'' + client_id + '\'');
-        req.connection.destroy();
+        if (res) {
+            res.statusCode = 502;
+            res.end(`no active client for '${subdomain}'`);
+            req.connection.destroy();
+        }
+        else if (sock) {
+            sock.destroy();
+        }
+
         return true;
     }
 
-    var finished = false;
+    let finished = false;
     if (sock) {
         sock.once('end', function() {
             finished = true;
         });
     }
-
-    if (res) {
+    else if (res) {
         // flag if we already finished before we get a socket
         // we can't respond to these requests
         on_finished(res, function(err) {
@@ -75,16 +83,19 @@ function maybe_bounce(req, res, sock, head) {
             req.connection.destroy();
         });
     }
+    // not something we are expecting, need a sock or a res
+    else {
+        req.connection.destroy();
+        return true;
+    }
 
     // TODO add a timeout, if we run out of sockets, then just 502
 
     // get client port
-    client.next_socket(function(socket, done) {
-        done = done || function() {};
-
+    client.next_socket(async (socket) => {
         // the request already finished or client disconnected
         if (finished) {
-            return done();
+            return;
         }
 
         // happens when client upstream is disconnected
@@ -103,9 +114,9 @@ function maybe_bounce(req, res, sock, head) {
         // and directly pipe the socket data
         // avoids having to rebuild the request and handle upgrades via the http client
         if (res === null) {
-            var arr = [req.method + ' ' + req.url + ' HTTP/' + req.httpVersion];
-            for (var i=0 ; i < (req.rawHeaders.length-1) ; i+=2) {
-                arr.push(req.rawHeaders[i] + ': ' + req.rawHeaders[i+1]);
+            const arr = [`${req.method} ${req.url} HTTP/${req.httpVersion}`];
+            for (let i=0 ; i < (req.rawHeaders.length-1) ; i+=2) {
+                arr.push(`${req.rawHeaders[i]}: ${req.rawHeaders[i+1]}`);
             }
 
             arr.push('');
@@ -113,40 +124,55 @@ function maybe_bounce(req, res, sock, head) {
 
             socket.pipe(sock).pipe(socket);
             socket.write(arr.join('\r\n'));
-            socket.once('end', function() {
-                done();
+
+            await new Promise((resolve) => {
+                socket.once('end', resolve);
             });
 
             return;
         }
 
-        var agent = new BindingAgent({
+        // regular http request
+
+        const agent = new BindingAgent({
             socket: socket
         });
 
-        var opt = {
+        const opt = {
             path: req.url,
             agent: agent,
             method: req.method,
             headers: req.headers
         };
 
-        var client_req = http.request(opt, function(client_res) {
-            // write response code and headers
-            res.writeHead(client_res.statusCode, client_res.headers);
-            
-            client_res.pipe(res);
-            on_finished(client_res, function(err) {
-                done();
-            });
-        });
+        await new Promise((resolve) => {
+            // what if error making this request?
+            const client_req = http.request(opt, function(client_res) {
+                // write response code and headers
+                res.writeHead(client_res.statusCode, client_res.headers);
 
-        req.pipe(client_req);
+                client_res.pipe(res);
+                on_finished(client_res, function(err) {
+                    resolve();
+                });
+            });
+
+            // happens if the other end dies while we are making the request
+            // so we just end the req and move on
+            // we can't really do more with the response here because headers
+            // may already be sent
+            client_req.on('error', (err) => {
+                req.connection.destroy();
+            });
+
+            req.pipe(client_req);
+        });
     });
 
     return true;
 }
 
+// create a new tunnel with `id`
 function new_client(id, opt, cb) {
 
     // can't ask for id already is use
@@ -155,43 +181,49 @@ function new_client(id, opt, cb) {
         id = rand_id();
     }
 
-    var popt = {
+    const popt = {
         id: id,
         max_tcp_sockets: opt.max_tcp_sockets
     };
 
-    var client = Proxy(popt, function(err, info) {
-        if (err) {
-            return cb(err);
-        }
+    const client = Proxy(popt);
 
-        ++stats.tunnels;
-        clients[id] = client;
-
-        info.id = id;
-
-        cb(err, info);
-    });
+    // add to clients map immediately
+    // avoiding races with other clients requesting same id
+    clients[id] = client;
 
     client.on('end', function() {
         --stats.tunnels;
         delete clients[id];
     });
+
+    client.start((err, info) => {
+        if (err) {
+            delete clients[id];
+            cb(err);
+            return;
+        }
+
+        ++stats.tunnels;
+
+        info.id = id;
+        cb(err, info);
+    });
 }
 
 module.exports = function(opt) {
     opt = opt || {};
 
-    var schema = opt.secure ? 'https' : 'http';
+    const schema = opt.secure ? 'https' : 'http';
 
-    var app = express();
+    const app = express();
 
     app.get('/', function(req, res, next) {
         if (req.query['new'] === undefined) {
             return next();
         }
 
-        var req_id = rand_id();
+        const req_id = rand_id();
         debug('making new client with id %s', req_id);
         new_client(req_id, opt, function(err, info) {
             if (err) {
@@ -199,7 +231,7 @@ module.exports = function(opt) {
                 return res.end(err.message);
             }
 
-            var url = schema + '://' + req_id + '.' + req.headers.host;
+            const url = schema + '://' + req_id + '.' + req.headers.host;
             info.url = url;
             res.json(info);
         });
@@ -217,12 +249,18 @@ module.exports = function(opt) {
         proxy.web(req, res);
     });
 
+    app.get('/api/status', function(req, res, next) {
+        res.json({
+            tunnels: stats.tunnels,
+        });
+    });
+
     app.get('/:req_id', function(req, res, next) {
-        var req_id = req.params.req_id;
+        const req_id = req.params.req_id;
 
         // limit requested hostnames to 63 characters
         if (! /^[a-z0-9]{4,63}$/.test(req_id)) {
-            var err = new Error('Invalid subdomain. Subdomains must be lowercase and between 4 and 63 alphanumeric characters.');
+            const err = new Error('Invalid subdomain. Subdomains must be lowercase and between 4 and 63 alphanumeric characters.');
             err.statusCode = 403;
             return next(err);
         }
@@ -233,7 +271,7 @@ module.exports = function(opt) {
                 return next(err);
             }
 
-            var url = schema + '://' + req_id + '.' + req.headers.host;
+            const url = schema + '://' + req_id + '.' + req.headers.host;
             info.url = url;
             res.json(info);
         });
@@ -241,13 +279,13 @@ module.exports = function(opt) {
     });
 
     app.use(function(err, req, res, next) {
-        var status = err.statusCode || err.status || 500;
+        const status = err.statusCode || err.status || 500;
         res.status(status).json({
             message: err.message
         });
     });
 
-    var server = http.createServer();
+    const server = http.createServer();
 
     server.on('request', function(req, res) {
         debug('request %s', req.url);