Switch to html-rewriter-wasm package for HTMLRewriter

Author: mrbbot

package-lock.json

@@ -45,6 +45,7 @@
     "dotenv": "^8.2.0",
     "env-paths": "^2.2.1",
     "formdata-node": "^2.4.0",
+    "html-rewriter-wasm": "^0.3.0",
     "http-cache-semantics": "^4.1.0",
     "ioredis": "^4.27.6",
     "micromatch": "^4.0.4",

package.json

@@ -1,6 +1,14 @@
 export { DurableObject } from "./do";
 export { FetchEvent, ScheduledEvent, ResponseWaitUntil } from "./events";
-export { HTMLRewriter, UnsafeHTMLRewriter } from "./rewriter";
+export {
+  HTMLRewriter,
+  Element,
+  Comment,
+  TextChunk,
+  Doctype,
+  DocumentEnd,
+  ContentTypeOptions,
+} from "./rewriter";
 export {
   URL,
   URLSearchParams,

src/modules/index.ts

@@ -1,30 +1,19 @@
-/// <reference types="@cloudflare/workers-types" />
 import { TextEncoder } from "util";
 import { Response } from "@mrbbot/node-fetch";
-import { ReadableStream } from "web-streams-polyfill/ponyfill/es6";
-// This import relies on dist having the same structure as src
 import {
-  HTMLRewriter as LOLHTMLRewriter,
-  registerPromise,
-} from "../../vendor/lol-html";
-import { Mutex } from "../kv/helpers";
-import { ProcessedOptions } from "../options";
+  HTMLRewriter as BaseHTMLRewriter,
+  Comment,
+  ContentTypeOptions,
+  Doctype,
+  DocumentEnd,
+  DocumentHandlers,
+  Element,
+  ElementHandlers,
+  TextChunk,
+} from "html-rewriter-wasm";
+import { ReadableStream } from "web-streams-polyfill/ponyfill/es6";
 import { Context, Module } from "./module";
 
-function wrapHandler<T>(handler?: (arg: T) => void | Promise<void>) {
-  if (handler === undefined) return undefined;
-  return function (arg: T) {
-    const result = handler(arg);
-    // If this handler is async and returns a promise, register it and return
-    // its handle so it can be awaited later in WebAssembly
-    if (typeof result === "object" && typeof result.then === "function") {
-      return registerPromise(result);
-    }
-    // Otherwise, return 0 to signal there's nothing to await
-    return 0;
-  };
-}
-
 // 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
@@ -49,54 +38,28 @@ export function transformToArray(chunk: any): Uint8Array {
   }
 }
 
-/* The WebAssembly version of lol-html used by Miniflare uses asyncify for
- * async handlers. When a handler returns a promise, the WebAssembly stack is
- * stored in temporary storage, the promise is awaited, then the stack is
- * restored and WebAssembly execution continues where it left off. This
- * temporary storage is currently per module instance and we only have a single
- * instance because of how wasm-pack generates package code for NodeJS.
- * TODO: ideally, we would allocate each transform call its own temporary
- *  space for the saved stack.
- *
- * This means if you have multiple concurrent transforms in progress, the saved
- * stacks will be overwritten and lol-html will be unhappy. Therefore, to be
- * "safe", we need to make sure only one transform operation is in-progress at
- * any time, hence the mutex.
- *
- * However, this problem only occurs when using async handlers with concurrent
- * transforms. If just using sync handlers, or not doing multiple rewrites
- * concurrently (very likely), there's no need for the mutex, so we can use the
- * "unsafe" version. The "safe" version is the default just so people don't see
- * confusing errors. See the docs for concrete examples of where the "unsafe"
- * version can be used.
- */
-
-const wasmModuleMutex = new Mutex();
-// Symbol gives us "protected" method only accessible/overridable by subclass
-const runCriticalSectionSymbol = Symbol("HTMLRewriter runCriticalSection");
-
-export class UnsafeHTMLRewriter {
+export class HTMLRewriter {
   #elementHandlers: [selector: string, handlers: any][] = [];
   #documentHandlers: any[] = [];
 
-  on(selector: string, handlers: Partial<ElementHandler>): this {
+  on(selector: string, handlers: ElementHandlers): this {
     // Ensure handlers register returned promises, and `this` is bound correctly
     const wrappedHandlers = {
-      element: wrapHandler(handlers.element?.bind(handlers)),
-      comments: wrapHandler(handlers.comments?.bind(handlers)),
-      text: wrapHandler(handlers.text?.bind(handlers)),
+      element: handlers.element?.bind(handlers),
+      comments: handlers.comments?.bind(handlers),
+      text: handlers.text?.bind(handlers),
     };
     this.#elementHandlers.push([selector, wrappedHandlers]);
     return this;
   }
 
-  onDocument(handlers: Partial<DocumentHandler>): this {
+  onDocument(handlers: DocumentHandlers): this {
     // Ensure handlers register returned promises, and `this` is bound correctly
     const wrappedHandlers = {
-      doctype: wrapHandler(handlers.doctype?.bind(handlers)),
-      comments: wrapHandler(handlers.comments?.bind(handlers)),
-      text: wrapHandler(handlers.text?.bind(handlers)),
-      end: wrapHandler(handlers.end?.bind(handlers)),
+      doctype: handlers.doctype?.bind(handlers),
+      comments: handlers.comments?.bind(handlers),
+      text: handlers.text?.bind(handlers),
+      end: handlers.end?.bind(handlers),
     };
     this.#documentHandlers.push(wrappedHandlers);
     return this;
@@ -108,7 +71,7 @@ export class UnsafeHTMLRewriter {
       start: async (controller) => {
         // Create a rewriter instance for this transformation that writes its
         // output to the transformed response's stream
-        const rewriter = new LOLHTMLRewriter((output: Uint8Array) => {
+        const rewriter = new BaseHTMLRewriter((output: Uint8Array) => {
           if (output.length === 0) {
             // Free the rewriter once it's finished doing its thing
             queueMicrotask(() => rewriter.free());
@@ -125,41 +88,32 @@ export class UnsafeHTMLRewriter {
           rewriter.onDocument(handlers);
         }
 
-        await this[runCriticalSectionSymbol](async () => {
-          // 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 the response body (may be null if empty)
+        if (response.body) {
+          for await (const chunk of response.body) {
+            await rewriter.write(transformToArray(chunk));
           }
-          await rewriter.end();
-        });
+        }
+        await rewriter.end();
       },
     });
 
     // Return a response with the transformed body, copying over headers, etc
     return new Response(transformedStream, response);
   }
-
-  [runCriticalSectionSymbol](closure: () => Promise<void>): Promise<void> {
-    return closure();
-  }
-}
-
-// See big comment above for what this does and why it's needed. It's possible
-// we'll remove this distinction in the future.
-export class HTMLRewriter extends UnsafeHTMLRewriter {
-  [runCriticalSectionSymbol](closure: () => Promise<void>): Promise<void> {
-    return wasmModuleMutex.run(closure);
-  }
 }
 
 export class HTMLRewriterModule extends Module {
-  buildSandbox(options: ProcessedOptions): Context {
-    return {
-      HTMLRewriter: options.htmlRewriterUnsafe
-        ? UnsafeHTMLRewriter
-        : HTMLRewriter,
-    };
+  buildSandbox(): Context {
+    return { HTMLRewriter };
   }
 }
+
+export {
+  Element,
+  Comment,
+  TextChunk,
+  Doctype,
+  DocumentEnd,
+  ContentTypeOptions,
+};

src/modules/rewriter.ts

@@ -94,8 +94,6 @@ export interface Options {
   envPath?: string;
   bindings?: Record<string, any>;
   wasmBindings?: Record<string, string>;
-
-  htmlRewriterUnsafe?: boolean;
 }
 
 export interface ProcessedOptions extends Options {
@@ -158,7 +156,6 @@ export function logOptions(log: Log, options: ProcessedOptions): void {
       : options.https === true
       ? "Self-Signed"
       : `Self-Signed: ${options.https}`,
-    "Unsafe HTMLRewriter": options.htmlRewriterUnsafe,
   };
   log.debug("Options:");
   for (const [key, value] of Object.entries(entries)) {

src/options/index.ts

@@ -2,21 +2,24 @@ import { URLSearchParams } from "url";
 import { TextDecoder, TextEncoder } from "util";
 import test, { Macro, ThrowsExpectation } from "ava";
 import { ReadableStream } from "web-streams-polyfill/ponyfill/es6";
-// Import UnsafeHTMLRewriter as HTMLRewriter since we're using it much more
-// (we want to be able to run tests in parallel) and typing "Unsafe" each time
-// is annoying.
 import {
-  UnsafeHTMLRewriter as HTMLRewriter,
+  Comment,
+  Doctype,
+  DocumentEnd,
+  Element,
+  HTMLRewriter,
   NoOpLog,
   Response,
-  HTMLRewriter as SafeHTMLRewriter,
+  TextChunk,
 } from "../../src";
 import {
   HTMLRewriterModule,
   transformToArray,
 } from "../../src/modules/rewriter";
 import { getObjectProperties, wait } from "../helpers";
 
+// TODO: remove most of these tests, they're now in html-rewriter-wasm
+
 // region: Uint8ArrayTransformStream
 const encoder = new TextEncoder();
 test("transformToArray: passes through Uint8Array", (t) => {
@@ -63,7 +66,7 @@ const mutationsMacro: Macro<
   [
     (
       rw: HTMLRewriter,
-      handler: (token: Element | Text | Comment) => void
+      handler: (token: Element | TextChunk | Comment) => void
     ) => HTMLRewriter,
     string,
     {
@@ -338,7 +341,7 @@ const textMutationsExpected = {
 };
 
 const textPropertiesMacro: Macro<
-  [(rw: HTMLRewriter, text: (text: Text) => void) => HTMLRewriter]
+  [(rw: HTMLRewriter, text: (text: TextChunk) => void) => HTMLRewriter]
 > = async (t, func) => {
   t.plan(6);
   const res = func(new HTMLRewriter(), (text) => {
@@ -364,7 +367,7 @@ test(
   textMutationsExpected
 );
 const textAsyncHandlerMacro: Macro<
-  [(rw: HTMLRewriter, text: (t: Text) => Promise<void>) => HTMLRewriter]
+  [(rw: HTMLRewriter, text: (t: TextChunk) => Promise<void>) => HTMLRewriter]
 > = async (t, func) => {
   const res = func(new HTMLRewriter(), async (text) => {
     if (text.text === "t") {
@@ -380,11 +383,16 @@ test.serial(
   (rw, text) => rw.on("p", { text })
 );
 const textClassHandlerMacro: Macro<
-  [(rw: HTMLRewriter, handler: { text: (text: Text) => void }) => HTMLRewriter]
+  [
+    (
+      rw: HTMLRewriter,
+      handler: { text: (text: TextChunk) => void }
+    ) => HTMLRewriter
+  ]
 > = async (t, func) => {
   class Handler {
     constructor(private content: string) {}
-    text(text: Text) {
+    text(text: TextChunk) {
       if (text.text === "t") text.after(this.content);
     }
   }
@@ -615,7 +623,7 @@ test("HTMLRewriter: handles streaming responses", async (t) => {
   t.is(chunks.join(""), '<html lang="en"><body><p>test</p></body></html>');
 });
 test("HTMLRewriter: handles empty response", async (t) => {
-  // Shouldn't call LOLHTMLRewriter.write, just LOLHTMLRewriter.end
+  // Shouldn't call BaseHTMLRewriter.write, just BaseHTMLRewriter.end
   const res = new HTMLRewriter().transform(new Response());
   t.is(await res.text(), "");
 });
@@ -663,7 +671,7 @@ test.serial(
     // Note this test requires the "safe" HTMLRewriter, see comments in
     // src/modules/rewriter.ts for more details
     const rewriter = (i: number) =>
-      new SafeHTMLRewriter()
+      new HTMLRewriter()
         .on("p", {
           async element(element) {
             await wait(50);
@@ -858,17 +866,9 @@ test("HTMLRewriter: throws error on unsupported selector", async (t) => {
 
 // region: MODULE
 
-test("HTMLRewriter: defaults to safe implementation", async (t) => {
-  const module = new HTMLRewriterModule(new NoOpLog());
-  const { HTMLRewriter: SandboxHTMLRewriter } = module.buildSandbox({});
-  t.is(SandboxHTMLRewriter, SafeHTMLRewriter);
-});
-test("HTMLRewriter: allows unsafe implementation", async (t) => {
+test("HTMLRewriter: included in sandbox", async (t) => {
   const module = new HTMLRewriterModule(new NoOpLog());
-  const { HTMLRewriter: SandboxHTMLRewriter } = module.buildSandbox({
-    htmlRewriterUnsafe: true,
-  });
-  // UnsafeHTMLRewriter has been imported as HTMLRewriter
+  const { HTMLRewriter: SandboxHTMLRewriter } = module.buildSandbox();
   t.is(SandboxHTMLRewriter, HTMLRewriter);
 });
 

test/modules/rewriter.spec.ts

@@ -63,7 +63,6 @@ test("logOptions: logs all options", (t) => {
     durableObjectsPersist: true,
     bindings: { KEY: "value" },
     https: true,
-    htmlRewriterUnsafe: true,
   });
 
   t.deepEqual(log.debugs, [
@@ -84,7 +83,6 @@ test("logOptions: logs all options", (t) => {
     "- Durable Objects Persistence: true",
     "- Bindings: KEY",
     "- HTTPS: Self-Signed",
-    "- Unsafe HTMLRewriter: true",
   ]);
 });