cancellation: Fix UAF in cancellation_resolver

include/async/cancellation.hpp

@@ -293,29 +293,52 @@ struct cancellation_resolver final : private abstract_cancellation_callback {
 		if (!(st & done_listen) || !(st & done_completion_path))
 			return;
 
-		if (!(st & done_cancellation_path)) {
-			// Try to unregister from the cancellation event once both listen() and complete() were called.
-			// Note that we enter this code path at most once since done_cancellation_path is the only missing
-			// bit in state_ and the next call to transition_() will necessarily set it.
-			assert(event_);
-
-			frg::unique_lock guard{event_->_mutex};
-			if (event_->_was_requested)
-				return;
-			auto it = event_->_cbs.iterator_to(this);
-			event_->_cbs.erase(it);
+		if (!(old_st & done_listen) || !(old_st & done_completion_path)) {
+			// If at least one of done_listen or done_completion_path was unset before,
+			// we perform the unregistration step. This happens exactly once.
+			unsigned int bits = done_unregister;
+			if (!(st & done_cancellation_path)) {
+				// Try to unregister from the cancellation event.
+				assert(event_);
+
+				frg::unique_lock guard{event_->_mutex};
+				if (!event_->_was_requested) {
+					auto it = event_->_cbs.iterator_to(this);
+					event_->_cbs.erase(it);
+					// call() can never run anymore.
+					bits |= done_cancellation_path;
+				}
+			}
+			// Set done_unregister (and potentially done_cancellation_path) last
+			// as accessing *this may be UB afterwards (see below).
+			st = state_.fetch_or(bits, std::memory_order_acq_rel) | bits;
+		} else {
+			// Both done_listen and done_completion_path were already set before.
+			assert(new_bits == done_cancellation_path);
 		}
 
-		// Call resume() when all code paths are done. This can only happen once.
+		// Call resume() once all code paths are done and unregistration is finished (i.e., all four bits are set).
+		// If the cancellation path and unregistering race, both of them perform an RMW on state_
+		// and exactly one of them observes the other's bit. resume() is hence called exactly once.
+		// If we are not the code path that calls resume_() here,
+		// accessing *this is UB as resume_() may concurrently destroy *this.
+		if (!(st & done_cancellation_path) || !(st & done_unregister))
+			return;
 		resume_(this);
 	}
 
+	// Note: resume() is called once all four of the bits below are set.
+	//       Hence, *this also only remains valid until this happens.
+	//       All code paths that access *this must guarantee that at least one bit is still unset.
+
 	// Set in state_ when listen() is done.
 	static constexpr unsigned int done_listen = 1u << 0;
 	// Set in state_ when complete() is done or when we know that it will never be called.
 	static constexpr unsigned int done_completion_path = 1u << 1;
 	// Set in state_ when tryCancel() is done or when we know that it will never be called.
 	static constexpr unsigned int done_cancellation_path = 1u << 2;
+	// Set in state_ when unregistration from the cancellation event is done or when we know that it is unnecessary.
+	static constexpr unsigned int done_unregister = 1u << 3;
 
 	cancellation_event *event_{nullptr};
 	std::atomic<unsigned int> state_{0};

tests/post-ack.cpp

@@ -1,4 +1,6 @@
+#include <atomic>
 #include <memory>
+#include <thread>
 #include <async/post-ack.hpp>
 #include <gtest/gtest.h>
 
@@ -74,3 +76,55 @@ TEST(PostAck, ImmovableType) {
 	ASSERT_EQ(ok1_ctr, 8);
 	ASSERT_EQ(ok2_ctr, 8);
 }
+
+// Stress test for the race between completion and cancellation of a poll operation;
+// resuming (and thereby destructing) the operation must not overlap with the racing
+// thread's accesses inside cancellation_resolver (cf. managarm/managarm#1509).
+// Best run under ASan/TSan to detect regressions.
+TEST(PostAck, PollCancelRace) {
+	for (int round = 0; round < 10000; ++round) {
+		async::post_ack_mechanism<int> mech;
+		async::post_ack_agent<int> agent;
+		agent.attach(&mech);
+		async::cancellation_event ce;
+
+		std::atomic<bool> consumerDone{false};
+		std::atomic<bool> producerDone{false};
+
+		auto consumer = [&] () -> async::detached {
+			{
+				auto handle = co_await agent.poll(ce);
+				if (handle)
+					handle.ack();
+			}
+			consumerDone.store(true, std::memory_order_release);
+		};
+
+		auto producer = [&] () -> async::detached {
+			co_await mech.post(round);
+			producerDone.store(true, std::memory_order_release);
+		};
+
+		std::atomic<int> ready{0};
+		auto spinUntilReady = [&] {
+			ready.fetch_add(1, std::memory_order_relaxed);
+			while (ready.load(std::memory_order_relaxed) != 3)
+				;
+		};
+		std::thread poller{[&] { spinUntilReady(); consumer(); }};
+		std::thread poster{[&] { spinUntilReady(); producer(); }};
+		std::thread canceller{[&] { spinUntilReady(); ce.cancel(); }};
+		poller.join();
+		poster.join();
+		canceller.join();
+
+		while (!consumerDone.load(std::memory_order_acquire))
+			std::this_thread::yield();
+
+		// If the poll was cancelled, detach() acks the outstanding post.
+		agent.detach();
+
+		while (!producerDone.load(std::memory_order_acquire))
+			std::this_thread::yield();
+	}
+}