http: add shouldUpgradeCallback to let servers control HTTP upgrades

Previously, you could either add no 'upgrade' event handler, in which
case all upgrades were ignored, or add an 'upgrade' handler and all
upgrade attempts would effectively succeed and skip normal request
handling. This change adds a new shouldUpgradeCallback option to HTTP
servers, which receives the request details and returns a boolean that
controls whether the request should be upgraded.

Author: pimterry

doc/api/http.md

@@ -1671,6 +1671,11 @@ per connection (in the case of HTTP Keep-Alive connections).
 <!-- YAML
 added: v0.1.94
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/00000
+    description: Whether this event is fired can now be controlled by the
+                 `shouldUpgradeCallback` and sockets will be destroyed
+                 if upgraded while no event handler is listening.
   - version: v10.0.0
     pr-url: https://github.com/nodejs/node/pull/19981
     description: Not listening to this event no longer causes the socket
@@ -1682,13 +1687,21 @@ changes:
 * `socket` {stream.Duplex} Network socket between the server and client
 * `head` {Buffer} The first packet of the upgraded stream (may be empty)
 
-Emitted each time a client requests an HTTP upgrade. Listening to this event
-is optional and clients cannot insist on a protocol change.
+Emitted each time a client's HTTP upgrade request is accepted. By default
+all HTTP upgrade requests are ignored unless you listen to this event, in which
+case they are all accepted. You can control this more precisely by using the
+server `shouldUpgradeCallback` option.
+
+Listening to this event is optional and clients cannot insist on a protocol
+change.
 
 After this event is emitted, the request's socket will not have a `'data'`
 event listener, meaning it will need to be bound in order to handle data
 sent to the server on that socket.
 
+If an upgrade is accepted by `shouldUpgradeCallback` but no event handler
+is registered, the socket is destroyed.
+
 This event is guaranteed to be passed an instance of the {net.Socket} class,
 a subclass of {stream.Duplex}, unless the user specifies a socket
 type other than {net.Socket}.
@@ -3537,6 +3550,9 @@ Found'`.
 <!-- YAML
 added: v0.1.13
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/00000
+    description: The `shouldUpgradeCallback` option is now supported.
   - version:
     - v20.1.0
     - v18.17.0
@@ -3626,6 +3642,12 @@ changes:
   * `ServerResponse` {http.ServerResponse} Specifies the `ServerResponse` class
     to be used. Useful for extending the original `ServerResponse`. **Default:**
     `ServerResponse`.
+  * `shouldUpgradeCallback` {Function} A callback which receives an incoming
+    request and returns a boolean, to control which upgrade attempts should be
+    accepted. Accepted upgrades will fire an `'upgrade'` event (or their sockets
+    will be destroyed, if no listener is registered) while rejected upgrades will
+    fire a `'request'` event like any non-upgrade request. This options defaults
+    to `() => server.listenerCount('upgrade') > 0`.
   * `uniqueHeaders` {Array} A list of response headers that should be sent only
     once. If the header's value is an array, the items will be joined
     using `; `.

lib/_http_server.js

@@ -91,6 +91,7 @@ const {
   validateBoolean,
   validateLinkHeaderValue,
   validateObject,
+  validateFunction,
 } = require('internal/validators');
 const Buffer = require('buffer').Buffer;
 const { setInterval, clearInterval } = require('timers');
@@ -522,6 +523,16 @@ function storeHTTPOptions(options) {
   } else {
     this.rejectNonStandardBodyWrites = false;
   }
+
+  const shouldUpgradeCallback = options.shouldUpgradeCallback;
+  if (shouldUpgradeCallback !== undefined) {
+    validateFunction(shouldUpgradeCallback, 'options.shouldUpgradeCallback');
+    this.shouldUpgradeCallback = shouldUpgradeCallback;
+  } else {
+    this.shouldUpgradeCallback = function() {
+      return this.listenerCount('upgrade') > 0;
+    };
+  }
 }
 
 function setupConnectionsTracking() {
@@ -957,15 +968,15 @@ function onParserExecuteCommon(server, socket, parser, state, ret, d) {
     parser = null;
 
     const eventName = req.method === 'CONNECT' ? 'connect' : 'upgrade';
-    if (eventName === 'upgrade' || server.listenerCount(eventName) > 0) {
+    if (server.listenerCount(eventName) > 0) {
       debug('SERVER have listener for %s', eventName);
       const bodyHead = d.slice(ret, d.length);
 
       socket.readableFlowing = null;
 
       server.emit(eventName, req, socket, bodyHead);
     } else {
-      // Got CONNECT method, but have no handler.
+      // Got upgrade or CONNECT method, but have no handler.
       socket.destroy();
     }
   } else if (parser.incoming && parser.incoming.method === 'PRI') {
@@ -1059,7 +1070,7 @@ function parserOnIncoming(server, socket, state, req, keepAlive) {
 
   if (req.upgrade) {
     req.upgrade = req.method === 'CONNECT' ||
-                  server.listenerCount('upgrade') > 0;
+                  !!server.shouldUpgradeCallback(req);
     if (req.upgrade)
       return 2;
   }

test/parallel/test-http-upgrade-server-callback.js

@@ -0,0 +1,167 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+
+const net = require('net');
+const http = require('http');
+
+function testUpgradeCallbackTrue() {
+  const server = http.createServer({
+    shouldUpgradeCallback: common.mustCall((req) => {
+      assert.strictEqual(req.url, '/websocket');
+      assert.strictEqual(req.headers.upgrade, 'websocket');
+
+      return true;
+    })
+  });
+
+  server.on('upgrade', function(req, socket, upgradeHead) {
+    assert.strictEqual(req.url, '/websocket');
+    assert.strictEqual(req.headers.upgrade, 'websocket');
+    assert.ok(socket instanceof net.Socket);
+    assert.ok(upgradeHead instanceof Buffer);
+
+    socket.write('HTTP/1.1 101 Web Socket Protocol Handshake\r\n' +
+      'Upgrade: WebSocket\r\n' +
+      'Connection: Upgrade\r\n' +
+      '\r\n\r\n');
+  });
+
+  server.on('request', common.mustNotCall());
+
+  server.listen(0, common.mustCall(() => {
+    const req = http.request({
+      port: server.address().port,
+      path: '/websocket',
+      headers: {
+        'Upgrade': 'websocket',
+        'Connection': 'Upgrade'
+      }
+    });
+
+    req.on('upgrade', common.mustCall((res, socket, upgradeHead) => {
+      assert.strictEqual(res.statusCode, 101);
+      assert.ok(socket instanceof net.Socket);
+      assert.ok(upgradeHead instanceof Buffer);
+      socket.end();
+      server.close();
+
+      testUpgradeCallbackFalse();
+    }));
+
+    req.on('response', common.mustNotCall());
+    req.end();
+  }));
+}
+
+
+function testUpgradeCallbackFalse() {
+  const server = http.createServer({
+    shouldUpgradeCallback: common.mustCall(() => {
+      return false;
+    })
+  });
+
+  server.on('upgrade', common.mustNotCall());
+
+  server.on('request', common.mustCall((req, res) => {
+    res.writeHead(200, { 'Content-Type': 'text/plain' });
+    res.write('received but not upgraded');
+    res.end();
+  }));
+
+  server.listen(0, common.mustCall(() => {
+    const req = http.request({
+      port: server.address().port,
+      path: '/websocket',
+      headers: {
+        'Upgrade': 'websocket',
+        'Connection': 'Upgrade'
+      }
+    });
+
+    req.on('upgrade', common.mustNotCall());
+
+    req.on('response', common.mustCall((res) => {
+      assert.strictEqual(res.statusCode, 200);
+      let data = '';
+      res.on('data', (chunk) => { data += chunk; });
+      res.on('end', common.mustCall(() => {
+        assert.strictEqual(data, 'received but not upgraded');
+        server.close();
+
+        testUpgradeCallbackTrueWithoutHandler();
+      }));
+    }));
+    req.end();
+  }));
+}
+
+
+function testUpgradeCallbackTrueWithoutHandler() {
+  const server = http.createServer({
+    shouldUpgradeCallback: common.mustCall(() => {
+      return true;
+    })
+  });
+
+  // N.b: no 'upgrade' handler
+  server.on('request', common.mustNotCall());
+
+  server.listen(0, common.mustCall(() => {
+    const req = http.request({
+      port: server.address().port,
+      path: '/websocket',
+      headers: {
+        'Upgrade': 'websocket',
+        'Connection': 'Upgrade'
+      }
+    });
+
+    req.on('upgrade', common.mustNotCall());
+    req.on('response', common.mustNotCall());
+
+    req.on('error', common.mustCall((e) => {
+      assert.strictEqual(e.code, 'ECONNRESET');
+      server.close();
+
+      testUpgradeCallbackError();
+    }));
+    req.end();
+  }));
+}
+
+
+function testUpgradeCallbackError() {
+  const server = http.createServer({
+    shouldUpgradeCallback: common.mustCall(() => {
+      throw new Error('should upgrade callback failed');
+    })
+  });
+
+  server.on('upgrade', common.mustNotCall());
+  server.on('request', common.mustNotCall());
+
+  server.listen(0, common.mustCall(() => {
+    const req = http.request({
+      port: server.address().port,
+      path: '/websocket',
+      headers: {
+        'Upgrade': 'websocket',
+        'Connection': 'Upgrade'
+      }
+    });
+
+    req.on('upgrade', common.mustNotCall());
+    req.on('response', common.mustNotCall());
+
+    process.on('uncaughtException', common.mustCall(() => {
+      process.exit(0);
+    }));
+
+    req.end();
+  }));
+}
+
+testUpgradeCallbackTrue();