[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 d07c3c7f9d473e7b068c43e7ebd8f322cc1f5316..c772c1f57883a2753dd88642a1b6a9da9763ef36 100644 (file)
--- 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 da37e4af067e9a2f70bdd5d51dcb582a8b008c3f..27b15b3a63f42e449c515cf4e8dd482887a8d3a8 100644 (file)
--- 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 29f9781f27ce3c803a645daadd7a486daf98dd91..b3daa5ab6e6b3d40164c9f8f3dc91ad186e81796 100644 (file)
--- 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 375e354f81add87e2bf02e6d3f550cd5c7c21654..2cee364c7ec9f3479a6df0cf28999209b23da8df 100644 (file)
--- 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;
 }