Check origin header for websocket connection

lib/Server.js

@@ -24,6 +24,22 @@ const webpackDevMiddleware = require('webpack-dev-middleware');
 const OptionsValidationError = require('./OptionsValidationError');
 const optionsSchema = require('./optionsSchema.json');
 
+// Workaround for sockjs@~0.3.19
+// sockjs will remove Origin header, however Origin header is required for checking host.
+// See https://github.com/webpack/webpack-dev-server/issues/1604 for more information
+{
+  // eslint-disable-next-line global-require
+  const SockjsSession = require('sockjs/lib/transport').Session;
+  const decorateConnection = SockjsSession.prototype.decorateConnection;
+  SockjsSession.prototype.decorateConnection = function(req) {
+    decorateConnection.call(this, req);
+    const connection = this.connection;
+    if (connection.headers && !('origin' in connection.headers) && 'origin' in req.headers) {
+      connection.headers.origin = req.headers.origin;
+    }
+  };
+}
+
 const clientStats = { errorDetails: false };
 const log = console.log; // eslint-disable-line no-console
 
@@ -510,17 +526,23 @@ Server.prototype.setContentHeaders = function (req, res, next) {
   next();
 };
 
-Server.prototype.checkHost = function (headers) {
+Server.prototype.checkHost = function (headers, headerToCheck) {
   // allow user to opt-out this security check, at own risk
   if (this.disableHostCheck) return true;
 
+  if (!headerToCheck) headerToCheck = 'host';
   // get the Host header and extract hostname
   // we don't care about port not matching
-  const hostHeader = headers.host;
+  const hostHeader = headers[headerToCheck];
   if (!hostHeader) return false;
 
   // use the node url-parser to retrieve the hostname from the host-header.
-  const hostname = url.parse(`//${hostHeader}`, false, true).hostname;
+  const hostname = url.parse(
+    // if hostHeader doesn't have scheme, add // for parsing.
+    /^(.+:)?\/\//.test(hostHeader) ? hostHeader : `//${hostHeader}`,
+    false,
+    true,
+  ).hostname;
 
   // always allow requests with explicit IPv4 or IPv6-address.
   // A note on IPv6 addresses: hostHeader will always contain the brackets denoting
@@ -581,8 +603,8 @@ Server.prototype.listen = function (port, hostname, fn) {
 
     sockServer.on('connection', (conn) => {
       if (!conn) return;
-      if (!this.checkHost(conn.headers)) {
-        this.sockWrite([conn], 'error', 'Invalid Host header');
+      if (!this.checkHost(conn.headers) || !this.checkHost(conn.headers, 'origin')) {
+        this.sockWrite([conn], 'error', 'Invalid Host/Origin header');
         conn.close();
         return;
       }

test/Validation.test.js

@@ -144,6 +144,19 @@ describe('Validation', () => {
       }
     });
 
+    it('should allow urls with scheme for checking origin', () => {
+      const options = {
+        public: 'test.host:80'
+      };
+      const headers = {
+        origin: 'https://test.host'
+      };
+      const server = new Server(compiler, options);
+      if (!server.checkHost(headers, 'origin')) {
+        throw new Error("Validation didn't fail");
+      }
+    });
+
     describe('allowedHosts', () => {
       it('should allow hosts in allowedHosts', () => {
         const testHosts = [