Additional safety guards in stream queue consumer iteration (#5932)

Co-authored-by: Harris Hancock <harris@cloudflare.com>

Author: jasnell

src/workerd/api/streams/queue-test.c++

@@ -1933,5 +1933,127 @@ KJ_TEST("ValueQueue error then destroy before consumer doesn't crash") {
 
 #pragma endregion Queue / Consumer Destruction Order Tests
 
+#pragma region Consumer Destroyed During Push Tests
+// These tests verify that the queue handles consumer destruction gracefully.
+// QueueImpl::push() takes a snapshot of consumers and iterates over it. If a consumer
+// is destroyed (removed from allConsumers) between the snapshot and the iteration,
+// the code must check allConsumers.contains() before dereferencing the pointer.
+// That said, the tests do not actually fail without the relevant fix because the
+// issue is extremely timing-dependent and difficult to trigger deterministically, even
+// with asan enabled. These tests at least document the intended behavior and may catch
+// future regressions.
+
+KJ_TEST("ByteQueue push skips consumer removed from queue during iteration") {
+  preamble([](jsg::Lock& js) {
+    ByteQueue queue(10);
+
+    // Create two consumers
+    auto consumer1 = kj::heap<ByteQueue::Consumer>(queue);
+    auto consumer2 = kj::heap<ByteQueue::Consumer>(queue);
+
+    // Destroy consumer2 BEFORE pushing. This directly tests that push()
+    // checks if consumers still exist before calling push on them.
+    // The snapshot taken at the start of push() would have included consumer2,
+    // but consumer2 is no longer in allConsumers when we iterate.
+    consumer2 = nullptr;
+
+    // Push data - should not crash even though consumer2 was in the queue
+    // when it was created but is now destroyed.
+    auto store = jsg::BackingStore::alloc(js, 4);
+    store.asArrayPtr().fill('x');
+    queue.push(js, kj::rc<ByteQueue::Entry>(jsg::BufferSource(js, kj::mv(store))));
+
+    // consumer1 should have received the data
+    KJ_ASSERT(consumer1->size() == 4);
+  });
+}
+
+KJ_TEST("ValueQueue push skips consumer removed from queue during iteration") {
+  preamble([](jsg::Lock& js) {
+    ValueQueue queue(10);
+
+    auto consumer1 = kj::heap<ValueQueue::Consumer>(queue);
+    auto consumer2 = kj::heap<ValueQueue::Consumer>(queue);
+
+    // Destroy consumer2 before pushing
+    consumer2 = nullptr;
+
+    queue.push(js, getEntry(js, 4));
+
+    KJ_ASSERT(consumer1->size() == 4);
+  });
+}
+
+KJ_TEST("ByteQueue push handles consumer destroyed by microtask between pushes") {
+  preamble([](jsg::Lock& js) {
+    ByteQueue queue(10);
+
+    auto consumer1 = kj::heap<ByteQueue::Consumer>(queue);
+    auto consumer2 = kj::heap<ByteQueue::Consumer>(queue);
+
+    // Set up a pending read on consumer1
+    auto prp = js.newPromiseAndResolver<ReadResult>();
+    consumer1->read(js,
+        ByteQueue::ReadRequest(kj::mv(prp.resolver),
+            {
+              .store = jsg::BufferSource(js, jsg::BackingStore::alloc(js, 4)),
+            }));
+
+    // The continuation destroys consumer2
+    MustCall<ReadContinuation> readContinuation([&consumer2](jsg::Lock& js, ReadResult&& result) {
+      consumer2 = nullptr;
+      return js.resolvedPromise(kj::mv(result));
+    });
+    prp.promise.then(js, readContinuation);
+
+    // First push - resolves consumer1's read, schedules microtask that will destroy consumer2
+    auto store1 = jsg::BackingStore::alloc(js, 4);
+    store1.asArrayPtr().fill('x');
+    queue.push(js, kj::rc<ByteQueue::Entry>(jsg::BufferSource(js, kj::mv(store1))));
+
+    // Run microtasks - this destroys consumer2
+    js.runMicrotasks();
+
+    // Second push - consumer2 is now destroyed, should not crash
+    auto store2 = jsg::BackingStore::alloc(js, 4);
+    store2.asArrayPtr().fill('y');
+    queue.push(js, kj::rc<ByteQueue::Entry>(jsg::BufferSource(js, kj::mv(store2))));
+
+    // consumer1 should have the second push's data buffered
+    KJ_ASSERT(consumer1->size() == 4);
+  });
+}
+
+KJ_TEST("ByteQueue maybeUpdateBackpressure skips destroyed consumers") {
+  preamble([](jsg::Lock& js) {
+    ByteQueue queue(10);
+
+    auto consumer1 = kj::heap<ByteQueue::Consumer>(queue);
+    auto consumer2 = kj::heap<ByteQueue::Consumer>(queue);
+
+    // Push some data so consumers have size
+    auto store = jsg::BackingStore::alloc(js, 4);
+    store.asArrayPtr().fill('x');
+    queue.push(js, kj::rc<ByteQueue::Entry>(jsg::BufferSource(js, kj::mv(store))));
+
+    KJ_ASSERT(consumer1->size() == 4);
+    KJ_ASSERT(consumer2->size() == 4);
+    KJ_ASSERT(queue.size() == 4);
+
+    // Destroy consumer2
+    consumer2 = nullptr;
+
+    // Trigger backpressure recalculation by pushing more data
+    auto store2 = jsg::BackingStore::alloc(js, 4);
+    store2.asArrayPtr().fill('y');
+    queue.push(js, kj::rc<ByteQueue::Entry>(jsg::BufferSource(js, kj::mv(store2))));
+
+    // Should not crash, and size should reflect only consumer1
+    KJ_ASSERT(consumer1->size() == 8);
+    KJ_ASSERT(queue.size() == 8);
+  });
+}
+#pragma endregion Consumer Destroyed During Push Tests
+
 }  // namespace
 }  // namespace workerd::api

src/workerd/api/streams/queue.h

@@ -10,6 +10,7 @@
 #include <workerd/util/ring-buffer.h>
 #include <workerd/util/small-set.h>
 #include <workerd/util/state-machine.h>
+#include <workerd/util/weak-refs.h>
 
 namespace workerd::api {
 
@@ -165,22 +166,18 @@ class QueueImpl final {
     // Detach all consumers before destruction to prevent UAF.
     // This can happen during isolate teardown when the destruction order
     // of JS wrapper objects doesn't follow the ownership hierarchy.
-    auto consumers = allConsumers.snapshot();
-    for (auto consumer: consumers) {
-      consumer->detachQueue();
-    }
+    allConsumers.forEach([&](ConsumerImpl& consumer) { consumer.detachQueue(); });
   }
 
   // Closes the queue. The close is forwarded on to all consumers.
   // If we are already closed or errored, do nothing here.
   void close(jsg::Lock& js) {
     if (state.isActive()) {
-      // We copy the list of consumers in case the consumers remove themselves
-      // from the queue during the close callback, invalidating the iterator.
-      auto consumers = allConsumers.snapshot();
-      for (auto consumer: consumers) {
-        consumer->close(js);
-      }
+#ifdef KJ_DEBUG
+      isClosingOrErroring = true;
+      KJ_DEFER(isClosingOrErroring = false);
+#endif
+      allConsumers.forEach([&](ConsumerImpl& consumer) { consumer.close(js); });
       state.template transitionTo<Closed>();
     }
   }
@@ -199,12 +196,11 @@ class QueueImpl final {
   // If we are already closed or errored, do nothing here.
   void error(jsg::Lock& js, jsg::Value reason) {
     if (state.isActive()) {
-      // We copy the list of consumers in case the consumers remove themselves
-      // from the queue during the error callback, invalidating the iterator.
-      auto consumers = allConsumers.snapshot();
-      for (auto consumer: consumers) {
-        consumer->error(js, reason.addRef(js));
-      }
+#ifdef KJ_DEBUG
+      isClosingOrErroring = true;
+      KJ_DEFER(isClosingOrErroring = false);
+#endif
+      allConsumers.forEach([&](ConsumerImpl& consumer) { consumer.error(js, reason.addRef(js)); });
       state.template transitionTo<Errored>(kj::mv(reason));
     }
   }
@@ -215,9 +211,9 @@ class QueueImpl final {
   void maybeUpdateBackpressure() {
     totalQueueSize = 0;
     if (state.isActive()) {
-      for (auto consumer: allConsumers.snapshot()) {
-        totalQueueSize = kj::max(totalQueueSize, consumer->size());
-      }
+      allConsumers.forEach([&](ConsumerImpl& consumer) {
+        totalQueueSize = kj::max(totalQueueSize, consumer.size());
+      });
     }
   }
 
@@ -229,15 +225,14 @@ class QueueImpl final {
   void push(jsg::Lock& js, kj::Rc<Entry> entry, kj::Maybe<ConsumerImpl&> skipConsumer = kj::none) {
     state.requireActiveUnsafe("The queue is closed or errored.");
 
-    for (auto consumer: allConsumers.snapshot()) {
+    allConsumers.forEach([&](ConsumerImpl& consumer) {
       KJ_IF_SOME(skip, skipConsumer) {
-        if (&skip == consumer) {
-          continue;
+        if (&skip == &consumer) {
+          return;
         }
       }
-
-      consumer->push(js, entry->clone(js));
-    }
+      consumer.push(js, entry->clone(js));
+    });
   }
 
   // The current size of consumer with the most stored data.
@@ -251,8 +246,10 @@ class QueueImpl final {
 
   bool wantsRead() const {
     if (state.isActive()) {
-      for (auto consumer: allConsumers.snapshot()) {
-        if (consumer->hasReadRequests()) return true;
+      for (const auto& weakRef: allConsumers) {
+        KJ_IF_SOME(consumer, weakRef->tryGet()) {
+          if (consumer.hasReadRequests()) return true;
+        }
       }
     }
     return false;
@@ -302,14 +299,31 @@ class QueueImpl final {
   // will be a very small number (often just one or two), so we use SmallSet to
   // optimize for that. This persists across state transitions so we can detach
   // consumers even after close()/error() transitions the queue to a terminal state.
-  SmallSet<ConsumerImpl*> allConsumers;
-
-  void addConsumer(ConsumerImpl* consumer) {
-    allConsumers.add(consumer);
+  //
+  // We store weak references to consumers to safely handle the case where a consumer
+  // is destroyed during iteration (e.g., resolving a read request triggers JS that
+  // destroys another consumer in the same queue). When iterating, we check if the WeakRef is still valid.
+  SmallSet<kj::Rc<WeakRef<ConsumerImpl>>> allConsumers;
+
+#ifdef KJ_DEBUG
+  // Debug flag to detect if addConsumer is called during close/error iteration.
+  // This should never happen - it would indicate a bug in the streams implementation.
+  bool isClosingOrErroring = false;
+#endif
+
+  void addConsumer(kj::Rc<WeakRef<ConsumerImpl>> weakRef) {
+    KJ_DASSERT(
+        !isClosingOrErroring, "Cannot add a consumer while the queue is being closed or errored");
+    allConsumers.add(kj::mv(weakRef));
   }
 
-  void removeConsumer(ConsumerImpl* consumer) {
-    allConsumers.remove(consumer);
+  void removeConsumer(ConsumerImpl& consumer) {
+    allConsumers.removeIf([&consumer](const kj::Rc<WeakRef<ConsumerImpl>>& ref) {
+      KJ_IF_SOME(c, ref->tryGet()) {
+        return &c == &consumer;
+      }
+      return false;  // Already invalid, will be cleaned up later
+    });
     maybeUpdateBackpressure();
   }
 
@@ -356,7 +370,7 @@ class ConsumerImpl final {
       : queue(queue),
         state(ConsumerState::template create<Ready>()),
         stateListener(stateListener) {
-    queue.addConsumer(this);
+    queue.addConsumer(selfRef.addRef());
   }
 
   explicit ConsumerImpl(kj::Maybe<ConsumerImpl::StateListener&> stateListener)
@@ -369,9 +383,13 @@ class ConsumerImpl final {
   ~ConsumerImpl() noexcept(false) {
     // queue may be none if the queue was destroyed before this consumer
     // (e.g., during isolate teardown) or if cloned from a closed stream.
+    // We must remove ourselves before invalidating selfRef, otherwise
+    // removeConsumer won't find us (tryGet() would return none).
     KJ_IF_SOME(q, queue) {
-      q.removeConsumer(this);
+      q.removeConsumer(*this);
     }
+    // Invalidate after removal so any concurrent iteration will skip us.
+    selfRef->invalidate();
   }
 
   // Called by QueueImpl destructor to detach this consumer from a queue
@@ -587,6 +605,11 @@ class ConsumerImpl final {
   kj::Maybe<QueueImpl&> queue;
   ConsumerState state;
   kj::Maybe<ConsumerImpl::StateListener&> stateListener;
+  // WeakRef to this consumer, used for safe registration with QueueImpl.
+  // When this consumer is destroyed, we invalidate the WeakRef so that
+  // any iteration over allConsumers in QueueImpl will safely skip us.
+  kj::Rc<WeakRef<ConsumerImpl>> selfRef =
+      kj::rc<WeakRef<ConsumerImpl>>(kj::Badge<ConsumerImpl>{}, *this);
 
   bool isClosing() {
     // Closing state is determined by whether there is a Close sentinel that has been

src/workerd/util/BUILD.bazel

@@ -386,7 +386,10 @@ kj_test(
 
 kj_test(
     src = "small-set-test.c++",
-    deps = [":small-set"],
+    deps = [
+        ":small-set",
+        ":util",
+    ],
 )
 
 wd_cc_library(