From: Vsevolod Stakhov <vsevolod@rspamd.com>
Date: Sun, 21 Dec 2025 20:05:27 +0000 (+0000)
Subject: [Feature] Add task registry for safe Lua task reference validation
X-Git-Tag: 3.14.3~32^2~4
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=584c2fdfe6720543ae8be9b6de5b6960ba0e6bdb;p=thirdparty%2Frspamd.git

[Feature] Add task registry for safe Lua task reference validation

Implement a global task registry that maps unique uint64_t keys to task
pointers. This prevents use-after-free bugs when Lua code holds references
to tasks that may have been freed (e.g., in async Redis callbacks).

Key changes:
- Add lua_key field to rspamd_task struct
- Implement task registry using khash (O(1) lookup)
- Store lua_key in Lua userdata instead of raw pointer
- Lookup via registry when extracting task from Lua
- Remove task from registry FIRST in rspamd_task_free()

The counter-based key approach avoids issues with:
- Pointer reuse after free (memory allocator may reuse addresses)
- Lua number precision (52-bit mantissa is sufficient for counter)
- NaN/subnormal float values that could cause issues

This fixes potential use-after-free in Redis script waitq callbacks
when Redis is unavailable longer than task lifetime.
---

diff --git a/src/libserver/cfg_utils.cxx b/src/libserver/cfg_utils.cxx
index d07c3c7f9d..c772c1f578 100644
--- a/src/libserver/cfg_utils.cxx
+++ b/src/libserver/cfg_utils.cxx
@@ -2753,6 +2753,7 @@ rspamd_init_libs(void)
 
 	auto *ctx = g_new0(struct rspamd_external_libs_ctx, 1);
 	ctx->crypto_ctx = rspamd_cryptobox_init();
+	rspamd_task_registry_init();
 	ottery_cfg = (struct ottery_config *) g_malloc0(ottery_get_sizeof_config());
 	ottery_config_init(ottery_cfg);
 	ctx->ottery_cfg = ottery_cfg;
@@ -3069,6 +3070,7 @@ void rspamd_deinit_libs(struct rspamd_external_libs_ctx *ctx)
 		}
 
 		rspamd_cryptobox_deinit(ctx->crypto_ctx);
+		rspamd_task_registry_destroy();
 
 		g_free(ctx);
 	}
diff --git a/src/libserver/task.c b/src/libserver/task.c
index da37e4af06..27b15b3a63 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -51,6 +51,74 @@ __KHASH_IMPL(rspamd_req_headers_hash, static inline,
 			 rspamd_ftok_t *, struct rspamd_request_header_chain *, 1,
 			 rspamd_ftok_icase_hash, rspamd_ftok_icase_equal)
 
+/* Task registry: maps lua_key -> task pointer for safe Lua references */
+KHASH_INIT(rspamd_task_registry, uint64_t, struct rspamd_task *, 1,
+		   kh_int64_hash_func, kh_int64_hash_equal);
+
+static khash_t(rspamd_task_registry) *task_registry = NULL;
+static uint64_t task_lua_key_counter = 0;
+
+void rspamd_task_registry_init(void)
+{
+	if (task_registry == NULL) {
+		task_registry = kh_init(rspamd_task_registry);
+	}
+}
+
+void rspamd_task_registry_destroy(void)
+{
+	if (task_registry != NULL) {
+		kh_destroy(rspamd_task_registry, task_registry);
+		task_registry = NULL;
+	}
+}
+
+struct rspamd_task *
+rspamd_task_by_lua_key(uint64_t lua_key)
+{
+	if (task_registry == NULL || lua_key == 0) {
+		return NULL;
+	}
+
+	khiter_t k = kh_get(rspamd_task_registry, task_registry, lua_key);
+	if (k != kh_end(task_registry)) {
+		return kh_value(task_registry, k);
+	}
+
+	return NULL;
+}
+
+static inline void
+rspamd_task_registry_add(struct rspamd_task *task)
+{
+	if (task_registry == NULL) {
+		rspamd_task_registry_init();
+	}
+
+	task->lua_key = ++task_lua_key_counter;
+
+	int ret;
+	khiter_t k = kh_put(rspamd_task_registry, task_registry, task->lua_key, &ret);
+	if (ret > 0) {
+		kh_value(task_registry, k) = task;
+	}
+}
+
+static inline void
+rspamd_task_registry_remove(struct rspamd_task *task)
+{
+	if (task_registry == NULL || task->lua_key == 0) {
+		return;
+	}
+
+	khiter_t k = kh_get(rspamd_task_registry, task_registry, task->lua_key);
+	if (k != kh_end(task_registry)) {
+		kh_del(rspamd_task_registry, task_registry, k);
+	}
+
+	task->lua_key = 0;
+}
+
 static GQuark
 rspamd_task_quark(void)
 {
@@ -125,6 +193,8 @@ rspamd_task_new(struct rspamd_worker *worker,
 	new_task->mail_esmtp_args = NULL;
 	new_task->rcpt_esmtp_args = NULL;
 
+	rspamd_task_registry_add(new_task);
+
 	return new_task;
 }
 
@@ -183,6 +253,8 @@ void rspamd_task_free(struct rspamd_task *task)
 	unsigned int i;
 
 	if (task) {
+		rspamd_task_registry_remove(task);
+
 		debug_task("free pointer %p", task);
 
 		if (task->rcpt_envelope) {
diff --git a/src/libserver/task.h b/src/libserver/task.h
index 29f9781f27..b3daa5ab6e 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -168,6 +168,7 @@ KHASH_INIT(rspamd_task_lua_cache, char *, struct rspamd_lua_cached_entry, 1, kh_
  */
 struct rspamd_task {
 	struct rspamd_worker *worker; /**< pointer to worker object						*/
+	uint64_t lua_key;             /**< unique key for Lua task registry				*/
 	enum rspamd_command cmd;      /**< command										*/
 	int sock;                     /**< socket descriptor								*/
 	uint32_t dns_requests;        /**< number of DNS requests per this task			*/
@@ -416,6 +417,13 @@ void rspamd_task_timeout(EV_P_ ev_timer *w, int revents);
  */
 void rspamd_worker_guard_handler(EV_P_ ev_io *w, int revents);
 
+/*
+ * Task registry for safe Lua task references
+ */
+void rspamd_task_registry_init(void);
+void rspamd_task_registry_destroy(void);
+struct rspamd_task *rspamd_task_by_lua_key(uint64_t lua_key);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 375e354f81..2cee364c7e 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -1492,7 +1492,11 @@ lua_check_task(lua_State *L, int pos)
 {
 	void *ud = rspamd_lua_check_udata(L, pos, rspamd_task_classname);
 	luaL_argcheck(L, ud != NULL, pos, "'task' expected");
-	return ud ? *((struct rspamd_task **) ud) : NULL;
+	if (ud) {
+		uint64_t lua_key = *((uint64_t *) ud);
+		return rspamd_task_by_lua_key(lua_key);
+	}
+	return NULL;
 }
 
 struct rspamd_task *
@@ -1500,7 +1504,11 @@ lua_check_task_maybe(lua_State *L, int pos)
 {
 	void *ud = rspamd_lua_check_udata_maybe(L, pos, rspamd_task_classname);
 
-	return ud ? *((struct rspamd_task **) ud) : NULL;
+	if (ud) {
+		uint64_t lua_key = *((uint64_t *) ud);
+		return rspamd_task_by_lua_key(lua_key);
+	}
+	return NULL;
 }
 
 static struct rspamd_image *
@@ -8092,9 +8100,9 @@ void luaopen_image(lua_State *L)
 
 void rspamd_lua_task_push(lua_State *L, struct rspamd_task *task)
 {
-	struct rspamd_task **ptask;
+	uint64_t *pkey;
 
-	ptask = lua_newuserdata(L, sizeof(gpointer));
+	pkey = lua_newuserdata(L, sizeof(uint64_t));
 	rspamd_lua_setclass(L, rspamd_task_classname, -1);
-	*ptask = task;
+	*pkey = task->lua_key;
 }