From 569ec1d62ec222c5a304e643f89e0ffe7f6d9857 Mon Sep 17 00:00:00 2001
From: James Coglan <jcoglan@gmail.com>
Date: Mon, 30 Oct 2017 22:53:34 +0000
Subject: [PATCH] Perform more stringent checks on request headers so that it
 should not be possible to get a driver for requests with invalid headers.

---
 lib/websocket/driver.js        | 13 +------
 lib/websocket/driver/server.js | 68 ++++++++++++++++++++++++----------
 2 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/lib/websocket/driver.js b/lib/websocket/driver.js
index 221ab27..e31ddc4 100644
--- a/lib/websocket/driver.js
+++ b/lib/websocket/driver.js
@@ -27,19 +27,8 @@ var Driver = {
     return Server.http.apply(Server, arguments);
   },
 
-  isSecureRequest: function(request) {
-    return Server.isSecureRequest(request);
-  },
-
   isWebSocket: function(request) {
-    if (request.method !== 'GET') return false;
-
-    var connection = request.headers.connection || '',
-        upgrade    = request.headers.upgrade || '';
-
-    return request.method === 'GET' &&
-           connection.toLowerCase().split(/ *, */).indexOf('upgrade') >= 0 &&
-           upgrade.toLowerCase() === 'websocket';
+    return Server.isWebSocket(request);
   },
 
   validateOptions: function(options, validKeys) {
diff --git a/lib/websocket/driver/server.js b/lib/websocket/driver/server.js
index 964ed28..73d9435 100644
--- a/lib/websocket/driver/server.js
+++ b/lib/websocket/driver/server.js
@@ -71,6 +71,54 @@ var instance = {
 for (var key in instance)
   Server.prototype[key] = instance[key];
 
+Server.http = function(request, options) {
+  options = options || {};
+  if (options.requireMasking === undefined) options.requireMasking = true;
+
+  var klass = this.getDriverClass(request),
+      url   = this.determineUrl(request);
+
+  return klass && new klass(request, url, options);
+};
+
+Server.isWebSocket = function(request) {
+  var klass = this.getDriverClass(request);
+  return !!klass;
+};
+
+Server.getDriverClass = function(request) {
+  var headers = request.headers;
+
+  var connection = headers['connection'] || '',
+      key        = headers['sec-websocket-key'],
+      key1       = headers['sec-websocket-key1'],
+      key2       = headers['sec-websocket-key2'],
+      origin     = headers['origin'],
+      upgrade    = headers['upgrade'] || '',
+      version    = headers['sec-websocket-version'];
+
+  if (request.method !== 'GET' ||
+      connection.toLowerCase().split(/ *, */).indexOf('upgrade') < 0 ||
+      upgrade.toLowerCase() !== 'websocket')
+    return null;
+
+  if (typeof version === 'string' || typeof key === 'string')
+    return (version === '13' && key.length > 0) ? Hybi : null;
+
+  if (typeof origin !== 'string' || origin.length === 0)
+    return null;
+
+  if (typeof key1 === 'string' && typeof key2 === 'string')
+    return Draft76;
+
+  return Draft75;
+};
+
+Server.determineUrl = function(request) {
+  var scheme = this.isSecureRequest(request) ? 'wss:' : 'ws:';
+  return scheme + '//' + request.headers.host + request.url;
+};
+
 Server.isSecureRequest = function(request) {
   if (request.connection && request.connection.authorized !== undefined) return true;
   if (request.socket && request.socket.secure) return true;
@@ -85,24 +133,4 @@ Server.isSecureRequest = function(request) {
   return false;
 };
 
-Server.determineUrl = function(request) {
-  var scheme = this.isSecureRequest(request) ? 'wss:' : 'ws:';
-  return scheme + '//' + request.headers.host + request.url;
-};
-
-Server.http = function(request, options) {
-  options = options || {};
-  if (options.requireMasking === undefined) options.requireMasking = true;
-
-  var headers = request.headers,
-      url     = this.determineUrl(request);
-
-  if (headers['sec-websocket-version'])
-    return new Hybi(request, url, options);
-  else if (headers['sec-websocket-key1'])
-    return new Draft76(request, url, options);
-  else
-    return new Draft75(request, url, options);
-};
-
 module.exports = Server;