Return byte streams when tee()ing byte streams, closes #317

mrbbot · mrbbot · commit 31792ba7d98f · 2022-09-13T18:52:00.000+01:00
Closes #375
diff --git a/packages/core/src/standards/http.ts b/packages/core/src/standards/http.ts
@@ -5,7 +5,6 @@ import { Blob } from "buffer";
 import type EventEmitter from "events";
 import http from "http";
 import {
-  ReadableByteStreamController,
   ReadableStream,
   ReadableStreamDefaultReader,
   UnderlyingByteSource,
@@ -55,7 +54,7 @@ import {
   buildNotBufferSourceError,
   isBufferSource,
 } from "./helpers";
-import { kContentLength } from "./streams";
+import { _isByteStream, kContentLength } from "./streams";
 
 // We need these for making Request's Headers immutable
 const fetchSymbols: {
@@ -162,24 +161,6 @@ const kInputGated = Symbol("kInputGated");
 const kFormDataFiles = Symbol("kFormDataFiles");
 const kCloned = Symbol("kCloned");
 
-/** @internal */
-export function _isByteStream(
-  stream: ReadableStream
-): stream is ReadableStream<Uint8Array> {
-  // Try to determine if stream is a byte stream by inspecting its state.
-  // It doesn't matter too much if the internal representation changes in the
-  // future: this code shouldn't throw. Currently we only use this as an
-  // optimisation to avoid creating a byte stream if it's already one.
-  for (const symbol of Object.getOwnPropertySymbols(stream)) {
-    if (symbol.description === "kState") {
-      // @ts-expect-error symbol properties are not included in type definitions
-      const controller = stream[symbol].controller;
-      return controller instanceof ReadableByteStreamController;
-    }
-  }
-  return false;
-}
-
 const enumerableBodyKeys: (keyof Body<any>)[] = ["body", "bodyUsed", "headers"];
 export class Body<Inner extends BaseRequest | BaseResponse> {
   /** @internal */
diff --git a/packages/core/src/standards/index.ts b/packages/core/src/standards/index.ts
@@ -7,7 +7,6 @@ export * from "./event";
 export {
   _headersFromIncomingRequest,
   _kInner,
-  _isByteStream,
   Body,
   withInputGating,
   withStringFormDataFiles,
@@ -48,6 +47,7 @@ export {
   FixedLengthStream,
   CompressionStream,
   DecompressionStream,
+  _isByteStream,
 } from "./streams";
 export type { ArrayBufferViewConstructor } from "./streams";
 export * from "./navigator";
diff --git a/packages/core/src/standards/streams.ts b/packages/core/src/standards/streams.ts
@@ -1,5 +1,6 @@
 import type { Transform } from "stream";
 import {
+  ReadableByteStreamController,
   ReadableStream,
   ReadableStreamBYOBReadResult,
   ReadableStreamBYOBReader,
@@ -15,6 +16,57 @@ import {
   isBufferSource,
 } from "./helpers";
 
+/** @internal */
+export function _isByteStream(
+  stream: ReadableStream
+): stream is ReadableStream<Uint8Array> {
+  // Try to determine if stream is a byte stream by inspecting its state.
+  // It doesn't matter too much if the internal representation changes in the
+  // future: this code shouldn't throw. Currently we only use this as an
+  // optimisation to avoid creating a byte stream if it's already one.
+  for (const symbol of Object.getOwnPropertySymbols(stream)) {
+    if (symbol.description === "kState") {
+      // @ts-expect-error symbol properties are not included in type definitions
+      const controller = stream[symbol].controller;
+      return controller instanceof ReadableByteStreamController;
+    }
+  }
+  return false;
+}
+
+function convertToByteStream(
+  stream: ReadableStream<Uint8Array>,
+  clone = false
+) {
+  let reader: ReadableStreamDefaultReader<Uint8Array>;
+  return new ReadableStream({
+    type: "bytes",
+    start() {
+      reader = stream.getReader();
+    },
+    async pull(controller) {
+      let result = await reader.read();
+      while (!result.done && result.value.byteLength === 0) {
+        result = await reader.read();
+      }
+      if (result.done) {
+        queueMicrotask(() => {
+          controller.close();
+          // Not documented in MDN but if there's an ongoing request that's
+          // waiting, we need to tell it that there were 0 bytes delivered so
+          // that it unblocks and notices the end of stream.
+          // @ts-expect-error `byobRequest` has type `undefined` in `@types/node`
+          controller.byobRequest?.respond(0);
+        });
+      } else if (result.value.byteLength > 0) {
+        if (clone) result.value = result.value.slice();
+        controller.enqueue(result.value);
+      }
+    },
+    cancel: (reason) => reader.cancel(reason),
+  });
+}
+
 export type ArrayBufferViewConstructor =
   | typeof Int8Array
   | typeof Uint8Array
@@ -83,6 +135,28 @@ ReadableStreamBYOBReader.prototype.readAtLeast = async function <
   return { value: value as any, done };
 };
 
+// See comment above about manipulating the prototype.
+// Rewrite tee() to return byte streams when tee()ing byte streams:
+// https://github.com/cloudflare/miniflare/issues/317
+const originalTee = ReadableStream.prototype.tee;
+ReadableStream.prototype.tee = function () {
+  if (!(this instanceof ReadableStream)) {
+    throw new TypeError("Illegal invocation");
+  }
+  if (_isByteStream(this)) {
+    const [stream1, stream2] = originalTee.call(this);
+    // We need to clone chunks here, as either of the tee()ed streams might be
+    // passed to `new Response()`, which will detach array buffers when reading:
+    // https://github.com/cloudflare/miniflare/issues/375
+    return [
+      convertToByteStream(stream1, true /* clone */),
+      convertToByteStream(stream2, true /* clone */),
+    ];
+  } else {
+    return originalTee.call(this);
+  }
+};
+
 const kTransformHook = Symbol("kTransformHook");
 const kFlushHook = Symbol("kFlushHook");
 
@@ -119,41 +193,7 @@ export class IdentityTransformStream extends TransformStream<
 
   get readable() {
     if (this.#readableByteStream !== undefined) return this.#readableByteStream;
-    const readable = super.readable;
-    let reader: ReadableStreamDefaultReader;
-    return (this.#readableByteStream = new ReadableStream({
-      type: "bytes",
-      start() {
-        reader = readable.getReader();
-      },
-      async pull(controller) {
-        let { done, value } = await reader.read();
-        // Make sure we eventually call a `controller` method, either because
-        // we're done, or there's data to enqueue
-        while (!done && value.byteLength === 0) {
-          const result = await reader.read();
-          done = result.done;
-          value = result.value;
-        }
-        if (done) {
-          queueMicrotask(() => {
-            controller.close();
-            // Not documented in MDN but if there's an ongoing request that's waiting,
-            // we need to tell it that there were 0 bytes delivered so that it unblocks
-            // and notices the end of stream.
-            // @ts-expect-error `byobRequest` has type `undefined` in `@types/node`
-            controller.byobRequest?.respond(0);
-          });
-        } else if (value.byteLength > 0) {
-          // Ensure chunk is non-empty before enqueuing:
-          // https://github.com/cloudflare/miniflare/issues/374
-          controller.enqueue(value);
-        }
-      },
-      cancel() {
-        return reader.cancel();
-      },
-    }));
+    return (this.#readableByteStream = convertToByteStream(super.readable));
   }
 }
 
diff --git a/packages/core/test/standards/http.spec.ts b/packages/core/test/standards/http.spec.ts
@@ -10,7 +10,6 @@ import {
   Response,
   _getBodyLength,
   _getURLList,
-  _isByteStream,
   createCompatFetch,
   fetch,
   logResponse,
@@ -100,24 +99,6 @@ test("Headers: getAll: returns separated Set-Cookie values", (t) => {
   t.deepEqual(headers.getAll("set-CoOkiE"), [cookie1, cookie2, cookie3]);
 });
 
-test("_isByteStream: determines if a ReadableStream is a byte stream", (t) => {
-  const regularStream = new ReadableStream({
-    pull(controller) {
-      controller.enqueue(new Uint8Array([1, 2, 3]));
-      controller.close();
-    },
-  });
-  const byteStream = new ReadableStream({
-    type: "bytes",
-    pull(controller) {
-      controller.enqueue(new Uint8Array([1, 2, 3]));
-      controller.close();
-    },
-  });
-  t.false(_isByteStream(regularStream));
-  t.true(_isByteStream(byteStream));
-});
-
 // These tests also implicitly test withInputGating
 test("Body: body isn't input gated by default", async (t) => {
   const inputGate = new InputGate();
diff --git a/packages/core/test/standards/streams.spec.ts b/packages/core/test/standards/streams.spec.ts
@@ -1,5 +1,5 @@
 import assert from "assert";
-import { arrayBuffer } from "stream/consumers";
+import { arrayBuffer, text } from "stream/consumers";
 import {
   ReadableStream,
   ReadableStreamBYOBReadResult,
@@ -13,6 +13,7 @@ import {
   IdentityTransformStream,
   Request,
   Response,
+  _isByteStream,
 } from "@miniflare/core";
 import { utf8Decode, utf8Encode } from "@miniflare/shared-test";
 import test, { Macro, ThrowsExpectation } from "ava";
@@ -56,6 +57,24 @@ async function* byobReadAtLeast<Ctor extends ArrayBufferViewConstructor>(
   }
 }
 
+test("_isByteStream: determines if a ReadableStream is a byte stream", (t) => {
+  const regularStream = new ReadableStream({
+    pull(controller) {
+      controller.enqueue(new Uint8Array([1, 2, 3]));
+      controller.close();
+    },
+  });
+  const byteStream = new ReadableStream({
+    type: "bytes",
+    pull(controller) {
+      controller.enqueue(new Uint8Array([1, 2, 3]));
+      controller.close();
+    },
+  });
+  t.false(_isByteStream(regularStream));
+  t.true(_isByteStream(byteStream));
+});
+
 test("ReadableStreamBYOBReader: readAtLeast: reads at least n bytes", async (t) => {
   const stream = chunkedStream([[1, 2, 3], [4], [5, 6]]);
   const reads = byobReadAtLeast(stream, 4, 8, Uint8Array);
@@ -189,6 +208,50 @@ test("ReadableStreamBYOBReader: readAtLeast: throws if minimum number of bytes e
   });
 });
 
+test("ReadableStream: tee: returns byte streams when teeing byte stream", async (t) => {
+  // https://github.com/cloudflare/miniflare/issues/317
+  const stream = chunkedStream([[1, 2, 3]]);
+  t.true(_isByteStream(stream));
+  const [stream1, stream2] = stream.tee();
+  t.true(_isByteStream(stream1));
+  t.true(_isByteStream(stream2));
+
+  // Check detaching array buffers on one stream doesn't affect the other
+  // (note Response#body will detach chunks as reading)
+  // https://github.com/cloudflare/miniflare/issues/375
+  const body1 = await arrayBuffer(new Response(stream1).body as any);
+  const body2 = await arrayBuffer(new Response(stream2).body as any);
+  t.deepEqual(new Uint8Array(body1), new Uint8Array([1, 2, 3]));
+  t.deepEqual(new Uint8Array(body2), new Uint8Array([1, 2, 3]));
+});
+test("ReadableStream: tee: returns regular stream when teeing regular stream", async (t) => {
+  const stream = new ReadableStream<string>({
+    start(controller) {
+      controller.enqueue("value");
+      controller.close();
+    },
+  });
+  t.false(_isByteStream(stream));
+  const [stream1, stream2] = stream.tee();
+  t.false(_isByteStream(stream1));
+  t.false(_isByteStream(stream2));
+  t.is(await text(stream1 as any), "value");
+  t.is(await text(stream2 as any), "value");
+});
+test("ReadableStream: tee: throws on illegal invocation", (t) => {
+  const stream = new ReadableStream<string>({
+    start(controller) {
+      controller.close();
+    },
+  });
+  // @ts-expect-error using comma expression to unbind this
+  // noinspection CommaExpressionJS
+  t.throws(() => (0, stream.tee)(), {
+    instanceOf: TypeError,
+    message: "Illegal invocation",
+  });
+});
+
 const identityMacro: Macro<[TransformStream]> = async (t, stream) => {
   const { readable, writable } = stream;
   const writer = writable.getWriter();
