From 5c3ac5254fe88870e1c75cc82bc1cb711a68ed67 Mon Sep 17 00:00:00 2001
From: Sam Frances <sam@samfrances.uk>
Date: Fri, 3 Nov 2017 15:12:54 +0000
Subject: [PATCH] Allow for auth plugin to be a class or non-class function,
 and resolve issue with ws library retiring upgradeReq property

---
 other/js/auth_plugin_examples.js | 23 +++++++++++++++++++++--
 other/js/package.json            |  5 ++++-
 other/js/utils.js                | 24 ++++++++++++++++++++++++
 other/js/websockify.js           | 21 +++++++++++++--------
 4 files changed, 62 insertions(+), 11 deletions(-)
 create mode 100644 other/js/utils.js

diff --git a/other/js/auth_plugin_examples.js b/other/js/auth_plugin_examples.js
index 6db3d7f..64b3ad3 100644
--- a/other/js/auth_plugin_examples.js
+++ b/other/js/auth_plugin_examples.js
@@ -50,10 +50,29 @@ exports.tokenAuth = function tokenAuth(source) {
      * a token provided as the argument to the `--auth-source` command line
      * argument.
      */
-    return function(info) {
-        let token = source;
+    return {
+        authenticate(info) {
+            const token = source;
+            return urlTokenMatch(info.req.url, token, true);
+        }
+    }
+}
+
+exports.TokenAuthClass = class TokenAuthClass {
+    /**
+     * Class-based equivalent of tokenAuth
+     */
+
+    constructor(source) {
+        this.source = source;
+    }
+
+    authenticate(info) {
+        const token = this.source;
+        console.log(token)
         return urlTokenMatch(info.req.url, token, true);
     }
+
 }
 
 exports.tokenAuthEnv = function tokenAuthEnv(source) {
diff --git a/other/js/package.json b/other/js/package.json
index cf84539..842ac75 100644
--- a/other/js/package.json
+++ b/other/js/package.json
@@ -8,7 +8,10 @@
     "type": "git",
     "url": "git://github.com/kanaka/websockify.git"
   },
-  "files": ["../../docs/LICENSE.LGPL-3","websockify.js"],
+  "files": [
+    "../../docs/LICENSE.LGPL-3",
+    "websockify.js"
+  ],
   "bin": {
     "websockify": "./websockify.js"
   },
diff --git a/other/js/utils.js b/other/js/utils.js
new file mode 100644
index 0000000..4b16094
--- /dev/null
+++ b/other/js/utils.js
@@ -0,0 +1,24 @@
+
+/**
+ * A decorator that will wrap a function or a class. In the case of a non-class
+ * function, the wrapped function will behave exactly the same as before.
+ * In the case of a class, the wrapper will allow instantiating the function]
+ * without using the |new| keyword. This is useful when you don't know
+ * ahead of time if the function you will be calling is a class or a non-class
+ * function
+ *
+ * @param  {function} function_or_class
+ */
+exports.factorify = function factorify(function_or_class) {
+    return (...args) => {
+        try {
+            return function_or_class(...args);
+        } catch (e) {
+            if (e instanceof TypeError) {
+                return new function_or_class(...args);
+            } else {
+                throw e;
+            }
+        }
+    }
+}
diff --git a/other/js/websockify.js b/other/js/websockify.js
index 21c1f67..8beda2a 100755
--- a/other/js/websockify.js
+++ b/other/js/websockify.js
@@ -20,6 +20,8 @@ var argv = require('optimist').argv,
     Buffer = require('buffer').Buffer,
     WebSocketServer = require('ws').Server,
 
+    utils = require('./utils'),
+
     webServer, wsServer,
     source_host, source_port, target_host, target_port,
     auth_plugin, websocket_server_opts,
@@ -27,9 +29,9 @@ var argv = require('optimist').argv,
 
 
 // Handle new WebSocket client
-new_client = function(client) {
+new_client = function(client, upgradeReq) {
     var clientAddr = client._socket.remoteAddress, log;
-    console.log(client.upgradeReq.url);
+    console.log(upgradeReq ? upgradeReq.url : client.upgradeReq.url);
     log = function (msg) {
         console.log(' ' + clientAddr + ': '+ msg);
     };
@@ -164,15 +166,18 @@ if (argv.cert) {
 }
 
 if (argv["auth-plugin"]) {
-    let auth_plugin_arg = argv["auth-plugin"].split(".")
-    let plugin_name = auth_plugin_arg.pop();
-    let module_path = auth_plugin_arg.join(".");
+    const auth_plugin_arg = argv["auth-plugin"].split(".")
+    const plugin_name = auth_plugin_arg.pop();
+    const module_path = auth_plugin_arg.join(".");
+
+    const auth_source = argv["auth-source"] || undefined;
+    const auth_plugin = utils.factorify(
+        require(module_path)[plugin_name]
+    )(auth_source);
 
-    let auth_plugin = require(module_path)[plugin_name];
-    let auth_source = argv["auth-source"] || undefined;
     websocket_server_opts = {
         server: webServer,
-        verifyClient: auth_plugin(auth_source)
+        verifyClient: auth_plugin.authenticate.bind(auth_plugin)
     };
 } else {
     websocket_server_opts = {server: webServer};