From 6a103f3e61e5c32177c32c33f02af01dc864b239 Mon Sep 17 00:00:00 2001
From: Maxim Rubis <siburny@gmail.com>
Date: Wed, 6 Jun 2018 08:05:34 -0400
Subject: [PATCH] Add client token authentication

---
 bin/server           |  3 +++
 lib/ClientManager.js | 13 ++++++++-----
 server.js            | 20 ++++++++++++++------
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/bin/server b/bin/server
index a6bba9d..c4cf740 100755
--- a/bin/server
+++ b/bin/server
@@ -42,6 +42,9 @@ const server = CreateServer({
     max_tcp_sockets: argv['max-sockets'],
     secure: argv.secure,
     domain: argv.domain,
+    tokens: {
+        // 'clientid': 'token'
+    }
 });
 
 server.listen(argv.port, argv.address, () => {
diff --git a/lib/ClientManager.js b/lib/ClientManager.js
index e1a7838..65686f1 100644
--- a/lib/ClientManager.js
+++ b/lib/ClientManager.js
@@ -26,15 +26,18 @@ class ClientManager {
     }
 
     // create a new tunnel with `id`
-    // if the id is already used, a random id is assigned
-    // if the tunnel could not be created, throws an error
-    async newClient(id) {
+    async newClient (id, token) {
         const clients = this.clients;
         const stats = this.stats;
 
-        // can't ask for id already is use
-        if (clients[id]) {
+        if (clients[id] !== undefined) {
+            if(this.opt.tokens[id] !== undefined) {
+                return Promise.reject(new Error('Client is already connected. Try again later.'));
+            }
+
             id = hri.random();
+        } else if(this.opt.tokens[id] !== undefined && this.opt.tokens[id] !== token) {
+            return Promise.reject(new Error('Token is not found or correct.'));
         }
 
         const maxSockets = this.opt.max_tcp_sockets;
diff --git a/server.js b/server.js
index 5a53c47..b5ef92e 100644
--- a/server.js
+++ b/server.js
@@ -104,13 +104,21 @@ export default function(opt) {
             return;
         }
 
-        debug('making new client with id %s', reqId);
-        const info = await manager.newClient(reqId);
+        try {
+            debug('making new client with id %s', reqId);
+            const info = await manager.newClient(reqId);
 
-        const url = schema + '://' + info.id + '.' + ctx.request.host;
-        info.url = url;
-        ctx.body = info;
-        return;
+            const url = schema + '://' + info.id + '.' + ctx.request.host;
+            info.url = url;
+            ctx.body = info;
+        }
+        catch(err) {
+            debug('Error: %s with id %s', err.message, reqId);
+            ctx.status = 403;
+            ctx.body = {
+                message: err.message,
+            };
+        }
     });
 
     const server = http.createServer();