From dd0cfb82c69cac103811fc7273b07d595a16fd9f Mon Sep 17 00:00:00 2001
From: jake champion <jake.champion@netlify.com>
Date: Mon, 20 Jan 2025 08:51:44 +0000
Subject: [PATCH] fix: free the underlying HTMLRewriter if the TransformStream
 has been cancelled

---
 src/html_rewriter_wrapper.ts | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/html_rewriter_wrapper.ts b/src/html_rewriter_wrapper.ts
index 00164a2..67e5450 100644
--- a/src/html_rewriter_wrapper.ts
+++ b/src/html_rewriter_wrapper.ts
@@ -24,7 +24,7 @@ export class HTMLRewriter {
       response = new Response(response.body, response);
     }
 
-    let rewriter: RawHTMLRewriter;
+    let rewriter: RawHTMLRewriter | undefined;
     const transformStream = new TransformStream<Uint8Array, Uint8Array>({
       start: (controller) => {
         // Create a rewriter instance for this transformation that writes its
@@ -39,10 +39,21 @@ export class HTMLRewriter {
           rewriter.on(selector, handlers);
         }
       },
-      transform: (chunk) => rewriter.write(chunk),
+      transform: (chunk) => rewriter?.write(chunk),
+      // deno-lint-ignore require-await
+      async cancel() {
+        if (rewriter) {
+          rewriter.end();
+          rewriter.free();
+          rewriter = undefined;
+        }
+      },
       flush: () => {
-        rewriter.end();
-        rewriter.free();
+        if (rewriter) {
+          rewriter.end();
+          rewriter.free();
+          rewriter = undefined;
+        }
       },
     });