Copy known-length when tee() ReadableStreams, closes #506 (#512)

Previously calling `ReadableStream#tee()` would discard the
fixed-length from `FixedLengthStream`s and the body-stream marker
from `Request`/`Response` streams. This meant attempting to upload
`tee()`ed streams as R2 multipart parts would fail. This change
updates the `tee()` implementation to copy these across.

Author: mrbbot

packages/core/src/standards/http.ts

@@ -57,6 +57,7 @@ import {
 import {
   _isByteStream,
   convertToRegularStream,
+  kBodyStreamBrand,
   kContentLength,
 } from "./streams";
 
@@ -162,7 +163,6 @@ export function _headersFromIncomingRequest(
 export const _kInner = Symbol("kInner");
 
 const kBodyStream = Symbol("kBodyStream");
-const kBodyStreamBrand = Symbol("kBodyStreamBrand");
 const kInputGated = Symbol("kInputGated");
 const kFormDataFiles = Symbol("kFormDataFiles");
 const kCloned = Symbol("kCloned");

packages/core/src/standards/streams.ts

@@ -16,6 +16,8 @@ import {
   isBufferSource,
 } from "./helpers";
 
+export const kBodyStreamBrand = Symbol("kBodyStreamBrand");
+
 /** @internal */
 export function _isByteStream(
   stream: ReadableStream
@@ -177,18 +179,32 @@ ReadableStream.prototype.tee = function () {
   if (!(this instanceof ReadableStream)) {
     throw new TypeError("Illegal invocation");
   }
+  let [stream1, stream2] = originalTee.call(this);
   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);
+    stream1 = convertToByteStream(stream1, true /* clone */);
+    stream2 = convertToByteStream(stream2, true /* clone */);
+  }
+
+  // Copy known-length stream markers
+  const branded = this as {
+    [kBodyStreamBrand]?: unknown;
+    [kContentLength]?: unknown;
+  };
+  const bodyBrand = branded[kBodyStreamBrand];
+  if (bodyBrand !== undefined) {
+    Object.defineProperty(stream1, kBodyStreamBrand, { value: bodyBrand });
+    Object.defineProperty(stream2, kBodyStreamBrand, { value: bodyBrand });
+  }
+  const contentLength = branded[kContentLength];
+  if (contentLength !== undefined) {
+    Object.defineProperty(stream1, kContentLength, { value: contentLength });
+    Object.defineProperty(stream2, kContentLength, { value: contentLength });
   }
+
+  return [stream1, stream2];
 };
 
 const kTransformHook = Symbol("kTransformHook");

packages/core/test/standards/streams.spec.ts

@@ -13,8 +13,10 @@ import {
   IdentityTransformStream,
   Request,
   Response,
+  _isBodyStream,
   _isByteStream,
   _isDisturbedStream,
+  _isFixedLengthStream,
 } from "@miniflare/core";
 import { utf8Decode, utf8Encode } from "@miniflare/shared-test";
 import test, { Macro, ThrowsExpectation } from "ava";
@@ -268,6 +270,24 @@ test("ReadableStream: tee: returns regular stream when teeing regular stream", a
   t.is(await text(stream1 as any), "value");
   t.is(await text(stream2 as any), "value");
 });
+test("ReadableStream: tee: returns known length streams when teeing known length streams", async (t) => {
+  const fixedLengthStream = new FixedLengthStream(3);
+  let [stream1, stream2] = fixedLengthStream.readable.tee();
+  t.true(_isFixedLengthStream(stream1));
+  t.true(_isFixedLengthStream(stream2));
+
+  const req = new Request("http://localhost", { method: "POST", body: "body" });
+  assert(req.body !== null);
+  [stream1, stream2] = req.body.tee();
+  t.true(_isBodyStream(stream1));
+  t.true(_isBodyStream(stream2));
+
+  const res = new Response("body");
+  assert(res.body !== null);
+  [stream1, stream2] = res.body.tee();
+  t.true(_isBodyStream(stream1));
+  t.true(_isBodyStream(stream2));
+});
 test("ReadableStream: tee: throws on illegal invocation", (t) => {
   const stream = new ReadableStream<string>({
     start(controller) {

packages/r2/test/multipart.spec.ts

@@ -177,8 +177,9 @@ test("R2MultipartUpload: uploadPart", async (t) => {
   const response = new Response("value7");
   assert(request.body !== null && response.body !== null);
   await upload.uploadPart(5, readable);
-  await upload.uploadPart(6, request.body);
-  await upload.uploadPart(7, response.body);
+  // Check `tee()`ing body inherits known length
+  await upload.uploadPart(6, request.body.tee()[0]);
+  await upload.uploadPart(7, response.body.tee()[1]);
   const unknownLengthReadable = new ReadableStream({
     type: "bytes",
     pull(controller) {