Merge pull request #232 from cloudflare/harris/fix-uaf-in-input-gate-…

…waiter

Fix UAF in InputGate::Waiter

src/workerd/io/io-gate-test.c++

@@ -182,6 +182,95 @@ KJ_TEST("InputGate nested critical section outlives parent") {
   rootWait.wait(ws);
 }
 
+KJ_TEST("InputGate deeply nested critical sections") {
+  kj::EventLoop loop;
+  kj::WaitScope ws(loop);
+
+  InputGate gate;
+
+  kj::Own<InputGate::CriticalSection> cs1;
+  kj::Own<InputGate::CriticalSection> cs2;
+  kj::Own<InputGate::CriticalSection> cs3;
+  kj::Own<InputGate::CriticalSection> cs4;
+
+  {
+    auto lock = gate.wait().wait(ws);
+    cs1 = lock.startCriticalSection();
+  }
+
+  {
+    auto lock = cs1->wait().wait(ws);
+    cs2 = lock.startCriticalSection();
+  }
+
+  {
+    auto lock = cs2->wait().wait(ws);
+    cs3 = lock.startCriticalSection();
+    cs4 = lock.startCriticalSection();
+  }
+
+  // Start cs2
+  cs2->wait();
+
+  // Add some waiters to cs2, some of which are waiting to start more nested critical sections
+  auto lock = cs2->wait().wait(ws);
+  auto waiter1 = cs2->wait();
+  auto waiter2 = cs2->wait();
+
+  // Both of these wait on cs2 indirectly, as they are nested under cs2
+  auto waiter3 = cs3->wait();
+  auto waiter4 = cs4->wait();
+
+  KJ_EXPECT(!waiter1.poll(ws));
+  KJ_EXPECT(!waiter2.poll(ws));
+  KJ_EXPECT(!waiter3.poll(ws));
+  KJ_EXPECT(!waiter4.poll(ws));
+
+  // Mark cs2 as complete with outstanding waiters, and drop our reference to it.
+  cs2->succeeded();
+  cs2 = nullptr;
+
+  // Our waiters should still be outstanding as we have not released the lock
+  KJ_EXPECT(!waiter1.poll(ws));
+  KJ_EXPECT(!waiter2.poll(ws));
+  KJ_EXPECT(!waiter3.poll(ws));
+  KJ_EXPECT(!waiter4.poll(ws));
+
+  // Drop some outstanding waiters
+  { auto drop = kj::mv(waiter2); }
+  { auto drop = kj::mv(waiter4); }
+
+  // Release the lock on cs2
+  { auto drop = kj::mv(lock); }
+
+  // cs3 should have started
+  KJ_ASSERT(!waiter1.poll(ws));
+  KJ_ASSERT(waiter3.poll(ws));
+  auto lock2 = waiter3.wait(ws);
+
+  // Add a waiter on cs3
+  auto waiter5 = cs3->wait();
+  KJ_ASSERT(!waiter5.poll(ws));
+
+  // Can't start new tasks on the root until both cs1 and cs3 have succeeded, and all outstanding
+  // tasks have either been dropped or completed.
+  auto waiter6 = gate.wait();
+  KJ_ASSERT(!waiter6.poll(ws));
+
+  cs1->succeeded();
+  cs3->succeeded();
+
+  // drop waiter5
+  { auto drop = kj::mv(waiter5); }
+
+  // Release the lock on cs3
+  { auto drop = kj::mv(lock2); }
+
+  // Our root task should be ready now.
+  KJ_ASSERT(waiter6.poll(ws));
+  waiter6.wait(ws);
+}
+
 KJ_TEST("InputGate critical section lock outlives critical section") {
   kj::EventLoop loop;
   kj::WaitScope ws(loop);
@@ -207,6 +296,17 @@ KJ_TEST("InputGate critical section lock outlives critical section") {
 
   // Lock should have been reparented, so should still work.
   KJ_ASSERT(lock.isFor(gate));
+
+  // The gate should still be locked
+  auto waiter = gate.wait();
+  KJ_EXPECT(!waiter.poll(ws));
+
+  // Drop the outstanding lock
+  { auto drop = kj::mv(lock); }
+
+  // Our waiter should resolve now
+  KJ_ASSERT(waiter.poll(ws));
+  KJ_EXPECT(waiter.wait(ws).isFor(gate));
 }
 
 KJ_TEST("InputGate broken") {

src/workerd/io/io-gate.c++

@@ -25,7 +25,7 @@ InputGate::~InputGate() noexcept {
 
 InputGate::Waiter::Waiter(
     kj::PromiseFulfiller<Lock>& fulfiller, InputGate& gate, bool isChildWaiter)
-    : fulfiller(fulfiller), gate(gate), isChildWaiter(isChildWaiter) {
+    : fulfiller(fulfiller), gate(&gate), isChildWaiter(isChildWaiter) {
   gate.hooks.inputGateWaiterAdded();
   if (isChildWaiter) {
     gate.waitingChildren.add(*this);
@@ -34,12 +34,12 @@ InputGate::Waiter::Waiter(
   }
 }
 InputGate::Waiter::~Waiter() noexcept(false) {
-  gate.hooks.inputGateWaiterRemoved();
+  gate->hooks.inputGateWaiterRemoved();
   if (link.isLinked()) {
     if (isChildWaiter) {
-      gate.waitingChildren.remove(*this);
+      gate->waitingChildren.remove(*this);
     } else {
-      gate.waiters.remove(*this);
+      gate->waiters.remove(*this);
     }
   }
 }
@@ -65,8 +65,14 @@ kj::Promise<void> InputGate::onBroken() {
 void InputGate::releaseLock() {
   if (isCriticalSection) {
     auto& self = static_cast<CriticalSection&>(*this);
-    if (self.state == CriticalSection::DONE) {
-      // Tasks and locks are reparented at this point.
+    if (self.state == CriticalSection::REPARENTED) {
+      // This lock was for a critical section that has already completed, therefore the lock
+      // should be considered "reparented", and we should forward the release to the parent.
+
+      // Ensure any waiters on us have already been reparented.
+      KJ_DASSERT(self.waitingChildren.size() == 0);
+      KJ_DASSERT(self.waiters.size() == 0);
+
       self.parentAsInputGate().releaseLock();
       return;
     }
@@ -102,6 +108,8 @@ kj::Maybe<InputGate::CriticalSection&> InputGate::Lock::getCriticalSection() {
 }
 
 bool InputGate::Lock::isFor(const InputGate& otherGate) const {
+  KJ_ASSERT(!otherGate.isCriticalSection);
+
   InputGate* ptr = gate;
   while (ptr->isCriticalSection) {
     ptr = &static_cast<CriticalSection&>(*ptr).parentAsInputGate();
@@ -132,7 +140,7 @@ InputGate::CriticalSection::~CriticalSection() noexcept(false) {
           "awaits a task that was initiated outside of the critical section. Since a critical "
           "section blocks all other tasks from completing, this leads to deadlock."));
       break;
-    case DONE:
+    case REPARENTED:
       // Common case.
       break;
   }
@@ -143,38 +151,20 @@ kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait() {
     case NOT_STARTED: {
       state = INITIAL_WAIT;
 
-      // Find the ancestor that we want to lock, skipping parents in the DONE state since children
-      // of DONE critical sections are reparented.
-      InputGate* target = nullptr;
-      for (CriticalSection* ptr = this; target == nullptr;) {
-        KJ_SWITCH_ONEOF(ptr->parent) {
-          KJ_CASE_ONEOF(p, InputGate*) {
-            target = p;
-          }
-          KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
-            if (c.get()->state == DONE) {
-              // Keep looping...
-              ptr = c;
-            } else {
-              target = c;
-            }
-          }
-        }
-      }
-
-      KJ_IF_MAYBE(e, target->brokenState.tryGet<kj::Exception>()) {
+      auto& target = parentAsInputGate();
+      KJ_IF_MAYBE(e, target.brokenState.tryGet<kj::Exception>()) {
         // Oops, we're broken.
         setBroken(*e);
         return kj::cp(*e);
       }
 
       // Add ourselves to this parent's child waiter list.
-      if (target->lockCount == 0) {
+      if (target.lockCount == 0) {
         state = RUNNING;
-        parentLock = Lock(*target);
+        parentLock = Lock(target);
         return wait();
       } else {
-        return kj::newAdaptedPromise<Lock, Waiter>(*target, true)
+        return kj::newAdaptedPromise<Lock, Waiter>(target, true)
             .then([this](Lock lock) {
           state = RUNNING;
           parentLock = kj::mv(lock);
@@ -194,7 +184,7 @@ kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait() {
     case RUNNING:
       // CriticalSection is active, so defer to InputGate implementation.
       return InputGate::wait();
-    case DONE:
+    case REPARENTED:
       // Once the CriticalSection has declared itself done, then any straggler tasks it initiated
       // are adopted by the parent.
       // WARNING: Don't use parentAsInputGate() here as that'll bypass the override of wait() if
@@ -219,17 +209,19 @@ InputGate::Lock InputGate::CriticalSection::succeeded() {
   // adopted by the parent.
   auto& parentGate = parentAsInputGate();
   for (auto& waiter: waitingChildren) {
-    waiters.remove(waiter);
+    waitingChildren.remove(waiter);
     parentGate.waitingChildren.add(waiter);
+    waiter.gate = &parentGate;
   }
   for (auto& waiter: waiters) {
     waiters.remove(waiter);
     parentGate.waiters.add(waiter);
+    waiter.gate = &parentGate;
   }
   parentGate.lockCount += lockCount;
   lockCount = 0;
 
-  state = DONE;
+  state = REPARENTED;
   auto result = KJ_ASSERT_NONNULL(kj::mv(parentLock));
   parentLock = nullptr;
   return result;
@@ -267,6 +259,25 @@ void InputGate::setBroken(const kj::Exception& e) {
   brokenState = kj::cp(e);
 }
 
+InputGate& InputGate::CriticalSection::parentAsInputGate() {
+  CriticalSection* ptr = this;
+  for(;;) {
+    KJ_SWITCH_ONEOF(ptr->parent) {
+      KJ_CASE_ONEOF(p, InputGate*) {
+        return *p;
+      }
+      KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
+        if (c.get()->state == REPARENTED) {
+          // Keep looping...
+          ptr = c;
+        } else {
+          return *c.get();
+        }
+      }
+    }
+  }
+}
+
 // =======================================================================================
 
 OutputGate::OutputGate(Hooks& hooks)

src/workerd/io/io-gate.h

@@ -73,8 +73,9 @@ class InputGate {
     // before the gate is unlocked.
 
     kj::Own<CriticalSection> startCriticalSection();
-    // Start a new critical section from this lock. No further Locks will be handed out by
-    // InputGate::lock() until the CriticalSection has been dropped.
+    // Start a new critical section from this lock. After `wait()` has been called on the returned
+    // critical section for the first time, no further Locks will be handed out by
+    // InputGate::wait() until the CriticalSection has been dropped.
     //
     // CriticalSections can be nested. If this Lock is itself part of a CriticalSection, the new
     // CriticalSection will be nested within it and the outer CriticalSection's wait() won't
@@ -120,7 +121,7 @@ class InputGate {
     ~Waiter() noexcept(false);
 
     kj::PromiseFulfiller<Lock>& fulfiller;
-    InputGate& gate;
+    InputGate* gate;
     bool isChildWaiter;
     kj::ListLink<Waiter> link;
   };
@@ -192,7 +193,7 @@ class InputGate::CriticalSection: private InputGate, public kj::Refcounted {
     RUNNING,
     // First lock has been obtained, waiting for success() or failed().
 
-    DONE
+    REPARENTED
     // success() or failed() has been called.
   };
 
@@ -207,17 +208,8 @@ class InputGate::CriticalSection: private InputGate, public kj::Refcounted {
 
   friend class InputGate;
 
-  InputGate& parentAsInputGate() {
-    KJ_SWITCH_ONEOF(parent) {
-      KJ_CASE_ONEOF(p, InputGate*) {
-        return *p;
-      }
-      KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
-        return *c;
-      }
-    }
-    KJ_UNREACHABLE;
-  }
+  InputGate& parentAsInputGate();
+  // Return a reference for the parent scope, skipping any reparented CriticalSections
 };
 
 class OutputGate {