From d5767113a3dd5e385ddb3627b4d1dab4cc998cb1 Mon Sep 17 00:00:00 2001
From: Amber Ehrlich <miuna.oshino@gmail.com>
Date: Sat, 13 Jul 2024 12:11:40 -0400
Subject: [PATCH] fix(coro): fix leak of the coroutine frame, add
 if_this_causes_an_invalid_read_your_promise_was_destroyed_before_your_awaitable____check_your_promise_lifetime

---
 include/dpp/coro/awaitable.h | 15 +++++++++++----
 include/dpp/coro/task.h      |  2 +-
 src/unittest/coro.cpp        | 13 +++++++++----
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/include/dpp/coro/awaitable.h b/include/dpp/coro/awaitable.h
index 20b091decd..4ba8ed79a9 100644
--- a/include/dpp/coro/awaitable.h
+++ b/include/dpp/coro/awaitable.h
@@ -294,6 +294,15 @@ class awaitable : public basic_awaitable<awaitable<T>> {
 	awaitable(awaitable&& rhs) noexcept : state_ptr(std::exchange(rhs.state_ptr, nullptr)) {
 	}
 
+	/**
+	 * @brief Title :)
+	 *
+	 * We use this in the destructor
+	 */
+	void if_this_causes_an_invalid_read_your_promise_was_destroyed_before_your_awaitable____check_your_promise_lifetime() {
+		abandon();
+	}
+
 	/**
 	 * @brief Destructor.
 	 *
@@ -656,9 +665,7 @@ auto awaitable<T>::abandon() -> uint8_t {
 
 template <typename T>
 awaitable<T>::~awaitable() {
-	if (state_ptr) {
-		state_ptr->state.fetch_or(state_flags::sf_broken, std::memory_order::acq_rel);
-	}
+	if_this_causes_an_invalid_read_your_promise_was_destroyed_before_your_awaitable____check_your_promise_lifetime();
 }
 
 template <typename T>
@@ -691,7 +698,7 @@ bool awaitable<T>::awaiter<Derived>::await_suspend(detail::std_coroutine::corout
 template <typename T>
 template <typename Derived>
 T awaitable<T>::awaiter<Derived>::await_resume() {
-	auto &promise = *std::exchange(awaitable_obj.state_ptr, nullptr);
+	auto &promise = *awaitable_obj.state_ptr;
 
 	promise.state.fetch_and(~detail::promise::sf_awaited, std::memory_order::acq_rel);
 	if (std::holds_alternative<std::exception_ptr>(promise.value)) {
diff --git a/include/dpp/coro/task.h b/include/dpp/coro/task.h
index 9225f2d053..8e6c862b3d 100644
--- a/include/dpp/coro/task.h
+++ b/include/dpp/coro/task.h
@@ -247,7 +247,7 @@ struct promise_base : basic_promise<R> {
 	 * Stores the exception pointer to rethrow on co_await. If the task object is destroyed and was not cancelled, throw instead
 	 */
 	void unhandled_exception() {
-		if ((this->state.load() == promise::state_flags::sf_broken) && !cancelled) {
+		if ((this->state.load() & promise::state_flags::sf_broken) && !cancelled) {
 			throw;
 		}
 		this->template set_exception<false>(std::current_exception());
diff --git a/src/unittest/coro.cpp b/src/unittest/coro.cpp
index ea35e919ca..eead3c85b6 100644
--- a/src/unittest/coro.cpp
+++ b/src/unittest/coro.cpp
@@ -147,6 +147,11 @@ dpp::task<void> task_offline_test() {
 	if (int ret = co_await simple_awaitable{test, 42}; ret != 42) {
 		set_status(test, ts_failed, "failed simple awaitable");
 	}
+
+	if (int ret = co_await []() -> dpp::task<int> { co_return 42; }(); ret != 42) {
+		set_status(test, ts_failed, "failed simple task");
+	}
+	
 	std::array<dpp::task<int>, 10> tasks;
 	auto start = clock::now();
 	for (int i = 0; i < 10; ++i) {
@@ -408,9 +413,9 @@ dpp::job coro_awaitable_test() {
 		{
 			dpp::promise<int> test;
 			dpp::awaitable     awaitable{test.get_awaitable()};
-			std::thread        th{[p = std::move(test)]() mutable {
+			std::thread        th{[&test]() mutable {
 				std::this_thread::sleep_for(std::chrono::seconds(2));
-				p.set_value(69);
+				test.set_value(69);
 			}};
 			th.detach();
 			if (std::optional<int> res = awaitable.sync_wait_for(std::chrono::seconds(5)); !res || *res != 69) {
@@ -450,9 +455,9 @@ dpp::job coro_awaitable_test() {
 		{
 			dpp::promise<void> test;
 			dpp::awaitable     awaitable{test.get_awaitable()};
-			std::thread        th{[p = std::move(test)]() mutable {
+			std::thread        th{[&test]() mutable {
 				std::this_thread::sleep_for(std::chrono::seconds(2));
-				p.set_exception(std::make_exception_ptr(dpp::voice_exception("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")));
+				test.set_exception(std::make_exception_ptr(dpp::voice_exception("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")));
 			}};
 			th.detach();
 			bool success = false;