Extract body buffering into stream-capture (with small fixes en route)

Author: pimterry

src/rules/requests/request-step-impls.ts

@@ -1524,7 +1524,9 @@ export class WaitForRequestBodyStepImpl extends WaitForRequestBodyStep {
     static readonly fromDefinition = () => new WaitForRequestBodyStepImpl();
 
     async handle(request: OngoingRequest): Promise<{ continue: true }> {
-        await request.body.asBuffer();
+        // Wait for body delivery to complete - not just buffering to end (they can be
+        // different, e.g. if buffer tops out the max buffering limits).
+        await request.body.waitForEnd();
         return { continue: true };
     }
 }

src/types.ts

@@ -214,7 +214,15 @@ export interface OngoingRequest extends Request, stream.Readable {
 
 export interface OngoingBody {
     asStream: () => stream.Readable;
+    /**
+     * Resolves with the body content when it's done, or an empty body if the body is dropped
+     * due to hitting the max buffering limits.
+     */
     asBuffer: () => Promise<Buffer>;
+    /**
+     * Resolves at the end of the body stream - regardless of buffering behaviour.
+     */
+    waitForEnd: () => Promise<void>;
     asDecodedBuffer: () => Promise<Buffer>;
     asText: () => Promise<string>;
     asJson: () => Promise<object>;

src/util/buffer-utils.ts

@@ -1,13 +1,6 @@
 import { Buffer } from 'buffer';
-import { EventEmitter } from 'events';
 import * as stream from 'stream';
 
-import { isNode } from './util';
-
-const MAX_BUFFER_SIZE = isNode
-    ? require('buffer').constants.MAX_LENGTH
-    : Infinity;
-
 export const asBuffer = (input: Buffer | Uint8Array | string) =>
     Buffer.isBuffer(input)
         ? input
@@ -16,132 +9,29 @@ export const asBuffer = (input: Buffer | Uint8Array | string) =>
     // Is Array:
         : 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
-    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. 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 => {
-    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;
-            }
-
-            // 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;
-};
-
 export const bufferToStream = (buffer: Buffer): stream.Readable => {
     const outputStream = new stream.PassThrough();
     outputStream.end(buffer);
     return outputStream;
 };
 
-export const streamToBuffer = (input: stream.Readable, maxSize = MAX_BUFFER_SIZE) => {
-    let chunks: Buffer[] = [];
-
-    const bufferPromise = <BufferInProgress> new Promise(
-        (resolve, reject) => {
-            function failWithAbortError() {
-                bufferPromise.failedWith = new Error('Aborted');
-                reject(bufferPromise.failedWith);
-            }
-
-            // If stream has already finished/aborted, resolve accordingly immediately:
-            if (input.readableEnded) return resolve(Buffer.from([]));
-            if (input.readableAborted) return setImmediate(failWithAbortError);
-
-            let currentSize = 0;
-            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) {
-                    // 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);
-
-            input.once('end', () => {
-                resolve(Buffer.concat(chunks));
-            });
-            input.once('aborted', failWithAbortError);
-            input.on('error', (e) => {
-                bufferPromise.failedWith = bufferPromise.failedWith || e;
-                reject(e);
-            });
-        }
-    );
-    bufferPromise.currentChunks = chunks;
-    bufferPromise.events = new EventEmitter();
-    return bufferPromise;
-};
+/**
+ * Reads a stream to completion, returning all its data as a single buffer. Rejects if the
+ * stream errors or is aborted.
+ */
+export const streamToBuffer = (input: stream.Readable) => new Promise<Buffer>((resolve, reject) => {
+    const failWithAbortError = () => reject(new Error('Aborted'));
+
+    // If the stream has already finished/aborted, resolve accordingly immediately:
+    if (input.readableEnded) return resolve(Buffer.from([]));
+    if (input.readableAborted) return setImmediate(failWithAbortError);
+
+    const chunks: Buffer[] = [];
+    input.on('data', (d: Buffer) => chunks.push(d));
+    input.once('end', () => resolve(Buffer.concat(chunks)));
+    input.once('aborted', failWithAbortError);
+    input.on('error', reject);
+});
 
 export function splitBuffer(input: Buffer, splitter: string, maxParts = Infinity) {
     const parts: Buffer[] = [];

src/util/request-utils.ts

@@ -33,13 +33,10 @@ import {
 import { makePropertyWritable } from './util';
 import { Mutable } from './type-utils';
 import {
-    bufferThenStream,
-    bufferToStream,
-    BufferInProgress,
     splitBuffer,
-    streamToBuffer,
     asBuffer
 } from './buffer-utils';
+import { StreamCapture } from './stream-capture';
 import {
     flattenPairedRawHeaders,
     getHeaderValue,
@@ -145,33 +142,28 @@ export async function decodeBodyBuffer(buffer: Buffer, headers: Headers) {
 // Parse an in-progress request or response stream, i.e. where the body or possibly even the headers have
 // not been fully received/sent yet.
 const parseBodyStream = (bodyStream: stream.Readable, maxSize: number, getHeaders: () => Headers): OngoingBody => {
-    let bufferPromise: BufferInProgress | null = null;
-    let completedBuffer: Buffer | null = null;
+    let capture: StreamCapture | null = null;
+
+    const startCapture = () => {
+        if (!capture) {
+            capture = new StreamCapture(bodyStream, maxSize);
+        }
+        return capture;
+    };
 
     let body = {
         // Returns a stream for the full body, not the live streaming body.
         // Each call creates a new stream, which starts with the already seen
         // and buffered data, and then continues with the live stream, if active.
         // Listeners to this stream *must* be attached synchronously after this call.
         asStream() {
-            // 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);
+            return startCapture().takeStream();
         },
         asBuffer() {
-            if (!bufferPromise) {
-                bufferPromise = streamToBuffer(bodyStream, maxSize);
-
-                bufferPromise
-                    .then((buffer) => completedBuffer = buffer)
-                    .catch(() => {}); // If we get no body, completedBuffer stays null
-            }
-            return bufferPromise;
+            return startCapture().buffer;
+        },
+        waitForEnd() {
+            return startCapture().completed;
         },
         async asDecodedBuffer() {
             const buffer = await body.asBuffer();