Check origin header for websocket connection

Author: timfjord

lib/Server.js

@@ -510,13 +510,14 @@ 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.
@@ -581,8 +582,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;
       }