[Fix] Add logging, preallocation and hash mixing to task registry

- Log error when detecting use-after-free attempt on task pointer
- Preallocate task set to 16 elements to reduce early rehashing
- Mix pointer bits using multiplicative hash for better distribution

diff --git a/src/libserver/task.c b/src/libserver/task.c
index c722c73be1ff4df229f997dc80fa60d09ac7d50b..4c4341eb5743d692152ba323fb91a76110c32e8b 100644 (file)
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -51,15 +51,33 @@ __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 pointer set for validating Lua task references */
+/*
+ * Task pointer set for validating Lua task references.
+ * Mix pointer bits to improve hash distribution since pointers
+ * are typically aligned and have predictable low bits.
+ */
+static inline uint64_t
+rspamd_task_hash_ptr(struct rspamd_task *task)
+{
+       uint64_t p = (uint64_t) (uintptr_t) task;
+       /* Mix bits: multiply by golden ratio prime and xor with shifted value */
+       p ^= p >> 33;
+       p *= 0xff51afd7ed558ccdULL;
+       p ^= p >> 33;
+       return p;
+}
+
 KHASH_SET_INIT_INT64(rspamd_task_set);
 
 static khash_t(rspamd_task_set) *task_registry = NULL;
 
+#define TASK_REGISTRY_INITIAL_SIZE 16
+
 void rspamd_task_registry_init(void)
 {
        if (task_registry == NULL) {
                task_registry = kh_init(rspamd_task_set);
+               kh_resize(rspamd_task_set, task_registry, TASK_REGISTRY_INITIAL_SIZE);
        }
 }
 
@@ -78,7 +96,7 @@ rspamd_task_is_valid(struct rspamd_task *task)
                return FALSE;
        }
 
-       khiter_t k = kh_get(rspamd_task_set, task_registry, (uint64_t) (uintptr_t) task);
+       khiter_t k = kh_get(rspamd_task_set, task_registry, rspamd_task_hash_ptr(task));
        return k != kh_end(task_registry);
 }
 
@@ -90,7 +108,7 @@ rspamd_task_registry_add(struct rspamd_task *task)
        }
 
        int ret;
-       kh_put(rspamd_task_set, task_registry, (uint64_t) (uintptr_t) task, &ret);
+       kh_put(rspamd_task_set, task_registry, rspamd_task_hash_ptr(task), &ret);
 }
 
 static inline void
@@ -100,7 +118,7 @@ rspamd_task_registry_remove(struct rspamd_task *task)
                return;
        }
 
-       khiter_t k = kh_get(rspamd_task_set, task_registry, (uint64_t) (uintptr_t) task);
+       khiter_t k = kh_get(rspamd_task_set, task_registry, rspamd_task_hash_ptr(task));
        if (k != kh_end(task_registry)) {
                kh_del(rspamd_task_set, task_registry, k);
        }

diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 92b588596c4e7ba087dc3be6cb4b96b243d5f0c5..d47575aa35f4b0455c0379a180fe4e996e854e44 100644 (file)
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -1497,6 +1497,7 @@ lua_check_task(lua_State *L, int pos)
                if (rspamd_task_is_valid(task)) {
                        return task;
                }
+               msg_err("detected use-after-free for task %p", task);
        }
        return NULL;
 }
@@ -1511,6 +1512,7 @@ lua_check_task_maybe(lua_State *L, int pos)
                if (rspamd_task_is_valid(task)) {
                        return task;
                }
+               msg_err("detected use-after-free for task %p", task);
        }
        return NULL;
 }