From a698357f23c08c6aebcd8210d61c4b664fbc4c18 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Fri, 2 Dec 2022 22:39:39 +0000 Subject: [PATCH] [Feature] Add possibility to execute lua scripts for blocked fuzzy clients --- src/fuzzy_storage.c | 112 +++++++++++++++++++++++++++++++++----------- 1 file changed, 84 insertions(+), 28 deletions(-) diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c index 9311193b5..8c721d487 100644 --- a/src/fuzzy_storage.c +++ b/src/fuzzy_storage.c @@ -51,31 +51,6 @@ static const gchar *local_db_name = "local"; -#define msg_err_fuzzy_update(...) rspamd_default_log_function (G_LOG_LEVEL_CRITICAL, \ - session->name, session->uid, \ - RSPAMD_LOG_FUNC, \ - __VA_ARGS__) -#define msg_warn_fuzzy_update(...) rspamd_default_log_function (G_LOG_LEVEL_WARNING, \ - session->name, session->uid, \ - RSPAMD_LOG_FUNC, \ - __VA_ARGS__) -#define msg_info_fuzzy_update(...) rspamd_default_log_function (G_LOG_LEVEL_INFO, \ - session->name, session->uid, \ - RSPAMD_LOG_FUNC, \ - __VA_ARGS__) -#define msg_err_fuzzy_collection(...) rspamd_default_log_function (G_LOG_LEVEL_CRITICAL, \ - "fuzzy_collection", session->uid, \ - RSPAMD_LOG_FUNC, \ - __VA_ARGS__) -#define msg_warn_fuzzy_collection(...) rspamd_default_log_function (G_LOG_LEVEL_WARNING, \ - "fuzzy_collection", session->uid, \ - RSPAMD_LOG_FUNC, \ - __VA_ARGS__) -#define msg_info_fuzzy_collection(...) rspamd_default_log_function (G_LOG_LEVEL_INFO, \ - "fuzzy_collection", session->uid, \ - RSPAMD_LOG_FUNC, \ - __VA_ARGS__) - /* Init functions */ gpointer init_fuzzy (struct rspamd_config *cfg); void start_fuzzy (struct rspamd_worker *worker); @@ -190,6 +165,7 @@ struct rspamd_fuzzy_storage_ctx { struct rspamd_hash_map_helper *skip_hashes; gint lua_pre_handler_cbref; gint lua_post_handler_cbref; + gint lua_blacklist_cbref; }; enum fuzzy_cmd_type { @@ -239,9 +215,13 @@ struct rspamd_updates_cbdata { static void rspamd_fuzzy_write_reply (struct fuzzy_session *session); -static gboolean rspamd_fuzzy_process_updates_queue ( - struct rspamd_fuzzy_storage_ctx *ctx, - const gchar *source, gboolean final); +static gboolean rspamd_fuzzy_process_updates_queue (struct rspamd_fuzzy_storage_ctx *ctx, + const gchar *source, gboolean final); +static gboolean rspamd_fuzzy_check_client (struct rspamd_fuzzy_storage_ctx *ctx, + rspamd_inet_addr_t *addr); +static void rspamd_fuzzy_maybe_call_blacklisted (struct rspamd_fuzzy_storage_ctx *ctx, + rspamd_inet_addr_t *addr, + const gchar *reason); static gboolean rspamd_fuzzy_check_ratelimit (struct fuzzy_session *session) @@ -320,6 +300,10 @@ rspamd_fuzzy_check_ratelimit (struct fuzzy_session *session) rspamd_inet_address_free (masked); + if (ratelimited) { + rspamd_fuzzy_maybe_call_blacklisted(session->ctx, session->addr, "ratelimit"); + } + return !ratelimited; } else { @@ -339,6 +323,32 @@ rspamd_fuzzy_check_ratelimit (struct fuzzy_session *session) return TRUE; } +static void +rspamd_fuzzy_maybe_call_blacklisted (struct rspamd_fuzzy_storage_ctx *ctx, + rspamd_inet_addr_t *addr, + const gchar *reason) +{ + if (ctx->lua_blacklist_cbref != -1) { + lua_State *L = ctx->cfg->lua_state; + gint err_idx, ret; + + lua_pushcfunction (L, &rspamd_lua_traceback); + err_idx = lua_gettop (L); + lua_rawgeti (L, LUA_REGISTRYINDEX, ctx->lua_blacklist_cbref); + /* client IP */ + rspamd_lua_ip_push (L, addr); + /* block reason */ + lua_pushstring (L, reason); + + if ((ret = lua_pcall (L, 2, 0, err_idx)) != 0) { + msg_err ("call to lua_blacklist_cbref " + "script failed (%d): %s", ret, lua_tostring (L, -1)); + } + + lua_settop (L, 0); + } +} + static gboolean rspamd_fuzzy_check_client (struct rspamd_fuzzy_storage_ctx *ctx, rspamd_inet_addr_t *addr) @@ -346,6 +356,8 @@ rspamd_fuzzy_check_client (struct rspamd_fuzzy_storage_ctx *ctx, if (ctx->blocked_ips != NULL) { if (rspamd_match_radix_map_addr (ctx->blocked_ips, addr) != NULL) { + + rspamd_fuzzy_maybe_call_blacklisted (ctx, addr, "blacklisted"); return FALSE; } } @@ -2013,6 +2025,37 @@ lua_fuzzy_add_post_handler (lua_State *L) return 0; } +static int +lua_fuzzy_add_blacklist_handler (lua_State *L) +{ + struct rspamd_worker *wrk, **pwrk = (struct rspamd_worker **) + rspamd_lua_check_udata (L, 1, "rspamd{worker}"); + struct rspamd_fuzzy_storage_ctx *ctx; + + if (!pwrk) { + return luaL_error (L, "invalid arguments, worker + function are expected"); + } + + wrk = *pwrk; + + if (wrk && lua_isfunction (L, 2)) { + ctx = (struct rspamd_fuzzy_storage_ctx *)wrk->ctx; + + if (ctx->lua_blacklist_cbref != -1) { + /* Should not happen */ + luaL_unref (L, LUA_REGISTRYINDEX, ctx->lua_blacklist_cbref); + } + + lua_pushvalue (L, 2); + ctx->lua_blacklist_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, @@ -2197,6 +2240,7 @@ init_fuzzy (struct rspamd_config *cfg) ctx->keypair_cache_size = DEFAULT_KEYPAIR_CACHE_SIZE; ctx->lua_pre_handler_cbref = -1; ctx->lua_post_handler_cbref = -1; + ctx->lua_blacklist_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, @@ -2688,6 +2732,9 @@ start_fuzzy (struct rspamd_worker *worker) rspamd_srv_send_command (worker, ctx->event_loop, &srv_cmd, -1, fuzzy_peer_rep, ctx); + /* + * Extra fields available for this particular worker + */ luaL_Reg fuzzy_lua_reg = { .name = "add_fuzzy_pre_handler", .func = lua_fuzzy_add_pre_handler, @@ -2698,6 +2745,11 @@ start_fuzzy (struct rspamd_worker *worker) .func = lua_fuzzy_add_post_handler, }; rspamd_lua_add_metamethod (ctx->cfg->lua_state, "rspamd{worker}", &fuzzy_lua_reg); + fuzzy_lua_reg = (luaL_Reg){ + .name = "add_fuzzy_blacklist_handler", + .func = lua_fuzzy_add_blacklist_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); @@ -2748,6 +2800,10 @@ start_fuzzy (struct rspamd_worker *worker) luaL_unref (ctx->cfg->lua_state, LUA_REGISTRYINDEX, ctx->lua_post_handler_cbref); } + if (ctx->lua_blacklist_cbref != -1) { + luaL_unref (ctx->cfg->lua_state, LUA_REGISTRYINDEX, ctx->lua_blacklist_cbref); + } + REF_RELEASE (ctx->cfg); rspamd_log_close (worker->srv->logger); rspamd_unset_crash_handler (worker->srv); -- 2.39.5