DeferredProxyCoroutine: Implement a guarantee against incorrect destruction order

Author: harrishancock

src/workerd/api/deferred-proxy-test.c++

@@ -229,5 +229,47 @@ KJ_TEST("kj::Promise<DeferredProxy<T>>: can be canceled while suspended after de
   KJ_EXPECT(unwind == 4);
 }
 
+KJ_TEST("kj::Promise<DeferredProxy<T>>: destroying inner PromiseNode before outer does not "
+    "segfault") {
+  // Destroy the inner promise before the outer promise to test our safeguard against incorrect
+  // destruction order causing segfaults.
+
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  auto coro = []() -> kj::Promise<DeferredProxy<void>> {
+    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
+    co_await kj::Promise<void>(kj::NEVER_DONE);
+  };
+
+  auto outer = coro();
+
+  // We could call `get()` on the outer node immediately, even before it reports it is ready, but
+  // we call `poll()` for good measure, in case the DeferredProxyCoroutine implementation ever
+  // changes to disallow `get()`-before-ready. We cannot use `wait()` for this purpose, because
+  // `wait()` would avoid the segfault by (correctly) destroying the outer PromiseNode before
+  // returning the result to us.
+  KJ_EXPECT(outer.poll(waitScope));
+
+  auto outerNode = kj::_::PromiseNode::from(kj::mv(outer));
+
+  // `poll()`, unlike `wait()`, does not call `setSelfPointer()` on the outer PromiseNode, which
+  // would cause an assertion failure inside the outer PromiseNode's `get()` implementation, so we
+  // have to do it ourselves.
+  outerNode->setSelfPointer(&outerNode);
+
+  kj::_::ExceptionOr<DeferredProxy<void>> result;
+  outerNode->get(result);
+
+  {
+    // Destroy the inner promise.
+    auto inner = kj::mv(KJ_ASSERT_NONNULL(result.value).proxyTask);
+  }
+
+  // Destroy the outer promise. At one time, this caused a segfault ... or at least it produced
+  // invalid accesses under valgrind. :/
+  outerNode = nullptr;
+}
+
 }  // namespace
 }  // namespace workerd::api

src/workerd/api/deferred-proxy.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <kj/async.h>
+#include <kj/debug.h>
 
 namespace workerd::api {
 
@@ -211,14 +212,19 @@ class DeferredProxyCoroutine: public kj::_::PromiseNode,
 
   // PromiseNode implementation
 
+  void setSelfPointer(kj::_::OwnPromiseNode* selfPtr) noexcept override {
+    this->selfPtr = selfPtr;
+  }
+
   void destroy() override {
     // The promise returned by `inner.get_return_object()` is what actually owns this coroutine
     // frame. We temporarily store that in `result` until our outer promise is fulfilled. So, to
     // destroy ourselves, we must manually drop `result`.
     //
-    // On the other hand, if our outer promise has already been fulfilled, and `result` delivered to
-    // wherever it is going, then someone else directly owns the coroutine now, not us, and it's
-    // okay for this drop to be a no-op.
+    // On the other hand, if our outer promise has already been fulfilled, then `result` has already
+    // been delivered to wherever it is going, and someone else directly owns the coroutine now, not
+    // us. In this case, this `destroy()` override will have already been called (and it will have
+    // been a no-op), because our own OwnPromiseNode will have already been dropped in `get()`.
 
     auto drop = kj::mv(result);
   }
@@ -228,6 +234,13 @@ class DeferredProxyCoroutine: public kj::_::PromiseNode,
   }
 
   void get(kj::_::ExceptionOrValue& output) noexcept override {
+    // Make sure that the outer PromiseNode (`this` one) is destroyed before the inner PromiseNode.
+    // kj-async should already provide us this guarantee, but since incorrect destruction order
+    // would cause invalid memory access, we provide a stronger guarantee. Also see the comment for
+    // the `result` data member.
+    KJ_ASSERT(selfPtr != nullptr);
+    KJ_DEFER(*selfPtr = nullptr);
+
     static_cast<decltype(result)&>(output) = kj::mv(result);
   }
 
@@ -249,6 +262,16 @@ class DeferredProxyCoroutine: public kj::_::PromiseNode,
 
   kj::_::ExceptionOr<DeferredProxy<T>> result;
   // Stores the result for the outer promise.
+  //
+  // WARNING: This object owns `this` PromiseNode! If `result` is ever moved away, as is done in
+  // `get()`, we must arrange to make sure that no one ever tries to use `this` PromiseNode again.
+  // Stated another way, we must guarantee that the outer PromiseNode (for `DeferredProxy<T>`) is
+  // always destroyed before the inner PromiseNode (for `T`). kj-async always does this anyway, but
+  // we implement an additional safeguard by immediately destroying our own `OwnPromiseNode` (which
+  // we have access to via `setSelfPointer()`) when we move `result` away in `get()`.
+
+  kj::_::OwnPromiseNode* selfPtr = nullptr;
+  // Used to drop ourselves in `get()` -- see comment for `result`.
 
   bool deferredProxyingHasBegun = false;
   // Set to true when deferred proxying has begun -- that is, when the outer DeferredProxy<T>