Improve conformance of HTMLRewriter to runtime

mrbbot · mrbbot · commit 39cc299edced · 2021-11-01T20:24:58.000Z
- Require `ArrayBuffer`/`ArrayBufferView` chunks
- Don't start rewriting until transformed `Response` body read
- Don't run `end()` handlers on `null` bodied `Response`s
diff --git a/packages/html-rewriter/src/rewriter.ts b/packages/html-rewriter/src/rewriter.ts
@@ -1,36 +1,25 @@
-import { ReadableStream } from "stream/web";
-import { TextEncoder } from "util";
+import { ReadableStream, TransformStream } from "stream/web";
 import { Response } from "@miniflare/core";
-import type { DocumentHandlers, ElementHandlers } from "html-rewriter-wasm";
+import type {
+  HTMLRewriter as BaseHTMLRewriter,
+  DocumentHandlers,
+  ElementHandlers,
+} from "html-rewriter-wasm";
 import { Response as BaseResponse } from "undici";
 
-// TODO: remove this function and these conversions, Workers only allows byte streams for Responses
-// Based on https://developer.mozilla.org/en-US/docs/Web/API/TransformStream#anything-to-uint8array_stream
-const encoder = new TextEncoder();
-// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
-function transformToArray(chunk: any): Uint8Array {
+function transformToArray(chunk: ArrayBuffer | ArrayBufferView): Uint8Array {
   if (chunk instanceof Uint8Array) {
     return chunk;
-  } else if (ArrayBuffer.isView(chunk)) {
-    return new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength);
   } else if (chunk instanceof ArrayBuffer) {
     return new Uint8Array(chunk);
-  } else if (
-    Array.isArray(chunk) &&
-    chunk.every((value) => typeof value === "number")
-  ) {
-    return new Uint8Array(chunk);
-  } else if (typeof chunk === "number") {
-    return new Uint8Array([chunk]);
-  } else if (chunk === null || chunk === undefined) {
-    throw new TypeError("Chunk must be defined");
   } else {
-    return encoder.encode(String(chunk));
+    return new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength);
   }
 }
 
 type SelectorElementHandlers = [selector: string, handlers: ElementHandlers];
 
+// noinspection SuspiciousTypeOfGuard
 export class HTMLRewriter {
   readonly #elementHandlers: SelectorElementHandlers[] = [];
   readonly #documentHandlers: DocumentHandlers[] = [];
@@ -46,8 +35,11 @@ export class HTMLRewriter {
   }
 
   transform(response: BaseResponse | Response): Response {
-    const transformedStream = new ReadableStream<Uint8Array>({
-      type: "bytes",
+    let rewriter: BaseHTMLRewriter;
+    const transformStream = new TransformStream<
+      ArrayBuffer | ArrayBufferView,
+      Uint8Array
+    >({
       start: async (controller) => {
         // Create a rewriter instance for this transformation that writes its
         // output to the transformed response's stream. Note that each
@@ -57,7 +49,7 @@ export class HTMLRewriter {
         const { HTMLRewriter: BaseHTMLRewriter } = await import(
           "html-rewriter-wasm"
         );
-        const rewriter = new BaseHTMLRewriter((output: Uint8Array) => {
+        rewriter = new BaseHTMLRewriter((output) => {
           // enqueue will throw on empty chunks
           if (output.length !== 0) controller.enqueue(output);
         });
@@ -68,29 +60,52 @@ export class HTMLRewriter {
         for (const handlers of this.#documentHandlers) {
           rewriter.onDocument(handlers);
         }
-
-        try {
-          // Transform the response body (may be null if empty)
-          if (response.body) {
-            for await (const chunk of response.body) {
-              await rewriter.write(transformToArray(chunk));
-            }
+      },
+      transform: async (chunk) => {
+        if (chunk instanceof ArrayBuffer || ArrayBuffer.isView(chunk)) {
+          try {
+            // Make sure we're passing a Uint8Array to Rust glue
+            return await rewriter.write(transformToArray(chunk));
+          } catch (e) {
+            // Make sure the rewriter is always freed (if transformer
+            // transform() throws an error, transformer flush() won't be called)
+            rewriter.free();
+            throw e;
           }
-          await rewriter.end();
-        } catch (e) {
-          controller.error(e);
+        } else if (typeof chunk === "string") {
+          throw new TypeError(
+            "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."
+          );
+        } else {
+          throw new TypeError(
+            "This TransformStream is being used as a byte stream, " +
+              "but received an object of non-ArrayBuffer/ArrayBufferView " +
+              "type on its writable side."
+          );
+        }
+      },
+      flush: async () => {
+        try {
+          // Runs document end handlers
+          return await rewriter.end();
         } finally {
-          // Make sure the rewriter/controller are always freed/closed
+          // Make sure the rewriter is always freed, regardless of whether
+          // rewriter.end() throws
           rewriter.free();
-          controller.close();
         }
       },
     });
 
     // Return a response with the transformed body, copying over headers, etc,
     // returning a @miniflare/core Response so we don't need to convert
-    // BaseResponse to one when dispatching fetch events
-    const res = new Response(transformedStream, response);
+    // BaseResponse to one when dispatching fetch events.
+    const body = response.body as ReadableStream<Uint8Array> | null;
+    const res = new Response(body?.pipeThrough(transformStream), response);
+    // If Content-Length is set, it's probably going to be wrong, since we're
+    // rewriting content, so remove it
     res.headers.delete("Content-Length");
     return res;
   }
diff --git a/packages/html-rewriter/test/rewriter.spec.ts b/packages/html-rewriter/test/rewriter.spec.ts
@@ -1,7 +1,6 @@
 import assert from "assert";
 import { ReadableStream } from "stream/web";
 import { setTimeout } from "timers/promises";
-import { URLSearchParams } from "url";
 import { TextDecoder, TextEncoder } from "util";
 import { Response } from "@miniflare/core";
 import { HTMLRewriter, HTMLRewriterPlugin } from "@miniflare/html-rewriter";
@@ -20,6 +19,8 @@ import {
 
 // TODO (someday): debug why removing .serial breaks some of these async tests
 
+const encoder = new TextEncoder();
+
 // region: ELEMENT HANDLERS
 
 const mutationsMacro: Macro<
@@ -556,7 +557,7 @@ test("HTMLRewriter: handles streaming responses", async (t) => {
         "</html>",
       ];
       for (const chunk of chunks) {
-        controller.enqueue(chunk);
+        controller.enqueue(encoder.encode(chunk));
         await setTimeout(50);
       }
       controller.close();
@@ -583,43 +584,106 @@ test("HTMLRewriter: handles streaming responses", async (t) => {
   t.true(chunks.length >= 2);
   t.is(chunks.join(""), '<html lang="en"><body><p>test</p></body></html>');
 });
-test("HTMLRewriter: handles streaming response with various chunk types", async (t) => {
-  const encoder = new TextEncoder();
-  const chunks: [input: any, output: Uint8Array][] = [
-    [new Uint8Array([1, 2, 3]), new Uint8Array([1, 2, 3])],
-    [new Int8Array([1, 2, 3]), new Uint8Array([1, 2, 3])],
-    [encoder.encode("test").buffer, encoder.encode("test")],
-    [[1, 2, 3], new Uint8Array([1, 2, 3])],
-    [1, new Uint8Array([1])],
-    ["test", encoder.encode("test")],
-    [
-      new URLSearchParams({ a: "1", b: "2", c: "3" }),
-      encoder.encode("a=1&b=2&c=3"),
-    ],
-    [{}, encoder.encode("[object Object]")],
-  ];
-
+test("HTMLRewriter: handles ArrayBuffer and ArrayBufferView chunks", async (t) => {
+  t.plan(2);
   const inputStream = new ReadableStream({
-    async start(controller) {
-      for (const [input] of chunks) controller.enqueue(input);
+    start(controller) {
+      const buffer = encoder.encode("<p>").buffer;
+      let array = encoder.encode("test");
+      const view1 = new Uint16Array(
+        array.buffer,
+        array.byteOffset,
+        array.byteLength / Uint16Array.BYTES_PER_ELEMENT
+      );
+      array = encoder.encode("</p>");
+      const view2 = new DataView(
+        array.buffer,
+        array.byteOffset,
+        array.byteLength
+      );
+      controller.enqueue(buffer);
+      controller.enqueue(view1);
+      controller.enqueue(view2);
+      controller.close();
+    },
+  });
+  const res = new HTMLRewriter()
+    .on("p", {
+      text(text) {
+        if (text.text) t.is(text.text, "test");
+      },
+    })
+    .transform(new Response(inputStream));
+  t.is(await res.text(), "<p>test</p>");
+});
+test("HTMLRewriter: throws on string chunks", async (t) => {
+  const inputStream = new ReadableStream({
+    start(controller) {
+      controller.enqueue("I'm a string");
       controller.close();
     },
   });
   const res = new HTMLRewriter().transform(new Response(inputStream));
-  assert(res.body);
-  for await (const chunk of res.body) {
-    const expected = chunks.shift();
-    t.deepEqual(chunk, expected?.[1]);
-  }
+  await t.throwsAsync(res.text(), {
+    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("HTMLRewriter: throws on non-ArrayBuffer/ArrayBufferView chunks", async (t) => {
+  const inputStream = new ReadableStream({
+    start(controller) {
+      controller.enqueue(42);
+      controller.close();
+    },
+  });
+  const res = new HTMLRewriter().transform(new Response(inputStream));
+  await t.throwsAsync(res.text(), {
+    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.",
+  });
 });
 test("HTMLRewriter: handles empty response", async (t) => {
   // Shouldn't call BaseHTMLRewriter.write, just BaseHTMLRewriter.end
-  const res = new HTMLRewriter().transform(new Response());
+  const res = new HTMLRewriter()
+    .onDocument({
+      end(end) {
+        end.append("end");
+      },
+    })
+    .transform(new Response());
+  // Workers don't run the end() handler on null responses
   t.is(await res.text(), "");
 });
 test("HTMLRewriter: handles empty string response", async (t) => {
-  const res = new HTMLRewriter().transform(new Response(""));
-  t.is(await res.text(), "");
+  const res = new HTMLRewriter()
+    .onDocument({
+      end(end) {
+        end.append("end");
+      },
+    })
+    .transform(new Response(""));
+  t.is(await res.text(), "end");
+});
+test("HTNLRewriter: doesn't transform response until needed", async (t) => {
+  const chunks: string[] = [];
+  const res = new HTMLRewriter()
+    .on("p", {
+      text(text) {
+        if (text.text) chunks.push(text.text);
+      },
+    })
+    .transform(new Response("<p>1</p><p>2</p><p>3</p>"));
+  await setTimeout(50);
+  t.deepEqual(chunks, []);
+  await res.arrayBuffer();
+  t.deepEqual(chunks, ["1", "2", "3"]);
 });
 test("HTMLRewriter: copies response status and headers", async (t) => {
   const res = new HTMLRewriter().transform(
