fix(server): fix header check for socket server (#2077)

* fix(server): fix header check for socket server

* test(server): tests to check header error and server methods

Author: knagaitsev

lib/Server.js

@@ -677,25 +677,22 @@ class Server {
     const SocketServerImplementation = this.socketServerImplementation;
     this.socketServer = new SocketServerImplementation(this);
 
-    this.socketServer.onConnection((connection) => {
+    this.socketServer.onConnection((connection, headers) => {
       if (!connection) {
         return;
       }
 
-      if (
-        !this.checkHost(connection.headers) ||
-        !this.checkOrigin(connection.headers)
-      ) {
+      if (headers && (!this.checkHost(headers) || !this.checkOrigin(headers))) {
         this.sockWrite([connection], 'error', 'Invalid Host/Origin header');
 
-        connection.close();
+        this.socketServer.close(connection);
 
         return;
       }
 
       this.sockets.push(connection);
 
-      connection.on('close', () => {
+      this.socketServer.onConnectionClose(connection, () => {
         const idx = this.sockets.indexOf(connection);
 
         if (idx >= 0) {

lib/servers/SockJSServer.js

@@ -57,8 +57,14 @@ module.exports = class SockJSServer extends BaseServer {
     connection.close();
   }
 
-  // f should return the resulting connection
+  // f should return the resulting connection and, optionally, the connection headers
   onConnection(f) {
-    this.socket.on('connection', f);
+    this.socket.on('connection', (connection) => {
+      f(connection, connection.headers);
+    });
+  }
+
+  onConnectionClose(connection, f) {
+    connection.on('close', f);
   }
 };

lib/servers/WebsocketServer.js

@@ -29,6 +29,12 @@ module.exports = class WebsocketServer extends BaseServer {
 
   // f should return the resulting connection
   onConnection(f) {
-    this.wsServer.on('connection', f);
+    this.wsServer.on('connection', (connection, req) => {
+      f(connection, req.headers);
+    });
+  }
+
+  onConnectionClose(connection, f) {
+    connection.on('close', f);
   }
 };

test/server/__snapshots__/serverMode-option.test.js.snap

@@ -0,0 +1,9 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`serverMode option with a bad host header results in an error 1`] = `
+Array [
+  "open",
+  "{\\"type\\":\\"error\\",\\"data\\":\\"Invalid Host/Origin header\\"}",
+  "close",
+]
+`;

test/server/serverMode-option.test.js

@@ -5,6 +5,7 @@
 */
 const request = require('supertest');
 const sockjs = require('sockjs');
+const SockJS = require('sockjs-client/dist/sockjs');
 const SockJSServer = require('../../lib/servers/SockJSServer');
 const config = require('../fixtures/simple-config/webpack.config');
 const testServer = require('../helpers/test-server');
@@ -77,6 +78,7 @@ describe('serverMode option', () => {
 
   describe('as a class (custom "sockjs" implementation)', () => {
     it('uses supplied server implementation', (done) => {
+      let sockPath;
       server = testServer.start(
         config,
         {
@@ -102,7 +104,7 @@ describe('serverMode option', () => {
                 prefix: this.server.sockPath,
               });
 
-              expect(server.options.sockPath).toEqual('/foo/test/bar/');
+              sockPath = server.options.sockPath;
             }
 
             send(connection, message) {
@@ -114,11 +116,20 @@ describe('serverMode option', () => {
             }
 
             onConnection(f) {
-              this.socket.on('connection', f);
+              this.socket.on('connection', (connection) => {
+                f(connection, connection.headers);
+              });
+            }
+
+            onConnectionClose(connection, f) {
+              connection.on('close', f);
             }
           },
         },
-        done
+        () => {
+          expect(sockPath).toEqual('/foo/test/bar/');
+          done();
+        }
       );
     });
   });
@@ -137,4 +148,79 @@ describe('serverMode option', () => {
       }).toThrow(/serverMode must be a string/);
     });
   });
+
+  describe('with a bad host header', () => {
+    beforeAll((done) => {
+      server = testServer.start(
+        config,
+        {
+          port,
+          serverMode: class MySockJSServer extends BaseServer {
+            constructor(serv) {
+              super(serv);
+              this.socket = sockjs.createServer({
+                // Use provided up-to-date sockjs-client
+                sockjs_url: '/__webpack_dev_server__/sockjs.bundle.js',
+                // Limit useless logs
+                log: (severity, line) => {
+                  if (severity === 'error') {
+                    this.server.log.error(line);
+                  } else {
+                    this.server.log.debug(line);
+                  }
+                },
+              });
+
+              this.socket.installHandlers(this.server.listeningApp, {
+                prefix: this.server.sockPath,
+              });
+            }
+
+            send(connection, message) {
+              connection.write(message);
+            }
+
+            close(connection) {
+              connection.close();
+            }
+
+            onConnection(f) {
+              this.socket.on('connection', (connection) => {
+                f(connection, {
+                  host: null,
+                });
+              });
+            }
+
+            onConnectionClose(connection, f) {
+              connection.on('close', f);
+            }
+          },
+        },
+        done
+      );
+    });
+
+    it('results in an error', (done) => {
+      const data = [];
+      const client = new SockJS(`http://localhost:${port}/sockjs-node`);
+
+      client.onopen = () => {
+        data.push('open');
+      };
+
+      client.onmessage = (e) => {
+        data.push(e.data);
+      };
+
+      client.onclose = () => {
+        data.push('close');
+      };
+
+      setTimeout(() => {
+        expect(data).toMatchSnapshot();
+        done();
+      }, 5000);
+    });
+  });
 });

test/server/servers/SockJSServer.test.js

@@ -35,10 +35,13 @@ describe('SockJSServer', () => {
     it('should recieve connection, send message, and close client', (done) => {
       const data = [];
 
-      socketServer.onConnection((connection) => {
+      let headers;
+      socketServer.onConnection((connection, h) => {
+        headers = h;
         data.push('open');
         socketServer.send(connection, 'hello world');
         setTimeout(() => {
+          // the server closes the connection with the client
           socketServer.close(connection);
         }, 1000);
       });
@@ -54,10 +57,33 @@ describe('SockJSServer', () => {
       };
 
       setTimeout(() => {
+        expect(headers.host).toMatchSnapshot();
         expect(data).toMatchSnapshot();
         done();
       }, 3000);
     });
+
+    it('should receive client close event', (done) => {
+      let receivedClientClose = false;
+      socketServer.onConnection((connection) => {
+        socketServer.onConnectionClose(connection, () => {
+          receivedClientClose = true;
+        });
+      });
+
+      // eslint-disable-next-line new-cap
+      const client = new SockJS(`http://localhost:${port}/sockjs-node`);
+
+      setTimeout(() => {
+        // the client closes itself, the server does not close it
+        client.close();
+      }, 1000);
+
+      setTimeout(() => {
+        expect(receivedClientClose).toBeTruthy();
+        done();
+      }, 3000);
+    });
   });
 
   afterAll((done) => {

test/server/servers/WebsocketServer.test.js

@@ -35,10 +35,13 @@ describe('WebsocketServer', () => {
     it('should recieve connection, send message, and close client', (done) => {
       const data = [];
 
-      socketServer.onConnection((connection) => {
+      let headers;
+      socketServer.onConnection((connection, h) => {
+        headers = h;
         data.push('open');
         socketServer.send(connection, 'hello world');
         setTimeout(() => {
+          // the server closes the connection with the client
           socketServer.close(connection);
         }, 1000);
       });
@@ -55,10 +58,33 @@ describe('WebsocketServer', () => {
       };
 
       setTimeout(() => {
+        expect(headers.host).toMatchSnapshot();
         expect(data).toMatchSnapshot();
         done();
       }, 3000);
     });
+
+    it('should receive client close event', (done) => {
+      let receivedClientClose = false;
+      socketServer.onConnection((connection) => {
+        socketServer.onConnectionClose(connection, () => {
+          receivedClientClose = 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(receivedClientClose).toBeTruthy();
+        done();
+      }, 3000);
+    });
   });
 
   afterAll((done) => {

test/server/servers/__snapshots__/SockJSServer.test.js.snap

@@ -1,6 +1,8 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`SockJSServer server should recieve connection, send message, and close client 1`] = `
+exports[`SockJSServer server should recieve connection, send message, and close client 1`] = `"localhost:8083"`;
+
+exports[`SockJSServer server should recieve connection, send message, and close client 2`] = `
 Array [
   "open",
   "hello world",

test/server/servers/__snapshots__/WebsocketServer.test.js.snap

@@ -1,6 +1,8 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`WebsocketServer server should recieve connection, send message, and close client 1`] = `
+exports[`WebsocketServer server should recieve connection, send message, and close client 1`] = `"localhost:8121"`;
+
+exports[`WebsocketServer server should recieve connection, send message, and close client 2`] = `
 Array [
   "open",
   "hello world",