[PATCH] shared/timeout: Keep callback data alive during removal

[Kyle Pepper <cochcoder@proton.me>]

A timeout callback can remove its own timer. The mainloop then destroys
the wrapper before timeout_callback reads its identifier again.

Hold a reference while dispatching and clear the identifier when the
timer is destroyed. This also prevents a removed callback from modifying
a replacement timer that reused its descriptor. Keep user destruction
synchronous for timerfd timers.

Add unit coverage for idle and timerfd timers with one-shot, repeating,
self-removing and replacement callbacks.

Fixes: 1f11ee8f5194 ("shared: Add timeout handling with mainloop support")
Assisted-by: Codex:GPT-6
---
 src/shared/timeout-mainloop.c |  26 +++++++--
 unit/test-mainloop.c          | 100 ++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+), 5 deletions(-)

diff --git a/src/shared/timeout-mainloop.c b/src/shared/timeout-mainloop.c
index e6f9c5e826..edcc801114 100644
--- a/src/shared/timeout-mainloop.c
+++ b/src/shared/timeout-mainloop.c
@@ -16,31 +16,46 @@
 
 struct timeout_data {
 	int id;
+	unsigned int ref_count;
 	timeout_func_t func;
 	timeout_destroy_func_t destroy;
 	unsigned int timeout;
 	void *user_data;
 };
 
+static void timeout_unref(struct timeout_data *data)
+{
+	if (!--data->ref_count)
+		free(data);
+}
+
 static void timeout_callback(int id, void *user_data)
 {
 	struct timeout_data *data = user_data;
+	bool repeat;
 
-	if (data->func(data->user_data) &&
-			!mainloop_modify_timeout(data->id, data->timeout))
-		return;
+	/* The callback may remove the timeout. */
+	data->ref_count++;
+	repeat = data->func(data->user_data);
+
+	if (data->id) {
+		if (!repeat || mainloop_modify_timeout(data->id, data->timeout))
+			mainloop_remove_timeout(data->id);
+	}
 
-	mainloop_remove_timeout(data->id);
+	timeout_unref(data);
 }
 
 static void timeout_destroy(void *user_data)
 {
 	struct timeout_data *data = user_data;
 
+	data->id = 0;
+
 	if (data->destroy)
 		data->destroy(data->user_data);
 
-	free(data);
+	timeout_unref(data);
 }
 
 unsigned int timeout_add(unsigned int timeout, timeout_func_t func,
@@ -49,6 +64,7 @@ unsigned int timeout_add(unsigned int timeout, timeout_func_t func,
 	struct timeout_data *data;
 
 	data = new0(struct timeout_data, 1);
+	data->ref_count = 1;
 	data->func = func;
 	data->user_data = user_data;
 	data->timeout = timeout;
diff --git a/unit/test-mainloop.c b/unit/test-mainloop.c
index 70ee03ee42..4cff654147 100644
--- a/unit/test-mainloop.c
+++ b/unit/test-mainloop.c
@@ -12,6 +12,7 @@
 #include <stdlib.h>
 
 #include "src/shared/mainloop.h"
+#include "src/shared/timeout.h"
 
 
 #define test_failed_msg(msg) \
@@ -153,6 +154,104 @@ static void test_timeout(struct test_data *data)
 	check(data->destroys == 1);
 }
 
+struct wrapper_data {
+	unsigned int id;
+	unsigned int calls;
+	unsigned int destroys;
+	unsigned int replacement_calls;
+	unsigned int replacement_destroys;
+	unsigned int interval;
+	bool repeat;
+	bool remove;
+	bool replace;
+};
+
+static void wrapper_destroy(void *user_data)
+{
+	struct wrapper_data *data = user_data;
+
+	data->destroys++;
+}
+
+static void replacement_destroy(void *user_data)
+{
+	struct wrapper_data *data = user_data;
+
+	data->replacement_destroys++;
+}
+
+static bool replacement_callback(void *user_data)
+{
+	struct wrapper_data *data = user_data;
+
+	data->replacement_calls++;
+	mainloop_exit_success();
+	return false;
+}
+
+static bool wrapper_callback(void *user_data)
+{
+	struct wrapper_data *data = user_data;
+	unsigned int replacement;
+
+	data->calls++;
+
+	if (data->remove) {
+		timeout_remove(data->id);
+
+		if (data->interval)
+			check(data->destroys == 1);
+
+		if (data->replace) {
+			replacement = timeout_add(1, replacement_callback,
+						  data, replacement_destroy);
+			check(replacement > 0);
+
+			if (data->interval)
+				check(replacement == data->id);
+		} else {
+			mainloop_exit_success();
+		}
+
+		return data->repeat;
+	}
+
+	if (data->repeat && data->calls < 3)
+		return true;
+
+	mainloop_exit_success();
+	return false;
+}
+
+static void test_timeout_wrapper(void)
+{
+	unsigned int interval, mode;
+
+	for (interval = 0; interval < 2; interval++) {
+		for (mode = 0; mode < 6; mode++) {
+			struct wrapper_data data = {
+				.interval = interval,
+				.repeat = mode & 1,
+				.remove = mode >= 2,
+				.replace = mode >= 4,
+			};
+
+			mainloop_init();
+			data.id = timeout_add(interval, wrapper_callback,
+					      &data, wrapper_destroy);
+			check(data.id > 0);
+			check(mainloop_run() == EXIT_SUCCESS);
+			check(data.calls ==
+					(!data.remove && data.repeat ? 3 : 1));
+			check(data.destroys == 1);
+			check(data.replacement_calls == data.replace);
+			check(data.replacement_destroys == data.replace);
+			fprintf(stderr, "Timeout wrapper %u/%u passed\n",
+				interval, mode);
+		}
+	}
+}
+
 #define define_test(name, function)				\
 	do {							\
 		static struct test_data data = {};		\
@@ -167,5 +266,6 @@ int main(int argc, char *argv[])
 	define_test("Test Modify File Descriptor", test_modify_fd);
 	define_test("Test Re-Add Other Callback", test_readd_other_callback);
 	define_test("Test Timeout", test_timeout);
+	test_timeout_wrapper();
 	return 0;
 }

base-commit: 34b6f8e0e998728ce9e8e64e280d9d390c511804
-- 
2.56.0