http: add shouldUpgradeCallback to let servers control HTTP upgrades

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/59824
+    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/59824
+    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();