feat(streams): LimitedTransformStream & LimitedBytesTransformStream (#…

…2007)
denoland · Mar 7, 2022 · f2682c6 · f2682c6
1 parent 6a69da9
commit f2682c6
Show file tree

Hide file tree

Showing 3 changed files with 165 additions and 6 deletions.
diff --git a/streams/buffer.ts b/streams/buffer.ts
@@ -166,3 +166,69 @@ export class Buffer {
     this.#reslice(m);
   }
 }
+
+/** A TransformStream that will only read & enqueue `size` amount of bytes.
+ * This operation is chunk based and not BYOB based,
+ * and as such will read more than needed.
+ *
+ * if options.error is set, then instead of terminating the stream,
+ * an error will be thrown.
+ *
+ * ```ts
+ * import { LimitedBytesTransformStream } from "./buffer.ts";
+ * const res = await fetch("https://example.com");
+ * const parts = res.body!
+ *   .pipeThrough(new LimitedBytesTransformStream(512 * 1024));
+ * ```
+ */
+export class LimitedBytesTransformStream
+  extends TransformStream<Uint8Array, Uint8Array> {
+  #read = 0;
+  constructor(size: number, options: { error?: boolean } = {}) {
+    super({
+      transform: (chunk, controller) => {
+        if ((this.#read + chunk.byteLength) > size) {
+          if (options.error) {
+            throw new RangeError(`Exceeded byte size limit of '${size}'`);
+          } else {
+            controller.terminate();
+          }
+        } else {
+          this.#read += chunk.byteLength;
+          controller.enqueue(chunk);
+        }
+      },
+    });
+  }
+}
+
+/** A TransformStream that will only read & enqueue `size` amount of chunks.
+ *
+ * if options.error is set, then instead of terminating the stream,
+ * an error will be thrown.
+ *
+ * ```ts
+ * import { LimitedTransformStream } from "./buffer.ts";
+ * const res = await fetch("https://example.com");
+ * const parts = res.body!.pipeThrough(new LimitedTransformStream(50));
+ * ```
+ */
+export class LimitedTransformStream<T> extends TransformStream<T, T> {
+  #read = 0;
+  constructor(size: number, options: { error?: boolean } = {}) {
+    super({
+      transform: (chunk, controller) => {
+        if ((this.#read + 1) > size) {
+          if (options.error) {
+            throw new RangeError(`Exceeded chunk limit of '${size}'`);
+          } else {
+            controller.terminate();
+          }
+        } else {
+          this.#read++;
+          controller.enqueue(chunk);
+        }
+      },
+    });
+  }
+}
diff --git a/streams/buffer_test.ts b/streams/buffer_test.ts
@@ -1,9 +1,13 @@
 // Copyright 2018-2022 the Deno authors. All rights reserved. MIT license.
 
-import { assert, assertEquals } from "../testing/asserts.ts";
-import { Buffer } from "./buffer.ts";
+import { assert, assertEquals, assertRejects } from "../testing/asserts.ts";
+import {
+  Buffer,
+  LimitedBytesTransformStream,
+  LimitedTransformStream,
+} from "./buffer.ts";
 
-Deno.test("Buffer Write & Read", async function () {
+Deno.test("[streams] Buffer Write & Read", async function () {
   const buf = new Buffer();
   const writer = buf.writable.getWriter();
   const reader = buf.readable.getReader({ mode: "byob" });
@@ -13,26 +17,114 @@ Deno.test("Buffer Write & Read", async function () {
   assertEquals(read.value, data);
 });
 
-Deno.test("Buffer Read empty", async function () {
+Deno.test("[streams] Buffer Read empty", async function () {
   const buf = new Buffer();
   const reader = buf.readable.getReader({ mode: "byob" });
   const read = await reader.read(new Uint8Array(5));
   assert(read.done);
   assertEquals(read.value!.byteLength, 0);
 });
 
-Deno.test("Buffer Write & get bytes", async function () {
+Deno.test("[streams] Buffer Write & get bytes", async function () {
   const buf = new Buffer();
   const writer = buf.writable.getWriter();
   const data = new Uint8Array([4, 21, 45, 19]);
   await writer.write(data);
   assertEquals(buf.bytes(), data);
 });
 
-Deno.test("Buffer truncate", async function () {
+Deno.test("[streams] Buffer truncate", async function () {
   const buf = new Buffer();
   const writer = buf.writable.getWriter();
   await writer.write(new Uint8Array([4, 21, 45, 19]));
   buf.truncate(3);
   assertEquals(buf.bytes(), new Uint8Array([4, 21, 45]));
 });
+
+Deno.test("[streams] LimitedBytesTransformStream", async function () {
+  const r = new ReadableStream({
+    start(controller) {
+      controller.enqueue(new Uint8Array([1, 2, 3]));
+      controller.enqueue(new Uint8Array([4, 5, 6]));
+      controller.enqueue(new Uint8Array([7, 8, 9]));
+      controller.enqueue(new Uint8Array([10, 11, 12]));
+      controller.enqueue(new Uint8Array([13, 14, 15]));
+      controller.enqueue(new Uint8Array([16, 17, 18]));
+      controller.close();
+    },
+  });
+
+  const chunks = [];
+  for await (const chunk of r.pipeThrough(new LimitedBytesTransformStream(7))) {
+    chunks.push(chunk);
+  }
+  assertEquals(chunks.length, 2);
+});
+
+Deno.test("[streams] LimitedBytesTransformStream error", async function () {
+  const r = new ReadableStream({
+    start(controller) {
+      controller.enqueue(new Uint8Array([1, 2, 3]));
+      controller.enqueue(new Uint8Array([4, 5, 6]));
+      controller.enqueue(new Uint8Array([7, 8, 9]));
+      controller.enqueue(new Uint8Array([10, 11, 12]));
+      controller.enqueue(new Uint8Array([13, 14, 15]));
+      controller.enqueue(new Uint8Array([16, 17, 18]));
+      controller.close();
+    },
+  });
+
+  await assertRejects(async () => {
+    for await (
+      const _chunk of r.pipeThrough(
+        new LimitedBytesTransformStream(7, { error: true }),
+      )
+    ) {
+      // needed to read
+    }
+  }, RangeError);
+});
+
+Deno.test("[streams] LimitedTransformStream", async function () {
+  const r = new ReadableStream({
+    start(controller) {
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.close();
+    },
+  });
+
+  const chunks = [];
+  for await (const chunk of r.pipeThrough(new LimitedTransformStream(3))) {
+    chunks.push(chunk);
+  }
+  assertEquals(chunks.length, 3);
+});
+
+Deno.test("[streams] LimitedTransformStream error", async function () {
+  const r = new ReadableStream({
+    start(controller) {
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.enqueue("foo");
+      controller.close();
+    },
+  });
+
+  await assertRejects(async () => {
+    for await (
+      const _chunk of r.pipeThrough(
+        new LimitedTransformStream(3, { error: true }),
+      )
+    ) {
+      // needed to read
+    }
+  }, RangeError);
+});
diff --git a/streams/mod.ts b/streams/mod.ts
@@ -1,5 +1,6 @@
 // Copyright 2018-2022 the Deno authors. All rights reserved. MIT license.
 
+export * from "./buffer.ts";
 export * from "./conversion.ts";
 export * from "./delimiter.ts";
 export * from "./merge.ts";