Implement FixedLengthStream, closes #123

Author: mrbbot

packages/core/src/plugins/core.ts

@@ -39,6 +39,7 @@ import {
   AbortSignal,
   DOMException,
   FetchEvent,
+  FixedLengthStream,
   Request,
   Response,
   ScheduledEvent,
@@ -365,6 +366,7 @@ export class CorePlugin extends Plugin<CoreOptions> implements CoreOptions {
       WritableStream,
       WritableStreamDefaultController,
       WritableStreamDefaultWriter,
+      FixedLengthStream,
 
       Event,
       EventTarget,

packages/core/src/standards/http.ts

@@ -48,6 +48,7 @@ import {
   buildNotBufferSourceError,
   isBufferSource,
 } from "./helpers";
+import { kContentLength } from "./streams";
 
 const inspect = Symbol.for("nodejs.util.inspect.custom");
 const nonEnumerable = Object.create(null);
@@ -334,6 +335,14 @@ export class Request extends Body<BaseRequest> {
     }
     this.#cf = cf ? nonCircularClone(cf) : undefined;
 
+    // If body is a FixedLengthStream, set Content-Length to its expected length
+    const contentLength: number | undefined = (init?.body as any)?.[
+      kContentLength
+    ];
+    if (contentLength !== undefined) {
+      this.headers.set("content-length", contentLength.toString());
+    }
+
     makeEnumerable(Request.prototype, this, enumerableRequestKeys);
   }
 
@@ -491,6 +500,12 @@ export class Response<
     this.#status = status;
     this.#webSocket = webSocket;
 
+    // If body is a FixedLengthStream, set Content-Length to its expected length
+    const contentLength: number | undefined = (body as any)?.[kContentLength];
+    if (contentLength !== undefined) {
+      this.headers.set("content-length", contentLength.toString());
+    }
+
     makeEnumerable(Response.prototype, this, enumerableResponseKeys);
     Object.defineProperty(this, kWaitUntil, nonEnumerable);
   }

packages/core/src/standards/index.ts

@@ -20,5 +20,6 @@ export {
   logResponse,
 } from "./http";
 export type { RequestInfo, RequestInit, ResponseInit, HRTime } from "./http";
-export * from "./streams";
+export { FixedLengthStream } from "./streams";
+export type { ArrayBufferViewConstructor } from "./streams";
 export * from "./timers";

packages/core/src/standards/streams.ts

@@ -1,7 +1,13 @@
 import {
   ReadableStreamBYOBReadResult,
   ReadableStreamBYOBReader,
+  TransformStream,
 } from "stream/web";
+import {
+  bufferSourceToArray,
+  buildNotBufferSourceError,
+  isBufferSource,
+} from "./helpers";
 
 export type ArrayBufferViewConstructor =
   | typeof Int8Array
@@ -70,3 +76,55 @@ ReadableStreamBYOBReader.prototype.readAtLeast = async function <
   const value = new ctor(buffer, byteOffset, read / bytesPerElement);
   return { value: value as any, done };
 };
+
+export const kContentLength = Symbol("kContentLength");
+
+export class FixedLengthStream extends TransformStream<Uint8Array, Uint8Array> {
+  constructor(expectedLength: number) {
+    // noinspection SuspiciousTypeOfGuard
+    if (typeof expectedLength !== "number" || expectedLength < 0) {
+      throw new TypeError(
+        "FixedLengthStream requires a non-negative integer expected length."
+      );
+    }
+
+    // Keep track of the number of bytes written
+    let written = 0;
+    super({
+      transform(chunk, controller) {
+        // Make sure this chunk is an ArrayBuffer(View)
+        if (isBufferSource(chunk)) {
+          const array = bufferSourceToArray(chunk);
+
+          // Throw if written too many bytes
+          written += array.byteLength;
+          if (written > expectedLength) {
+            return controller.error(
+              new TypeError(
+                "Attempt to write too many bytes through a FixedLengthStream."
+              )
+            );
+          }
+
+          controller.enqueue(array);
+        } else {
+          controller.error(new TypeError(buildNotBufferSourceError(chunk)));
+        }
+      },
+      flush(controller) {
+        // Throw if not written enough bytes on close
+        if (written < expectedLength) {
+          controller.error(
+            new TypeError(
+              "FixedLengthStream did not see all expected bytes before close()."
+            )
+          );
+        }
+      },
+    });
+
+    // When used as Request/Response body, override the Content-Length header
+    // with the expectedLength
+    (this.readable as any)[kContentLength] = expectedLength;
+  }
+}

packages/core/test/plugins/core.spec.ts

@@ -295,6 +295,7 @@ test("CorePlugin: setup: includes web standards", async (t) => {
   t.true(typeof globals.WritableStream === "function");
   t.true(typeof globals.WritableStreamDefaultController === "function");
   t.true(typeof globals.WritableStreamDefaultWriter === "function");
+  t.true(typeof globals.FixedLengthStream === "function");
 
   t.true(typeof globals.Event === "function");
   t.true(typeof globals.EventTarget === "function");

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

@@ -1,9 +1,13 @@
 import assert from "assert";
 import { ReadableStream, ReadableStreamBYOBReadResult } from "stream/web";
-import { ArrayBufferViewConstructor } from "@miniflare/core";
-import test from "ava";
-
-import "@miniflare/core";
+import {
+  ArrayBufferViewConstructor,
+  FixedLengthStream,
+  Request,
+  Response,
+} from "@miniflare/core";
+import { utf8Encode } from "@miniflare/shared-test";
+import test, { ThrowsExpectation } from "ava";
 
 function chunkedStream(chunks: number[][]): ReadableStream<Uint8Array> {
   return new ReadableStream({
@@ -176,3 +180,142 @@ test("ReadableStreamBYOBReader: readAtLeast: throws if minimum number of bytes e
     message: "Minimum bytes to read (4) exceeds size of buffer (3).",
   });
 });
+
+test("FixedLengthStream: requires non-negative integer expected length", (t) => {
+  const expectations: ThrowsExpectation = {
+    instanceOf: TypeError,
+    message:
+      "FixedLengthStream requires a non-negative integer expected length.",
+  };
+  // @ts-expect-error intentionally testing incorrect types
+  t.throws(() => new FixedLengthStream(), expectations);
+  t.throws(() => new FixedLengthStream(-42), expectations);
+  new FixedLengthStream(0);
+});
+test("FixedLengthStream: throws if too many bytes written", async (t) => {
+  const { readable, writable } = new FixedLengthStream(3);
+  const writer = writable.getWriter();
+  // noinspection ES6MissingAwait
+  void writer.write(new Uint8Array([1, 2]));
+  // noinspection ES6MissingAwait
+  void writer.write(new Uint8Array([3, 4]));
+
+  const reader = readable.getReader();
+  t.deepEqual((await reader.read()).value, new Uint8Array([1, 2]));
+  await t.throwsAsync(reader.read(), {
+    instanceOf: TypeError,
+    message: "Attempt to write too many bytes through a FixedLengthStream.",
+  });
+});
+test("FixedLengthStream: throws if too few bytes written", async (t) => {
+  const { readable, writable } = new FixedLengthStream(3);
+  const writer = writable.getWriter();
+  // noinspection ES6MissingAwait
+  void writer.write(new Uint8Array([1, 2]));
+  // noinspection ES6MissingAwait
+  const closePromise = writer.close();
+
+  const reader = readable.getReader();
+  t.deepEqual((await reader.read()).value, new Uint8Array([1, 2]));
+  await t.throwsAsync(closePromise, {
+    instanceOf: TypeError,
+    message: "FixedLengthStream did not see all expected bytes before close().",
+  });
+});
+test("FixedLengthStream: behaves as identity transform if just right number of bytes written", async (t) => {
+  const { readable, writable } = new FixedLengthStream(3);
+  const writer = writable.getWriter();
+  // noinspection ES6MissingAwait
+  void writer.write(new Uint8Array([1, 2]));
+  // noinspection ES6MissingAwait
+  void writer.write(new Uint8Array([3]));
+  // noinspection ES6MissingAwait
+  void writer.close();
+
+  const reader = readable.getReader();
+  t.deepEqual((await reader.read()).value, new Uint8Array([1, 2]));
+  t.deepEqual((await reader.read()).value, new Uint8Array([3]));
+  t.true((await reader.read()).done);
+});
+test("FixedLengthStream: throws on string chunks", async (t) => {
+  const { readable, writable } = new FixedLengthStream(5);
+  const writer = writable.getWriter();
+  // noinspection ES6MissingAwait
+  void writer.write(
+    // @ts-expect-error intentionally testing incorrect types
+    "how much chunk would a chunk-chuck chuck if a chunk-chuck could chuck chunk?"
+  );
+
+  const reader = readable.getReader();
+  await t.throwsAsync(reader.read(), {
+    instanceOf: TypeError,
+    message:
+      "This TransformStream is being used as a byte stream, " +
+      "but received a string on its writable side. " +
+      "If you wish to write a string, you'll probably want to " +
+      "explicitly UTF-8-encode it with TextEncoder.",
+  });
+});
+test("FixedLengthStream: throws on non-ArrayBuffer/ArrayBufferView chunks", async (t) => {
+  const { readable, writable } = new FixedLengthStream(5);
+  const writer = writable.getWriter();
+  // @ts-expect-error intentionally testing incorrect types
+  // noinspection ES6MissingAwait
+  void writer.write(42);
+
+  const reader = readable.getReader();
+  await t.throwsAsync(reader.read(), {
+    instanceOf: TypeError,
+    message:
+      "This TransformStream is being used as a byte stream, " +
+      "but received an object of non-ArrayBuffer/ArrayBufferView " +
+      "type on its writable side.",
+  });
+});
+function buildFixedLengthReadableStream(length: number) {
+  const { readable, writable } = new FixedLengthStream(length);
+  const writer = writable.getWriter();
+  if (length > 0) void writer.write(utf8Encode("".padStart(length, "x")));
+  void writer.close();
+  return readable;
+}
+test("FixedLengthStream: sets Content-Length header on Request", async (t) => {
+  let body = buildFixedLengthReadableStream(3);
+  let req = new Request("http://localhost", { method: "POST", body });
+  t.is(req.headers.get("Content-Length"), "3");
+  t.is(await req.text(), "xxx");
+
+  // Check overrides existing Content-Length header
+  body = buildFixedLengthReadableStream(3);
+  req = new Request("http://localhost", {
+    method: "POST",
+    body,
+    headers: { "Content-Length": "2" },
+  });
+  t.is(req.headers.get("Content-Length"), "3");
+  t.is(await req.text(), "xxx");
+
+  // Check still includes header with 0 expected length
+  body = buildFixedLengthReadableStream(0);
+  req = new Request("http://localhost", { method: "POST", body });
+  t.is(req.headers.get("Content-Length"), "0");
+  t.is(await req.text(), "");
+});
+test("FixedLengthStream: sets Content-Length header on Response", async (t) => {
+  let body = buildFixedLengthReadableStream(3);
+  let res = new Response(body);
+  t.is(res.headers.get("Content-Length"), "3");
+  t.is(await res.text(), "xxx");
+
+  // Check overrides existing Content-Length header
+  body = buildFixedLengthReadableStream(3);
+  res = new Response(body, { headers: { "Content-Length": "2" } });
+  t.is(res.headers.get("Content-Length"), "3");
+  t.is(await res.text(), "xxx");
+
+  // Check still includes header with 0 expected length
+  body = buildFixedLengthReadableStream(0);
+  res = new Response(body);
+  t.is(res.headers.get("Content-Length"), "0");
+  t.is(await res.text(), "");
+});