Enable per-users learn cache

4 files changed, 44 insertions, 20 deletions

diff --git a/src/libstat/backends/redis_backend.c b/src/libstat/backends/redis_backend.c
index f339668ca..5d6ba7d98 100644
--- a/src/libstat/backends/redis_backend.c
+++ b/src/libstat/backends/redis_backend.c
@@ -130,32 +130,40 @@ rspamd_redis_expand_object (const gchar *pattern,
 	stcf = ctx->stcf;
 	L = task->cfg->lua_state;
 
-	if (ctx->cbref_user == -1) {
-		rcpt = rspamd_task_get_principal_recipient (task);
-	}
-	else {
-		/* Execute lua function to get userdata */
-		lua_pushcfunction (L, &rspamd_lua_traceback);
-		err_idx = lua_gettop (L);
-
-		lua_rawgeti (L, LUA_REGISTRYINDEX, ctx->cbref_user);
-		ptask = lua_newuserdata (L, sizeof (struct rspamd_task *));
-		*ptask = task;
-		rspamd_lua_setclass (L, "rspamd{task}", -1);
-
-		if (lua_pcall (L, 1, 1, err_idx) != 0) {
-			tb = lua_touserdata (L, -1);
-			msg_err_task ("call to user extraction script failed: %v", tb);
-			g_string_free (tb, TRUE);
+	if (ctx->enable_users) {
+		if (ctx->cbref_user == -1) {
+			rcpt = rspamd_task_get_principal_recipient (task);
 		}
 		else {
-			rcpt = rspamd_mempool_strdup (task->task_pool, lua_tostring (L, -1));
+			/* Execute lua function to get userdata */
+			lua_pushcfunction (L, &rspamd_lua_traceback);
+			err_idx = lua_gettop (L);
+
+			lua_rawgeti (L, LUA_REGISTRYINDEX, ctx->cbref_user);
+			ptask = lua_newuserdata (L, sizeof (struct rspamd_task *));
+			*ptask = task;
+			rspamd_lua_setclass (L, "rspamd{task}", -1);
+
+			if (lua_pcall (L, 1, 1, err_idx) != 0) {
+				tb = lua_touserdata (L, -1);
+				msg_err_task ("call to user extraction script failed: %v", tb);
+				g_string_free (tb, TRUE);
+			}
+			else {
+				rcpt = rspamd_mempool_strdup (task->task_pool, lua_tostring (L, -1));
+			}
+
+			/* Result + error function */
+			lua_pop (L, 2);
 		}
 
-		/* Result + error function */
-		lua_pop (L, 2);
+		if (rcpt) {
+			rspamd_mempool_set_variable (task->task_pool, "stat_user", rcpt, NULL);
+		}
 	}
 
+
+
 	/* Length calculation */
 	while (*p) {
 		switch (state) {
diff --git a/src/libstat/backends/sqlite3_backend.c b/src/libstat/backends/sqlite3_backend.c
index 7000226e2..ee9e53214 100644
--- a/src/libstat/backends/sqlite3_backend.c
+++ b/src/libstat/backends/sqlite3_backend.c
@@ -335,6 +335,8 @@ rspamd_sqlite3_get_user (struct rspamd_stat_sqlite3_db *db,
 
 
 	if (user != NULL) {
+		rspamd_mempool_set_variable (task->task_pool, "stat_user", user, NULL);
+
 		rc = rspamd_sqlite3_run_prstmt (task->task_pool, db->sqlite, db->prstmt,
 				RSPAMD_STAT_BACKEND_GET_USER, user, &id);
 
diff --git a/src/libstat/learn_cache/redis_cache.c b/src/libstat/learn_cache/redis_cache.c
index c1dd89fc4..424748039 100644
--- a/src/libstat/learn_cache/redis_cache.c
+++ b/src/libstat/learn_cache/redis_cache.c
@@ -160,9 +160,16 @@ rspamd_stat_cache_redis_generate_id (struct rspamd_task *task)
 	guint i;
 	guchar out[rspamd_cryptobox_HASHBYTES];
 	gchar *b32out;
+	gchar *user = NULL;
 
 	rspamd_cryptobox_hash_init (&st, NULL, 0);
 
+	user = rspamd_mempool_get_variable (task->task_pool, "stat_user");
+	/* Use dedicated hash space for per users cache */
+	if (user != NULL) {
+		rspamd_cryptobox_hash_update (&st, user, strlen (user));
+	}
+
 	for (i = 0; i < task->tokens->len; i ++) {
 		tok = g_ptr_array_index (task->tokens, i);
 		rspamd_cryptobox_hash_update (&st, tok->data, tok->datalen);
diff --git a/src/libstat/learn_cache/sqlite3_cache.c b/src/libstat/learn_cache/sqlite3_cache.c
index 7b46afefc..77a2cafef 100644
--- a/src/libstat/learn_cache/sqlite3_cache.c
+++ b/src/libstat/learn_cache/sqlite3_cache.c
@@ -187,6 +187,7 @@ rspamd_stat_cache_sqlite3_check (struct rspamd_task *task,
 	rspamd_cryptobox_hash_state_t st;
 	rspamd_ftok_t *word;
 	guchar *out;
+	gchar *user = NULL;
 	guint i, j;
 	gint rc;
 	gint64 flag;
@@ -196,6 +197,12 @@ rspamd_stat_cache_sqlite3_check (struct rspamd_task *task,
 
 		rspamd_cryptobox_hash_init (&st, NULL, 0);
 
+		user = rspamd_mempool_get_variable (task->task_pool, "stat_user");
+		/* Use dedicated hash space for per users cache */
+		if (user != NULL) {
+			rspamd_cryptobox_hash_update (&st, user, strlen (user));
+		}
+
 		for (i = 0; i < task->text_parts->len; i ++) {
 			part = g_ptr_array_index (task->text_parts, i);