HTMLRewriter: cancel a still-open JS stream input when the rewrite fails or its output is cancelled

src/runtime/api/html_rewriter.rs

@@ -6,8 +6,9 @@ use core::ptr::NonNull;
 use std::rc::Rc;
 
 use bun_jsc::{
-    self as jsc, CallFrame, GlobalRef, JSGlobalObject, JSPromise, JSValue, JsCell, JsResult,
-    ProtectedJSValue, SystemError, bun_string_jsc,
+    self as jsc, CallFrame, CommonAbortReason, CommonAbortReasonExt as _, GlobalRef,
+    JSGlobalObject, JSPromise, JSValue, JsCell, JsResult, ProtectedJSValue, SystemError,
+    bun_string_jsc,
 };
 // Note: `bun_jsc::VirtualMachine` is a *module* re-export
 // (`pub use self::virtual_machine as VirtualMachine;`). The struct lives at
@@ -845,19 +846,25 @@ impl RewriterPipe {
     }
 
     /// Sever the wired input source: null the upstream's raw `sink` backref
-    /// so it can no longer dispatch into this pipe. With `cancel_upstream`,
-    /// also close a native producer afterwards, so a failed or cancelled
-    /// rewrite stops a fetch mid-download and closes a file fd instead of
-    /// draining to upstream EOF; EOF paths pass `false`. Only called from
-    /// terminal paths on the JS thread. Idempotent: the handle is `None`
-    /// after the first call. Returns the severed handle for
+    /// so it can no longer dispatch into this pipe. With a `cancel` reason
+    /// (a failed or cancelled rewrite; EOF paths pass `None`) the upstream
+    /// is also told its consumer went away: a native producer is closed, so
+    /// a fetch stops mid-download and a file fd closes instead of draining
+    /// to EOF, and a JS stream's controller is closed before the detach, so
+    /// the source's `cancel(reason)` runs. Only called from terminal paths
+    /// on the JS thread, after `done` is set (closing a JS controller ends
+    /// the sink again through `end_from_js`). Idempotent: the handle is
+    /// `None` after the first call. Returns the severed handle for
     /// [`Self::release_input_roots`], which the caller runs after its
     /// terminal work.
-    fn detach_input_source(&self, cancel_upstream: bool) -> SourceHandle {
+    fn detach_input_source(&self, cancel: Option<JSValue>) -> SourceHandle {
         let mut src = self.input_source.replace(SourceHandle::None);
         let mut upstream = src;
+        if let (Some(reason), SourceHandle::JSController(_)) = (cancel, upstream) {
+            upstream.cancel(reason);
+        }
         JSSink::<RewriterPipe>::detach(&mut src, &self.global);
-        if cancel_upstream {
+        if cancel.is_some() {
             match upstream {
                 SourceHandle::ByteStream(_) | SourceHandle::FileReader(_) => {
                     upstream.close(None);
@@ -1420,7 +1427,7 @@ impl RewriterPipe {
         // pipe; otherwise the controller's destructor would later dispatch
         // `__controllerDetached`/`__finalize` on freed memory. The upstream
         // already ended, so there is nothing to cancel.
-        let src = self.detach_input_source(false);
+        let src = self.detach_input_source(None);
 
         if self.js_pump_reaction_pending.get() {
             // The pump closes the sink before its promise settles; the `.then()`
@@ -1465,9 +1472,12 @@ impl RewriterPipe {
     pub fn cancel_from_output(&self, _err: Option<SysError>) {
         let _pin = self.pin();
         self.detach_output();
-        let src = self.detach_input_source(true);
         self.phase.set(RewritePhase::Done);
         self.done.set(true);
+        // The reader's cancel reason does not travel through the output
+        // `ByteStream`; the input's `cancel()` gets the same `AbortError` the
+        // output hands a native sink wired to it.
+        let src = self.detach_input_source(Some(CommonAbortReason::UserAbort.to_js(&self.global)));
         self.pending.with_mut(|p| {
             p.result = Writable::Done;
             p.run();
@@ -1714,11 +1724,11 @@ impl RewriterPipe {
     }
 
     /// Put `err` on the output `Response`'s body / ByteStream.
-    fn fail(&self, err: webcore::body::ValueError) {
+    fn fail(&self, mut err: webcore::body::ValueError) {
         let _pin = self.pin();
         self.phase.set(RewritePhase::Done);
         self.done.set(true);
-        let src = self.detach_input_source(true);
+        let src = self.detach_input_source(Some(err.to_js(&self.global)));
         // Settle any `flush(true)`/`write()` promise a direct-stream `pull()`
         // is parked on so the pump promise can settle (mirrors
         // `cancel_from_output`).
@@ -1730,7 +1740,6 @@ impl RewriterPipe {
         if let Some(out) = self.output.get() {
             // Output emitted before the failure still precedes the error.
             self.flush_output();
-            let mut err = err;
             out.on_data(StreamResult::Err(err.to_stream_error(&self.global)));
             self.detach_output();
         } else if let Some(response) = self.response.get().as_deref() {

src/runtime/webcore/streams.rs

@@ -920,13 +920,27 @@ impl SourceHandle {
 
     /// [`close`](Self::close), with `reason` as the piped JS stream's cancel() argument.
     pub fn abort(&mut self, reason: CommonAbortReason) {
+        match *self {
+            SourceHandle::JSController(_) => {
+                let global = VirtualMachine::get().global();
+                if global.has_exception() {
+                    return;
+                }
+                self.cancel(reason.to_js(global));
+            }
+            _ => self.close(None),
+        }
+    }
+
+    /// [`close`](Self::close), with an arbitrary JS `reason` for the piped JS stream's cancel().
+    pub fn cancel(&mut self, reason: JSValue) {
         match *self {
             SourceHandle::JSController(cpp) => {
                 let global = VirtualMachine::get().global();
                 if global.has_exception() {
                     return;
                 }
-                Self::close_js_controller(global, cpp, reason.to_js(global));
+                Self::close_js_controller(global, cpp, reason);
             }
             _ => self.close(None),
         }

test/js/workerd/html-rewriter.test.js

@@ -13,7 +13,7 @@ import {
   tls,
   tmpdirSync,
 } from "harness";
-import { createServer as createTcpServer } from "net";
+import { createConnection, createServer as createTcpServer } from "net";
 import path, { join } from "path";
 import { setImmediate as setImmediatePromise } from "timers/promises";
 var setTimeoutAsync = (fn, delay) => {
@@ -600,6 +600,113 @@ describe("HTMLRewriter", () => {
       expect(afterFlush).toBe(1);
     });
 
+    // The rewrite's consumer going away while a `type: 'direct'` input is
+    // still open is a cancel of that input: its `cancel(reason)` runs once,
+    // and a later write() reports 0 bytes instead of throwing. A detach alone
+    // (what the source's own close() does) no longer calls cancel(), so the
+    // failed or cancelled pipe has to close the controller itself.
+    describe("a still-open direct ReadableStream input is cancelled when", () => {
+      function directInput(events) {
+        const hold = Promise.withResolvers();
+        const settled = Promise.withResolvers();
+        const stream = new ReadableStream({
+          type: "direct",
+          async pull(controller) {
+            try {
+              controller.write("<p>one</p>");
+              controller.flush();
+              await hold.promise;
+              events.push("write after: " + controller.write("<p>two</p>"));
+            } catch (e) {
+              events.push("pull threw: " + e.message);
+            } finally {
+              settled.resolve();
+            }
+          },
+          cancel(reason) {
+            events.push("cancel: " + (reason instanceof Error ? `${reason.name}: ${reason.message}` : reason));
+          },
+        });
+        return { response: new Response(stream), hold, settled: settled.promise };
+      }
+
+      it("the output reader cancels", async () => {
+        const events = [];
+        const { response, hold, settled } = directInput(events);
+        const res = new HTMLRewriter().on("p", { element() {} }).transform(response);
+        const reader = res.body.getReader();
+        const first = await reader.read();
+        expect(new TextDecoder().decode(first.value)).toBe("<p>one</p>");
+        await reader.cancel(new Error("reader went away"));
+        hold.resolve();
+        await settled;
+        expect(events).toEqual(["cancel: AbortError: The operation was aborted.", "write after: 0"]);
+      });
+
+      it("a handler throws", async () => {
+        const events = [];
+        const { response, hold, settled } = directInput(events);
+        const res = new HTMLRewriter()
+          .on("p", {
+            element() {
+              throw new Error("handler threw");
+            },
+          })
+          .transform(response);
+        await expect(res.text()).rejects.toThrow("handler threw");
+        hold.resolve();
+        await settled;
+        expect(events).toEqual(["cancel: Error: handler threw", "write after: 0"]);
+      });
+
+      it("the HTTP client disconnects from a Bun.serve response", async () => {
+        const events = [];
+        const flushed = Promise.withResolvers();
+        let settled;
+        using server = Bun.serve({
+          port: 0,
+          idleTimeout: 0,
+          fetch() {
+            const input = directInput(events);
+            settled = input.settled;
+            flushed.resolve(input.hold);
+            return new HTMLRewriter().on("p", { element() {} }).transform(input.response);
+          },
+        });
+        const socket = createConnection(server.port, "127.0.0.1");
+        await once(socket, "connect");
+        socket.write("GET / HTTP/1.1\r\nHost: localhost\r\n\r\n");
+        await once(socket, "data");
+        const hold = await flushed.promise;
+        socket.destroy();
+        while (server.pendingRequests > 0) await setImmediatePromise();
+        hold.resolve();
+        await settled;
+        expect(events).toEqual(["cancel: AbortError: The operation was aborted.", "write after: 0"]);
+      });
+
+      // A default ReadableStream input was already cancelled here (the pump
+      // cancels its stream on any close), but with `undefined`; it gets the
+      // same reason now.
+      it("the output reader cancels (default ReadableStream input gets the reason too)", async () => {
+        const cancelled = [];
+        const stream = new ReadableStream({
+          start(controller) {
+            controller.enqueue("<p>one</p>");
+          },
+          cancel(reason) {
+            cancelled.push(reason instanceof Error ? `${reason.name}: ${reason.message}` : reason);
+          },
+        });
+        const res = new HTMLRewriter().on("p", { element() {} }).transform(new Response(stream));
+        const reader = res.body.getReader();
+        const first = await reader.read();
+        expect(new TextDecoder().decode(first.value)).toBe("<p>one</p>");
+        await reader.cancel();
+        expect(cancelled).toEqual(["AbortError: The operation was aborted."]);
+      });
+    });
+
     // Two rewriters chained, both suspending: `init()`'s own comment names
     // "another transform()" as a supported consumer of a pending body.
     it("chained suspending transforms", async () => {