Handle websocket authentication slightly differently to make errors easier to work with

Author: DaneEveritt

resources/scripts/components/server/WebsocketHandler.tsx

@@ -11,7 +11,7 @@ export default () => {
 
     const updateToken = (uuid: string, socket: Websocket) => {
         getWebsocketToken(uuid)
-            .then(data => socket.setToken(data.token))
+            .then(data => socket.setToken(data.token, true))
             .catch(error => console.error(error));
     };
 
@@ -24,7 +24,7 @@ export default () => {
 
         const socket = new Websocket();
 
-        socket.on('SOCKET_OPEN', () => setConnectionState(true));
+        socket.on('auth success', () => setConnectionState(true));
         socket.on('SOCKET_CLOSE', () => setConnectionState(false));
         socket.on('SOCKET_ERROR', () => setConnectionState(false));
         socket.on('status', (status) => setServerStatus(status));
@@ -38,7 +38,10 @@ export default () => {
 
         getWebsocketToken(server.uuid)
             .then(data => {
+                // Connect and then set the authentication token.
                 socket.setToken(data.token).connect(data.socket);
+
+                // Once that is done, set the instance.
                 setInstance(socket);
             })
             .catch(error => console.error(error));

resources/scripts/plugins/Websocket.ts

@@ -23,9 +23,9 @@ export class Websocket extends EventEmitter {
     private token: string = '';
 
     // Connects to the websocket instance and sets the token for the initial request.
-    connect (url: string) {
+    connect (url: string): this {
         this.url = url;
-        this.socket = new Sockette(`${this.url}?token=${this.token}`, {
+        this.socket = new Sockette(`${this.url}`, {
             onmessage: e => {
                 try {
                     let { event, args } = JSON.parse(e.data);
@@ -34,11 +34,19 @@ export class Websocket extends EventEmitter {
                     console.warn('Failed to parse incoming websocket message.', ex);
                 }
             },
-            onopen: () => this.emit('SOCKET_OPEN'),
-            onreconnect: () => this.emit('SOCKET_RECONNECT'),
+            onopen: () => {
+                this.emit('SOCKET_OPEN');
+                this.authenticate();
+            },
+            onreconnect: () => {
+                this.emit('SOCKET_RECONNECT');
+                this.authenticate();
+            },
             onclose: () => this.emit('SOCKET_CLOSE'),
             onerror: () => this.emit('SOCKET_ERROR'),
         });
+
+        return this;
     }
 
     // Returns the URL connected to for the socket.
@@ -48,11 +56,11 @@ export class Websocket extends EventEmitter {
 
     // Sets the authentication token to use when sending commands back and forth
     // between the websocket instance.
-    setToken (token: string): this {
+    setToken (token: string, isUpdate = false): this {
         this.token = token;
 
-        if (this.url) {
-            this.send('auth', token);
+        if (isUpdate) {
+            this.authenticate();
         }
 
         return this;
@@ -63,6 +71,12 @@ export class Websocket extends EventEmitter {
         return this.token;
     }
 
+    authenticate () {
+        if (this.url && this.token) {
+            this.send('auth', this.token);
+        }
+    }
+
     close (code?: number, reason?: string) {
         this.url = null;
         this.token = '';