From ee3e805bc634818655e35cb51f414eac6e495186 Mon Sep 17 00:00:00 2001
From: Joonas Lahtinen <joonas.lahtinen@iki.fi>
Date: Sat, 9 Nov 2019 14:03:36 +0200
Subject: [PATCH] fix(server): respect sockPath on transportMode: 'ws' (#2310)

---
 lib/servers/WebsocketServer.js              | 18 +++++++-
 test/server/servers/WebsocketServer.test.js | 46 +++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/lib/servers/WebsocketServer.js b/lib/servers/WebsocketServer.js
index 21bfb0c344..8a4cfb208b 100644
--- a/lib/servers/WebsocketServer.js
+++ b/lib/servers/WebsocketServer.js
@@ -3,6 +3,7 @@
 /* eslint-disable
   class-methods-use-this
 */
+const url = require('url');
 const ws = require('ws');
 const BaseServer = require('./BaseServer');
 
@@ -10,15 +11,30 @@ module.exports = class WebsocketServer extends BaseServer {
   constructor(server) {
     super(server);
     this.wsServer = new ws.Server({
-      server: this.server.listeningApp,
+      noServer: true,
       path: this.server.sockPath,
     });
 
+    this.server.listeningApp.on('upgrade', (req, sock, head) => {
+      if (!this.shouldUpgrade(req)) {
+        return;
+      }
+
+      this.wsServer.handleUpgrade(req, sock, head, (connection) => {
+        this.wsServer.emit('connection', connection, req);
+      });
+    });
+
     this.wsServer.on('error', (err) => {
       this.server.log.error(err.message);
     });
   }
 
+  shouldUpgrade(req) {
+    const pathname = url.parse(req.url).pathname;
+    return pathname.match(`${this.server.sockPath}/?`);
+  }
+
   send(connection, message) {
     // prevent cases where the server is trying to send data while connection is closing
     if (connection.readyState !== 1) {
diff --git a/test/server/servers/WebsocketServer.test.js b/test/server/servers/WebsocketServer.test.js
index e8e59a786b..b3a076abd1 100644
--- a/test/server/servers/WebsocketServer.test.js
+++ b/test/server/servers/WebsocketServer.test.js
@@ -85,6 +85,52 @@ describe('WebsocketServer', () => {
         done();
       }, 3000);
     });
+
+    it('should match sockPath', (done) => {
+      let receivedConnection = false;
+      socketServer.onConnection(() => {
+        receivedConnection = true;
+      });
+
+      // eslint-disable-next-line new-cap
+      const client = new ws(`http://localhost:${port}/ws-server`);
+
+      setTimeout(() => {
+        // the client closes itself, the server does not close it
+        client.close();
+      }, 1000);
+
+      setTimeout(() => {
+        expect(receivedConnection).toBeTruthy();
+        done();
+      }, 3000);
+    });
+
+    it('should ignore other paths', (done) => {
+      let receivedConnection = false;
+      socketServer.onConnection(() => {
+        receivedConnection = true;
+      });
+
+      // eslint-disable-next-line new-cap
+      const client = new ws(`http://localhost:${port}/ws-server-NOT`);
+
+      let receivedError = false;
+      client.on('error', () => {
+        receivedError = true;
+      });
+
+      setTimeout(() => {
+        // the client closes itself, the server does not close it
+        client.close();
+      }, 1000);
+
+      setTimeout(() => {
+        expect(receivedError).toBeTruthy();
+        expect(receivedConnection).toBeFalsy();
+        done();
+      }, 3000);
+    });
   });
 
   afterAll((done) => {