transformer

Author: JakeChampion

src/html_rewriter_wrapper.ts

@@ -27,13 +27,13 @@ class HTMLRewriter {
     }
 
     let rewriter: RawHTMLRewriter;
-    const readable = new ReadableStream<Uint8Array>({
+
+    const transfomer = new TransformStream({
       start: async (controller) => {
+        await HTMLRewriter.initPromise;
         // Create a rewriter instance for this transformation that writes its
         // output to the transformed response's stream. Note that each
         // BaseHTMLRewriter can only be used once.
-        await HTMLRewriter.initPromise;
-        // console.log('creating rewriter')
         rewriter = new RawHTMLRewriter((chunk: Uint8Array) => {
           // enqueue will throw on empty chunks
           if (chunk.length !== 0) controller.enqueue(chunk);
@@ -42,31 +42,18 @@ class HTMLRewriter {
         for (const [selector, handlers] of this.elementHandlers) {
           rewriter.on(selector, handlers);
         }
-
-        // Pipe the response body to the rewriter
-        const reader = body.getReader();
-        try {
-          while (true) {
-            // console.log('reading')
-            const { done, value } = await reader.read();
-            if (done) break;
-            rewriter.write(value);
-          }
-
-          rewriter.end();
-        } catch (error) {
-          // rewriter.end()
-          controller.error(error);
-        } finally {
-          rewriter.free();
-          reader.releaseLock();
-          controller.close();
-        }
+      },
+      transform(chunk) {
+        rewriter.write(chunk);
+      },
+      flush() {
+        rewriter.end();
+        rewriter.free();
       },
     });
 
     // Return a response with the transformed body, copying over headers, etc
-    const res = new Response(readable, response);
+    const res = new Response(body.pipeThrough(transfomer), 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");