From 44a9b03e8ebd5667d026ce3505d2b28adf946540 Mon Sep 17 00:00:00 2001
From: James M Snell <jasnell@gmail.com>
Date: Mon, 22 Apr 2024 16:08:08 -0700
Subject: [PATCH] Hold strong refs in certain streams Writer/Reader methods

In some edge cases, if the Writer/Reader is the only thing remaining
holding a strong reference to a WritableStream or ReadableStream,
the calls to `close()`, `abort()`, `cancel()`, etc can cause the
Writer/Reader to drop it's strong reference in mid operation, leading
to issues. Let's defensively hold an additional strong reference while
the operations are pending
---
 src/workerd/api/streams/readable.c++  | 11 +++++++++++
 src/workerd/api/streams/writable.c++  | 16 ++++++++++++++++
 src/workerd/api/tests/streams-test.js | 13 +++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/src/workerd/api/streams/readable.c++ b/src/workerd/api/streams/readable.c++
index 36a015e70fd2..96f59b91d910 100644
--- a/src/workerd/api/streams/readable.c++
+++ b/src/workerd/api/streams/readable.c++
@@ -21,6 +21,7 @@ ReaderImpl::~ReaderImpl() noexcept(false) {
     // There's a very good likelihood that this is called during GC or other
     // cleanup so we have to make sure that releasing the reader does not also
     // trigger resolution of the close promise.
+    auto ref = stream.addRef();
     stream->getController().releaseReader(reader, kj::none);
   }
 }
@@ -61,6 +62,11 @@ jsg::Promise<void> ReaderImpl::cancel(
       KJ_FAIL_ASSERT("this reader was never attached");
     }
     KJ_CASE_ONEOF(stream, Attached) {
+      // In some edge cases, this reader is the last thing holding a strong
+      // reference to the stream. Calling cancel might cause the readers strong
+      // reference to be cleared, so let's make sure we keep a reference to
+      // the stream at least until the call to cancel completes.
+      auto ref = stream.addRef();
       return stream->getController().cancel(js, maybeReason);
     }
     KJ_CASE_ONEOF(r, Released) {
@@ -141,6 +147,11 @@ void ReaderImpl::releaseLock(jsg::Lock& js) {
       KJ_FAIL_ASSERT("this reader was never attached");
     }
     KJ_CASE_ONEOF(stream, Attached) {
+      // In some edge cases, this reader is the last thing holding a strong
+      // reference to the stream. Calling releaseLock might cause the readers strong
+      // reference to be cleared, so let's make sure we keep a reference to
+      // the stream at least until the call to releaseLock completes.
+      auto ref = stream.addRef();
       stream->getController().releaseReader(reader, js);
       state.init<Released>();
       return;
diff --git a/src/workerd/api/streams/writable.c++ b/src/workerd/api/streams/writable.c++
index 38ed6ce56829..956f8a32dc98 100644
--- a/src/workerd/api/streams/writable.c++
+++ b/src/workerd/api/streams/writable.c++
@@ -17,6 +17,7 @@ WritableStreamDefaultWriter::~WritableStreamDefaultWriter() noexcept(false) {
     // Because this can be called during gc or other cleanup, it is important
     // that releasing the writer does not cause the closed promise be resolved
     // since that requires v8 heap allocations.
+    auto ref = stream.addRef();
     stream->getController().releaseWriter(*this, kj::none);
   }
 }
@@ -39,6 +40,11 @@ jsg::Promise<void> WritableStreamDefaultWriter::abort(
       KJ_FAIL_ASSERT("this writer was never attached");
     }
     KJ_CASE_ONEOF(stream, Attached) {
+      // In some edge cases, this writer is the last thing holding a strong
+      // reference to the stream. Calling abort can cause the writers strong
+      // reference to be cleared, so let's make sure we keep a reference to
+      // the stream at least until the call to abort completes.
+      auto ref = stream.addRef();
       return stream->getController().abort(js, reason);
     }
     KJ_CASE_ONEOF(r, Released) {
@@ -68,6 +74,11 @@ jsg::Promise<void> WritableStreamDefaultWriter::close(jsg::Lock& js) {
       KJ_FAIL_ASSERT("this writer was never attached");
     }
     KJ_CASE_ONEOF(stream, Attached) {
+      // In some edge cases, this writer is the last thing holding a strong
+      // reference to the stream. Calling close can cause the writers strong
+      // reference to be cleared, so let's make sure we keep a reference to
+      // the stream at least until the call to close completes.
+      auto ref = stream.addRef();
       return stream->getController().close(js);
     }
     KJ_CASE_ONEOF(r, Released) {
@@ -142,6 +153,11 @@ void WritableStreamDefaultWriter::releaseLock(jsg::Lock& js) {
       KJ_FAIL_ASSERT("this writer was never attached");
     }
     KJ_CASE_ONEOF(stream, Attached) {
+      // In some edge cases, this writer is the last thing holding a strong
+      // reference to the stream. Calling releaseWriter can cause the writers
+      // strong reference to be cleared, so let's make sure we keep a reference
+      // to the stream at least until the call to releaseLock completes.
+      auto ref = stream.addRef();
       stream->getController().releaseWriter(*this, js);
       state.init<Released>();
       return;
diff --git a/src/workerd/api/tests/streams-test.js b/src/workerd/api/tests/streams-test.js
index 4029a7ba1497..b08ad700dd8a 100644
--- a/src/workerd/api/tests/streams-test.js
+++ b/src/workerd/api/tests/streams-test.js
@@ -414,6 +414,19 @@ export const readableStreamFromNoopAsyncGen = {
   }
 };
 
+export const abortWriterAfterGc = {
+  async test() {
+    function getWriter() {
+      const { writable } = new IdentityTransformStream();
+      return writable.getWriter();
+    }
+
+    const writer = getWriter();
+    gc();
+    await writer.abort();
+  }
+};
+
 export default {
   async fetch(request, env) {
     strictEqual(request.headers.get('content-length'), '10');