Create a body-hiding upgrade stream wrapper

Author: pimterry

doc/api/http.md

@@ -1708,14 +1708,13 @@ 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.
-
-In the (uncommon) case that the incoming request has a body, this event will
-not be emitted until the body has been fully received. This differs from normal
-`'request'` event handling, but ensures that all data received on the socket
-is part of the upgraded protocol, not the initial request.
+In the uncommon case that the incoming request has a body, this body will be
+parsed as normal separate to the upgrade stream, and the raw socket data will
+only begin after it has completed. To ensure that reading from the socket isn't
+blocked bv waiting for the request body to be read, any reads on the socket
+will start the request body flowing automatically. If you want to read the
+request body, ensure that you do so (i.e. you attach `'data'` listeners)
+before starting to read from the upgraded socket.
 
 If an upgrade is accepted by `shouldUpgradeCallback` but no event handler
 is registered then the socket is destroyed, resulting in an immediate

lib/_http_server.js

@@ -26,14 +26,19 @@ const {
   Error,
   MathMin,
   NumberIsFinite,
+  ObjectGetPrototypeOf,
   ObjectKeys,
   ObjectSetPrototypeOf,
+  Proxy,
   ReflectApply,
+  ReflectGet,
+  ReflectSet,
   Symbol,
   SymbolAsyncDispose,
   SymbolFor,
 } = primordials;
 
+const { Duplex } = require('stream');
 const net = require('net');
 const EE = require('events');
 const assert = require('internal/assert');
@@ -43,6 +48,7 @@ const {
   continueExpression,
   chunkExpression,
   kIncomingMessage,
+  kSocket,
   HTTPParser,
   isLenient,
   _checkInvalidHeaderChar: checkInvalidHeaderChar,
@@ -106,6 +112,7 @@ const onResponseFinishChannel = dc.channel('http.server.response.finish');
 
 const kServerResponse = Symbol('ServerResponse');
 const kServerResponseStatistics = Symbol('ServerResponseStatistics');
+const kUpgradeStream = Symbol('UpgradeStream');
 
 const kOptimizeEmptyRequests = Symbol('OptimizeEmptyRequestsOption');
 
@@ -953,43 +960,159 @@ function socketOnError(e) {
   }
 }
 
+class UpgradeStream extends Duplex {
+  constructor(socket, req) {
+    super({
+      allowHalfOpen: socket.allowHalfOpen,
+    });
+
+    this[kSocket] = socket;
+    this[kIncomingMessage] = req;
+
+    // Proxy error, end & closure events immediately.
+    socket.on('error', (err) => this.destroy(err));
+    this.on('error', (err) => socket.destroy(err));
+
+    socket.on('close', () => this.destroy());
+    this.on('close', () => socket.destroy());
+
+    socket.on('end', () => {
+      this.push(null);
+
+      // Match the socket behaviour, where 'end' will fire despite no 'data'
+      // listeners if a socket with no pending data ends:
+      if (this.readableLength === 0) {
+        this.resume();
+      }
+    });
+
+    // Other events (most notably, reading) all only
+    // activate after requestBodyCompleted is called.
+
+    // This stream wraps itself in a proxy which forwards all non-stream
+    // property lookups back to the underlying socket, for backward
+    // compatibiility.
+    // eslint-disable-next-line no-constructor-return
+    return new Proxy(this, {
+      __proto__: null,
+      has(target, prop) {
+        return prop in target || prop in socket;
+      },
+      get(target, prop, receiver) {
+        if (prop in target) {
+          return ReflectGet(target, prop, receiver);
+        }
+
+        return socket[prop];
+      },
+      set(target, prop, value, receiver) {
+        if (prop in target) {
+          return ReflectSet(target, prop, value, receiver);
+        }
+
+        socket[prop] = value;
+        return true;
+      },
+      getPrototypeOf(target) {
+        return ObjectGetPrototypeOf(socket);
+      },
+    });
+  }
+
+  requestBodyCompleted(upgradeHead) {
+    // When the request body is completed, we begin streaming all the
+    // post-body data for the upgraded protocol:
+    if (upgradeHead?.length > 0) {
+      if (!this.push(upgradeHead)) {
+        this[kSocket].pause();
+      }
+    }
+
+    this[kSocket].on('data', (data) => {
+      if (!this.push(data)) {
+        this[kSocket].pause();
+      }
+    });
+  }
+
+  _read(size) {
+    // Reading the upgrade stream starts the request stream flowing. It's
+    // important that this happens, even if there are no listeners, or it
+    // would be impossible to read this without explicitly reading all the
+    // request body first, which is backward incompatible & awkward.
+    this[kIncomingMessage].resume();
+
+    this[kSocket].resume();
+  }
+
+  _final(callback) {
+    this[kSocket].end(callback);
+  }
+
+  _write(chunk, encoding, callback) {
+    this[kSocket].write(chunk, encoding, callback);
+  }
+}
+
 function onParserExecuteCommon(server, socket, parser, state, ret, d) {
   if (ret instanceof Error) {
     prepareError(ret, parser, d);
     debug('parse error', ret);
     socketOnError.call(socket, ret);
   } else if (parser.incoming?.upgrade) {
-    if (!parser.incoming.complete) {
-      // The request is still incomplete - this means it's an upgrade with a request
-      // body that hasn't yet been received. We have to wait for the body.
-      return;
-    }
-
     // Upgrade or CONNECT
     const req = parser.incoming;
     debug('SERVER upgrade or connect', req.method);
+    const eventName = req.method === 'CONNECT' ? 'connect' : 'upgrade';
 
-    d ||= parser.getCurrentBuffer();
+    let upgradeStream;
+    if (req.complete) {
+      d ||= parser.getCurrentBuffer();
+
+      socket.removeListener('data', state.onData);
+      socket.removeListener('end', state.onEnd);
+      socket.removeListener('close', state.onClose);
+      socket.removeListener('drain', state.onDrain);
+      socket.removeListener('error', socketOnError);
+      socket.removeListener('timeout', socketOnTimeout);
+
+      unconsume(parser, socket);
+      parser.finish();
+      freeParser(parser, req, socket);
+      parser = null;
+
+      // If the request is complete (no body, or all body read upfront) then
+      // we just emit the socket directly as the upgrade stream.
+      upgradeStream = socket;
+    } else {
+      // If the body hasn't been fully parsed yet, we emit immediately but
+      // we add a wrapper around the socket to not expose incoming data
+      // until the request body has finished.
+
+      if (socket[kUpgradeStream]) {
+        // We've already emitted the incomplete upgrade - nothing do to
+        // until actual body parsing completion.
+        return;
+      }
 
-    socket.removeListener('data', state.onData);
-    socket.removeListener('end', state.onEnd);
-    socket.removeListener('close', state.onClose);
-    socket.removeListener('drain', state.onDrain);
-    socket.removeListener('error', socketOnError);
-    socket.removeListener('timeout', socketOnTimeout);
-    unconsume(parser, socket);
-    parser.finish();
-    freeParser(parser, req, socket);
-    parser = null;
+      d ||= Buffer.from([]);
+
+      upgradeStream = new UpgradeStream(socket, req);
+      socket[kUpgradeStream] = upgradeStream;
+    }
 
-    const eventName = req.method === 'CONNECT' ? 'connect' : 'upgrade';
     if (server.listenerCount(eventName) > 0) {
       debug('SERVER have listener for %s', eventName);
-      const bodyHead = d.slice(ret, d.length);
 
-      socket.readableFlowing = null;
+      const bodyHead = d.slice(ret, d.length);
 
-      server.emit(eventName, req, socket, bodyHead);
+      if (req.complete && socket[kUpgradeStream]) {
+        // Previously emitted, now completed - just activate the stream
+        socket[kUpgradeStream].requestBodyCompleted(bodyHead);
+      } else {
+        socket.readableFlowing = null;
+        server.emit(eventName, req, upgradeStream, bodyHead);
+      }
     } else {
       // Got upgrade or CONNECT method, but have no handler.
       socket.destroy();