Improve body buffering/streaming to avoid possible deadlock

It was possible to deadlock request body streaming with
beforeRequest. This could happen because the passthrough handler
creates a request stream (with asStream) ready to send upstream (this
is required, to ensure that's registered before any possible response is
sent, or node can dump the body) but beforeRequest itself isn't called
until the completed body is buffered.

This deadlocked because the stream could hit its high watermark (pausing
the request stream entirely) before the body was buffered. The stream
would never continue, because it wasn't connected to a sink yet, but it
couldn't be connected because beforeRequest needed to run before
creating the upstream connection.

This is resolved through two changes:
* asStream no longer actually streams data - it only reads into asBuffer
  for as long as possible (at least until truncation). Real streaming
  (and backpressure) only begins once the buffer truncates or a sink
  appears.
* Truncation stops the input entirely and delegates entirely to
  asStream. This means that backpressure kicks in at this point.

The end result should be that we only store maxSize + highWatermark in
our buffers at any time, applying backpressure for anything else, but
avoid this deadlock before that point by keeping the stream inactive
while buffering is happening. That means if you're just waiting on
waitForCompletedRequest then stream backpressure is not happening.

The only remaining possible deadlock would happen if a handler created a
stream *and started streaming from it* but then paused it so it couldn't
drain fully, and then called waitForCompletedRequest. Don't do that - if
you start streaming, make sure you will finish.

This is all very complicated, but I _think_ this should improve things
significantly.

Author: pimterry

src/util/buffer-utils.ts

@@ -1,4 +1,5 @@
 import * as _ from 'lodash';
+import { EventEmitter } from 'events';
 import * as stream from 'stream';
 
 import { isNode } from './util';
@@ -16,28 +17,60 @@ export const asBuffer = (input: Buffer | Uint8Array | string) =>
         : Buffer.from(input);
 
 export type BufferInProgress = Promise<Buffer> & {
-    currentChunks: Buffer[] // Stores the body chunks as they arrive
-    failedWith?: Error // Stores the error that killed the stream, if one did
+    currentChunks: Buffer[]; // Stores the body chunks as they arrive
+    failedWith?: Error; // Stores the error that killed the stream, if one did
+    events: EventEmitter; // Emits events - notably 'truncate' if data is truncated
 };
 
-// Takes a buffer and a stream, returns a simple stream that outputs the buffer then the stream.
+// Takes a buffer and a stream, returns a simple stream that outputs the buffer then the stream. The stream
+// is lazy, so doesn't read data in from the buffer or input until something here starts reading.
 export const bufferThenStream = (buffer: BufferInProgress, inputStream: stream.Readable): stream.Readable => {
-    const outputStream = new stream.PassThrough();
+    let active = false;
+
+    const outputStream = new stream.PassThrough({
+        // Note we use the default highWaterMark, which means this applies backpressure, pushing buffering
+        // onto the OS & backpressure on network instead of accepting data before we're ready to stream it.
+
+        // Without changing behaviour, we listen for read() events, and don't start streaming until we get one.
+        read(size) {
+            // On the first actual read of this stream, we pull from the buffer
+            // and then hook it up to the input.
+            if (!active) {
+                if (buffer.failedWith) {
+                    outputStream.destroy(buffer.failedWith);
+                } else {
+                    // First stream everything that's been buffered so far:
+                    outputStream.write(Buffer.concat(buffer.currentChunks));
+
+                    // Then start streaming all future incoming data:
+                    inputStream.pipe(outputStream);
+
+                    if (inputStream.readableEnded) outputStream.end();
+                    if (inputStream.readableAborted) outputStream.destroy();
+
+                    // Forward any future errors from the input stream:
+                    inputStream.on('error', (e) => {
+                        outputStream.emit('error', e)
+                    });
+
+                    // Silence 'unhandled rejection' warnings here, since we'll handle
+                    // them on the stream instead
+                    buffer.catch(() => {});
+                }
+                active = true;
+            }
 
-    // Forward the buffered data so far
-    outputStream.write(Buffer.concat(buffer.currentChunks));
-    // After the data, forward errors from the buffer
-    if (buffer.failedWith) {
-        // Announce async, to ensure listeners have time to get set up
-        setTimeout(() => outputStream.emit('error', buffer.failedWith));
-    } else {
-        // Forward future data as it arrives
-        inputStream.pipe(outputStream);
-        // Forward any future errors from the input stream
-        inputStream.on('error', (e) => outputStream.emit('error', e));
-        // Silence 'unhandled rejection' warnings here, since we'll handle them on the stream instead
-        buffer.catch(() => {});
-    }
+            // Except for the first activation logic (above) do the default transform read() steps just
+            // like a normal PassThrough stream.
+            return stream.Transform.prototype._read.call(this, size);
+        }
+    });
+
+    buffer.events.on('truncate', (chunks) => {
+        // If the stream hasn't started yet, start it now, so it grabs the buffer
+        // data before it gets truncated:
+        if (!active) outputStream.read(0);
+    });
 
     return outputStream;
 };
@@ -63,20 +96,41 @@ export const streamToBuffer = (input: stream.Readable, maxSize = MAX_BUFFER_SIZE
             if (input.readableAborted) return failWithAbortError();
 
             let currentSize = 0;
-            input.on('data', (d: Buffer) => {
+            const onData = (d: Buffer) => {
                 currentSize += d.length;
+                chunks.push(d);
 
                 // If we go over maxSize, drop the whole stream, so the buffer
                 // resolves empty. MaxSize should be large, so this is rare,
                 // and only happens as an alternative to crashing the process.
                 if (currentSize > maxSize) {
-                    chunks = []; // Drop all the data so far
-                    return; // Don't save any more data
+                    // Announce truncation, so that other mechanisms (asStream) can
+                    // capture this data if they're interested in it.
+                    bufferPromise.events.emit('truncate', chunks);
+
+                    // Drop all the data so far & stop reading
+                    bufferPromise.currentChunks = chunks = [];
+                    input.removeListener('data', onData);
+
+                    // We then resolve immediately - the buffer is done, even if the body
+                    // might still be streaming in we're not listening to it. This means
+                    // that requests can 'complete' for event/callback purposes while
+                    // they're actually still streaming, but only in this scenario where
+                    // the data is too large to really be used by the events/callbacks.
+
+                    // If we don't resolve, then cases which intentionally don't consume
+                    // the raw stream but do consume the buffer (beforeRequest) would
+                    // deadlock: beforeRequest must complete to begin streaming the
+                    // full body to the target clients.
+
+                    resolve(Buffer.from([]));
                 }
+            };
+            input.on('data', onData);
 
-                chunks.push(d);
+            input.once('end', () => {
+                resolve(Buffer.concat(chunks));
             });
-            input.once('end', () => resolve(Buffer.concat(chunks)));
             input.once('aborted', failWithAbortError);
             input.on('error', (e) => {
                 bufferPromise.failedWith = bufferPromise.failedWith || e;
@@ -85,6 +139,7 @@ export const streamToBuffer = (input: stream.Readable, maxSize = MAX_BUFFER_SIZE
         }
     );
     bufferPromise.currentChunks = chunks;
+    bufferPromise.events = new EventEmitter();
     return bufferPromise;
 };
 

src/util/request-utils.ts

@@ -174,9 +174,14 @@ const parseBodyStream = (bodyStream: stream.Readable, maxSize: number, getHeader
         // and buffered data, and then continues with the live stream, if active.
         // Listeners to this stream *must* be attached synchronously after this call.
         asStream() {
-            return completedBuffer
-                ? bufferToStream(completedBuffer)
-                : bufferThenStream(body.asBuffer(), bodyStream);
+            // If we've already buffered the whole body, just stream it out:
+            if (completedBuffer) return bufferToStream(completedBuffer);
+
+            // Otherwise, we want to start buffering now, and wrap that with
+            // a stream that can live-stream the buffered data on demand:
+            const buffer = body.asBuffer();
+            buffer.catch(() => {}); // Errors will be handled via the stream, so silence unhandled rejections here.
+            return bufferThenStream(buffer, bodyStream);
         },
         asBuffer() {
             if (!bufferPromise) {

test/integration/proxying/http-proxying.spec.ts

@@ -1045,6 +1045,37 @@ nodeOnly(() => {
                 expect(await proxiedRequestData.body.getText()).to.equal('');
             });
 
+            it("should still proxy larger request bodies even if beforeRequest is used", async () => {
+                const message = "A large request body";
+
+                const remoteEndpoint = await remoteServer.forAnyRequest().thenReply(200);
+                const proxyEndpoint = await server.forPost(remoteServer.url).thenPassThrough({
+                    beforeRequest: async (req) => {
+                        const bodyText = await req.body.getText();
+                        // Body is too long, and so should be truncated when examined (as here)
+                        expect(bodyText).to.equal('');
+                        return {};
+                    }
+                });
+
+                let response = await request.post({
+                    url: remoteServer.url,
+                    body: message,
+                    resolveWithFullResponse: true
+                });
+
+                expect(response.statusCode).to.equal(200);
+
+                // Even though it's truncated for buffering, the request data should still be proxied
+                // through successfully to the upstream server:
+                const resultingRequest = (await remoteEndpoint.getSeenRequests())[0];
+                expect(await resultingRequest.body.getText()).to.equal(message);
+
+                // But it's truncated in event data, not buffered
+                const proxiedRequestData = (await proxyEndpoint.getSeenRequests())[0];
+                expect(await proxiedRequestData.body.getText()).to.equal('');
+            });
+
             it("should still proxy larger response bodies", async () => {
                 await remoteServer.forAnyRequest().thenReply(200, "A large response body");
                 const proxyEndpoint = await server.forGet(remoteServer.url).thenPassThrough();