commit d33893c: [Feature] Add support for scripts in fuzzy storage

Vsevolod Stakhov vsevolod at highsecure.ru
Wed Mar 4 20:07:08 UTC 2020


Author: Vsevolod Stakhov
Date: 2020-03-04 16:20:57 +0000
URL: https://github.com/rspamd/rspamd/commit/d33893c59fb0a22f098c98445aa8b936655382e3

[Feature] Add support for scripts in fuzzy storage

---
 src/fuzzy_storage.c | 198 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 198 insertions(+)

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index 073e66965..107cf7944 100644
--- a/src/fuzzy_storage.c
+++ b/src/fuzzy_storage.c
@@ -36,6 +36,7 @@
 #include "libutil/hash.h"
 #include "libserver/maps/map_private.h"
 #include "contrib/uthash/utlist.h"
+#include "lua/lua_common.h"
 #include "unix-std.h"
 
 #include <math.h>
@@ -180,6 +181,8 @@ struct rspamd_fuzzy_storage_ctx {
 	struct rspamd_worker *worker;
 	const ucl_object_t *skip_map;
 	struct rspamd_hash_map_helper *skip_hashes;
+	gint lua_pre_handler_cbref;
+	gint lua_post_handler_cbref;
 	guchar cookie[COOKIE_SIZE];
 };
 
@@ -783,6 +786,71 @@ rspamd_fuzzy_check_callback (struct rspamd_fuzzy_reply *result, void *ud)
 		break;
 	}
 
+	if (session->ctx->lua_pre_handler_cbref != -1) {
+		/* Start lua post handler */
+		lua_State *L = session->ctx->cfg->lua_state;
+		gint err_idx, ret;
+
+		lua_pushcfunction (L, &rspamd_lua_traceback);
+		err_idx = lua_gettop (L);
+		/* Preallocate stack (small opt) */
+		lua_checkstack (L, err_idx + 9);
+		/* function */
+		lua_rawgeti (L, LUA_REGISTRYINDEX, session->ctx->lua_pre_handler_cbref);
+		/* client IP */
+		rspamd_lua_ip_push (L, session->addr);
+		/* client command */
+		lua_pushinteger (L, cmd->cmd);
+		/* command value (push as rspamd_text) */
+		(void)lua_new_text (L, result->digest, sizeof (result->digest), FALSE);
+		/* is shingle */
+		lua_pushboolean (L, is_shingle);
+		/* result value */
+		lua_pushinteger (L, result->v1.value);
+		/* result probability */
+		lua_pushnumber (L, result->v1.prob);
+		/* result flag */
+		lua_pushinteger (L, result->v1.flag);
+		/* result timestamp */
+		lua_pushinteger (L, result->ts);
+		/* TODO: add additional data maybe (encryption, pubkey, etc) */
+
+		if ((ret = lua_pcall (L, 8, LUA_MULTRET, err_idx)) != 0) {
+			msg_err ("call to lua_post_handler lua "
+					 "script failed (%d): %s", ret, lua_tostring (L, -1));
+		}
+		else {
+			/* Return values order:
+			 * the first reply will be on err_idx + 1
+			 * if it is true, then we need to read the former ones:
+			 * 2-nd will be reply code
+			 * 3-rd will be probability (or 0.0 if missing)
+			 */
+			ret = lua_toboolean (L, err_idx + 1);
+
+			if (ret) {
+				/* Artificial reply */
+				result->v1.value = lua_tointeger (L, err_idx + 2);
+
+				if (lua_isnumber (L, err_idx + 3)) {
+					result->v1.prob = lua_tonumber (L, err_idx + 3);
+				}
+				else {
+					result->v1.prob = 0.0f;
+				}
+
+				lua_settop (L, 0);
+				rspamd_fuzzy_make_reply (cmd, result, session, encrypted, is_shingle);
+				REF_RELEASE (session);
+
+				return;
+			}
+		}
+
+		lua_settop (L, 0);
+	}
+
+
 	rspamd_fuzzy_make_reply (cmd, result, session, encrypted, is_shingle);
 
 	/* Refresh hash if found with strong confidence */
@@ -883,6 +951,62 @@ rspamd_fuzzy_process_command (struct fuzzy_session *session)
 	result.v1.flag = cmd->flag;
 	result.v1.tag = cmd->tag;
 
+	if (session->ctx->lua_pre_handler_cbref != -1) {
+		/* Start lua pre handler */
+		lua_State *L = session->ctx->cfg->lua_state;
+		gint err_idx, ret;
+
+		lua_pushcfunction (L, &rspamd_lua_traceback);
+		err_idx = lua_gettop (L);
+		/* Preallocate stack (small opt) */
+		lua_checkstack (L, err_idx + 5);
+		/* function */
+		lua_rawgeti (L, LUA_REGISTRYINDEX, session->ctx->lua_pre_handler_cbref);
+		/* client IP */
+		rspamd_lua_ip_push (L, session->addr);
+		/* client command */
+		lua_pushinteger (L, cmd->cmd);
+		/* command value (push as rspamd_text) */
+		(void)lua_new_text (L, cmd->digest, sizeof (cmd->digest), FALSE);
+		/* is shingle */
+		lua_pushboolean (L, is_shingle);
+		/* TODO: add additional data maybe (encryption, pubkey, etc) */
+
+		if ((ret = lua_pcall (L, 4,LUA_MULTRET, err_idx)) != 0) {
+			msg_err ("call to lua_pre_handler lua "
+						  "script failed (%d): %s", ret, lua_tostring (L, -1));
+		}
+		else {
+			/* Return values order:
+			 * the first reply will be on err_idx + 1
+			 * if it is true, then we need to read the former ones:
+			 * 2-nd will be reply code
+			 * 3-rd will be probability (or 0.0 if missing)
+			 */
+			ret = lua_toboolean (L, err_idx + 1);
+
+			if (ret) {
+				/* Artificial reply */
+				result.v1.value = lua_tointeger (L, err_idx + 2);
+
+				if (lua_isnumber (L, err_idx + 3)) {
+					result.v1.prob = lua_tonumber (L, err_idx + 3);
+				}
+				else {
+					result.v1.prob = 0.0f;
+				}
+
+				lua_settop (L, 0);
+				rspamd_fuzzy_make_reply (cmd, &result, session, encrypted, is_shingle);
+
+				return;
+			}
+		}
+
+		lua_settop (L, 0);
+	}
+
+
 	if (G_UNLIKELY (cmd == NULL || up_len == 0)) {
 		result.v1.value = 500;
 		result.v1.prob = 0.0f;
@@ -1496,6 +1620,56 @@ rspamd_fuzzy_stat_to_ucl (struct rspamd_fuzzy_storage_ctx *ctx, gboolean ip_stat
 	return obj;
 }
 
+static int
+lua_fuzzy_add_pre_handler (lua_State *L)
+{
+	struct rspamd_worker *wrk = *(struct rspamd_worker **)
+			rspamd_lua_check_udata (L, 1, "rspamd{worker}");
+	struct rspamd_fuzzy_storage_ctx *ctx;
+
+	if (wrk && lua_isfunction (L, 2)) {
+		ctx = (struct rspamd_fuzzy_storage_ctx *)wrk->ctx;
+
+		if (ctx->lua_pre_handler_cbref != -1) {
+			/* Should not happen */
+			luaL_unref (L, LUA_REGISTRYINDEX, ctx->lua_pre_handler_cbref);
+		}
+
+		lua_pushvalue (L, 2);
+		ctx->lua_pre_handler_cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+	}
+	else {
+		return luaL_error (L, "invalid arguments, worker + function are expected");
+	}
+
+	return 0;
+}
+
+static int
+lua_fuzzy_add_post_handler (lua_State *L)
+{
+	struct rspamd_worker *wrk = *(struct rspamd_worker **)
+			rspamd_lua_check_udata (L, 1, "rspamd{worker}");
+	struct rspamd_fuzzy_storage_ctx *ctx;
+
+	if (wrk && lua_isfunction (L, 2)) {
+		ctx = (struct rspamd_fuzzy_storage_ctx *)wrk->ctx;
+
+		if (ctx->lua_post_handler_cbref != -1) {
+			/* Should not happen */
+			luaL_unref (L, LUA_REGISTRYINDEX, ctx->lua_post_handler_cbref);
+		}
+
+		lua_pushvalue (L, 2);
+		ctx->lua_post_handler_cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+	}
+	else {
+		return luaL_error (L, "invalid arguments, worker + function are expected");
+	}
+
+	return 0;
+}
+
 static gboolean
 rspamd_fuzzy_storage_stat (struct rspamd_main *rspamd_main,
 		struct rspamd_worker *worker, gint fd,
@@ -1676,6 +1850,8 @@ init_fuzzy (struct rspamd_config *cfg)
 	ctx->magic = rspamd_fuzzy_storage_magic;
 	ctx->sync_timeout = DEFAULT_SYNC_TIMEOUT;
 	ctx->keypair_cache_size = DEFAULT_KEYPAIR_CACHE_SIZE;
+	ctx->lua_pre_handler_cbref = -1;
+	ctx->lua_post_handler_cbref = -1;
 	ctx->keys = g_hash_table_new_full (fuzzy_kp_hash, fuzzy_kp_equal,
 			NULL, fuzzy_key_dtor);
 	rspamd_mempool_add_destructor (cfg->cfg_pool,
@@ -2067,6 +2243,20 @@ start_fuzzy (struct rspamd_worker *worker)
 	rspamd_srv_send_command (worker, ctx->event_loop, &srv_cmd, -1,
 			fuzzy_peer_rep, ctx);
 
+	luaL_Reg fuzzy_lua_reg = {
+			.name = "add_fuzzy_pre_handler",
+			.func = lua_fuzzy_add_pre_handler,
+	};
+	rspamd_lua_add_metamethod (ctx->cfg->lua_state, "rspamd{worker}", &fuzzy_lua_reg);
+	fuzzy_lua_reg = (luaL_Reg){
+			.name = "add_fuzzy_post_handler",
+			.func = lua_fuzzy_add_post_handler,
+	};
+	rspamd_lua_add_metamethod (ctx->cfg->lua_state, "rspamd{worker}", &fuzzy_lua_reg);
+
+	rspamd_lua_run_postloads (ctx->cfg->lua_state, ctx->cfg, ctx->event_loop,
+			worker);
+
 	ev_loop (ctx->event_loop, 0);
 	rspamd_worker_block_signals ();
 
@@ -2104,6 +2294,14 @@ start_fuzzy (struct rspamd_worker *worker)
 		rspamd_lru_hash_destroy (ctx->ratelimit_buckets);
 	}
 
+	if (ctx->lua_pre_handler_cbref != -1) {
+		luaL_unref (ctx->cfg->lua_state, LUA_REGISTRYINDEX, ctx->lua_pre_handler_cbref);
+	}
+
+	if (ctx->lua_post_handler_cbref != -1) {
+		luaL_unref (ctx->cfg->lua_state, LUA_REGISTRYINDEX, ctx->lua_post_handler_cbref);
+	}
+
 	REF_RELEASE (ctx->cfg);
 	rspamd_log_close (worker->srv->logger);
 


More information about the Commits mailing list