fix use-after-free in inplace_stop_source::request_stop()

include/stdexec/stop_token.hpp

@@ -157,12 +157,15 @@ STDEXEC_P2300_NAMESPACE_BEGIN()
 
     void __remove_callback_(STDEXEC::__stok::__inplace_stop_callback_base*) const noexcept;
 
+    void __sync_with_request_stop_() noexcept;
+
     static constexpr uint8_t __stop_requested_flag_ = 1;
     static constexpr uint8_t __locked_flag_         = 2;
 
     mutable STDEXEC::__std::atomic<uint8_t>                __state_{0};
     mutable STDEXEC::__stok::__inplace_stop_callback_base* __callbacks_ = nullptr;
     std::thread::id                                        __notifying_thread_;
+    bool*                                                  __destroyed_ = nullptr;
   };
 
   // [stoptoken.inplace], class inplace_stop_token
@@ -266,17 +269,51 @@ STDEXEC_P2300_NAMESPACE_BEGIN()
 
   inline inplace_stop_source::~inplace_stop_source()
   {
+    if (stop_requested())
+      __sync_with_request_stop_();
+
     STDEXEC_ASSERT((__state_.load(STDEXEC::__std::memory_order_relaxed) & __locked_flag_) == 0);
     STDEXEC_ASSERT(__callbacks_ == nullptr);
   }
 
+  // Kept separate so the destructor's no-stop path stays cheap to inline.
+  inline void inplace_stop_source::__sync_with_request_stop_() noexcept
+  {
+    auto __old_state        = __lock_();
+    auto __destroyed        = __destroyed_;
+    auto __notifying_thread = __notifying_thread_;
+    __unlock_(__old_state);
+
+    if (__destroyed != nullptr && std::this_thread::get_id() == __notifying_thread)
+    {
+      // A stop callback is destroying us from inside request_stop().
+      *__destroyed = true;
+    }
+    else
+    {
+      // Wait for request_stop() to finish on the other thread.
+      STDEXEC::__stok::__spin_wait __spin;
+      while (__destroyed != nullptr)
+      {
+        __spin.__wait();
+        __old_state = __lock_();
+        __destroyed = __destroyed_;
+        __unlock_(__old_state);
+      }
+    }
+  }
+
   inline auto inplace_stop_source::request_stop() noexcept -> bool
   {
     if (!__try_lock_unless_stop_requested_(true))
       return false;
 
     __notifying_thread_ = std::this_thread::get_id();
 
+    // ~inplace_stop_source() sets this if a callback destroys us.
+    bool __destroyed = false;
+    __destroyed_     = &__destroyed;
+
     // We are responsible for executing callbacks.
     while (__callbacks_ != nullptr)
     {
@@ -299,9 +336,14 @@ STDEXEC_P2300_NAMESPACE_BEGIN()
         __callbk->__callback_completed_.store(true, STDEXEC::__std::memory_order_release);
       }
 
+      // The callback destroyed this stop source, so we can't touch *this again.
+      if (__destroyed)
+        return true;
+
       __lock_();
     }
 
+    __destroyed_ = nullptr;
     __state_.store(__stop_requested_flag_, STDEXEC::__std::memory_order_release);
     return true;
   }

test/stdexec/concepts/test_concepts_stop_tokens.cpp

@@ -22,6 +22,11 @@
 #if STDEXEC_USE_MODULES()
 import std;
 #else
+#  include <atomic>
+#  include <chrono>
+#  include <memory>
+#  include <optional>
+#  include <thread>
 #  include <type_traits>
 #endif
 
@@ -63,4 +68,77 @@ namespace
     ::STDEXEC::inplace_stop_callback cb{source.get_token(), on_stop_request{}};
     STATIC_REQUIRE(std::is_same_v<decltype(cb), ::STDEXEC::inplace_stop_callback<on_stop_request>>);
   }
+
+  struct stop_state;
+
+  struct destroy_on_stop
+  {
+    std::unique_ptr<stop_state>* state;
+
+    void operator()() noexcept;
+  };
+
+  struct stop_state
+  {
+    ::STDEXEC::inplace_stop_source                                   source;
+    std::optional<::STDEXEC::inplace_stop_callback<on_stop_request>> other;
+    std::optional<::STDEXEC::inplace_stop_callback<destroy_on_stop>> callback;
+  };
+
+  void destroy_on_stop::operator()() noexcept
+  {
+    state->reset();
+  }
+
+  TEST_CASE("a stop callback can destroy the inplace_stop_source", "[stop_token]")
+  {
+    auto state = std::make_unique<stop_state>();
+    state->other.emplace(state->source.get_token(), on_stop_request{});
+    state->callback.emplace(state->source.get_token(), destroy_on_stop{&state});
+
+    CHECK(state->source.request_stop());
+    CHECK(state == nullptr);
+  }
+
+  struct slow_state;
+
+  struct remove_and_sleep
+  {
+    slow_state*        state;
+    std::atomic<bool>* removed;
+
+    void operator()() noexcept;
+  };
+
+  struct slow_state
+  {
+    ::STDEXEC::inplace_stop_source                                    source;
+    std::optional<::STDEXEC::inplace_stop_callback<remove_and_sleep>> callback;
+  };
+
+  void remove_and_sleep::operator()() noexcept
+  {
+    // Like an operation that completes from inside its stop callback.
+    auto* removed_flag = removed;
+    state->callback.reset();
+    removed_flag->store(true);
+
+    // Keep request_stop() running while the main thread destroys the source.
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+  }
+
+  TEST_CASE("inplace_stop_source destructor waits for request_stop on another thread",
+            "[stop_token]")
+  {
+    std::atomic<bool> removed{false};
+    auto              state = std::make_unique<slow_state>();
+    state->callback.emplace(state->source.get_token(), remove_and_sleep{state.get(), &removed});
+
+    std::thread thread([source = &state->source] { source->request_stop(); });
+    while (!removed.load())
+      std::this_thread::yield();
+
+    state.reset();  // request_stop() is still running on the other thread
+    thread.join();
+  }
 }  // namespace