From d78ac2d1060c2b049b4eebc176d3823f82f9ea76 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Mon, 2 Nov 2009 19:37:06 +0300 Subject: [PATCH] * Add ability to add maps from lua scripts and access theirs elements * Add whitelist module for whitelisting score for some ip/from addresses --- src/lua/lua_common.c | 2 + src/lua/lua_common.h | 2 + src/lua/lua_config.c | 159 +++++++++++++++++++++++++++++++++- src/lua/lua_task.c | 36 ++++++++ src/main.c | 36 ++++---- src/plugins/lua/whitelist.lua | 73 ++++++++++++++++ 6 files changed, 289 insertions(+), 19 deletions(-) create mode 100644 src/plugins/lua/whitelist.lua diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c index 05829db97..7e7caefda 100644 --- a/src/lua/lua_common.c +++ b/src/lua/lua_common.c @@ -191,6 +191,8 @@ init_lua () (void)luaopen_logger (L); (void)luaopen_config (L); (void)luaopen_metric (L); + (void)luaopen_radix (L); + (void)luaopen_hash_table (L); (void)luaopen_task (L); (void)luaopen_textpart (L); (void)luaopen_message (L); diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h index ab06166b3..9cd966e23 100644 --- a/src/lua/lua_common.h +++ b/src/lua/lua_common.h @@ -24,6 +24,8 @@ int luaopen_message (lua_State *L); int luaopen_task (lua_State *L); int luaopen_config (lua_State *L); int luaopen_metric (lua_State *L); +int luaopen_radix (lua_State *L); +int luaopen_hash_table (lua_State *L); int luaopen_textpart (lua_State *L); void init_lua_filters (struct config_file *cfg); diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c index c25103f77..defc9508f 100644 --- a/src/lua/lua_config.c +++ b/src/lua/lua_config.c @@ -25,18 +25,24 @@ #include "lua_common.h" #include "../expressions.h" +#include "../map.h" +#include "../radix.h" /* Config file methods */ LUA_FUNCTION_DEF (config, get_module_opt); LUA_FUNCTION_DEF (config, get_metric); LUA_FUNCTION_DEF (config, get_all_opt); LUA_FUNCTION_DEF (config, register_function); +LUA_FUNCTION_DEF (config, add_radix_map); +LUA_FUNCTION_DEF (config, add_hash_map); static const struct luaL_reg configlib_m[] = { LUA_INTERFACE_DEF (config, get_module_opt), LUA_INTERFACE_DEF (config, get_metric), LUA_INTERFACE_DEF (config, get_all_opt), LUA_INTERFACE_DEF (config, register_function), + LUA_INTERFACE_DEF (config, add_radix_map), + LUA_INTERFACE_DEF (config, add_hash_map), {"__tostring", lua_class_tostring}, {NULL, NULL} }; @@ -50,6 +56,24 @@ static const struct luaL_reg metriclib_m[] = { {NULL, NULL} }; +/* Radix tree */ +LUA_FUNCTION_DEF (radix, get_key); + +static const struct luaL_reg radixlib_m[] = { + LUA_INTERFACE_DEF (radix, get_key), + {"__tostring", lua_class_tostring}, + {NULL, NULL} +}; + +/* Hash table */ +LUA_FUNCTION_DEF (hash_table, get_key); + +static const struct luaL_reg hashlib_m[] = { + LUA_INTERFACE_DEF (hash_table, get_key), + {"__tostring", lua_class_tostring}, + {NULL, NULL} +}; + static struct config_file * lua_check_config (lua_State * L) { @@ -66,6 +90,22 @@ lua_check_metric (lua_State * L) return *((struct metric **)ud); } +static radix_tree_t * +lua_check_radix (lua_State * L) +{ + void *ud = luaL_checkudata (L, 1, "rspamd{radix}"); + luaL_argcheck (L, ud != NULL, 1, "'radix' expected"); + return **((radix_tree_t ***)ud); +} + +static GHashTable * +lua_check_hash_table (lua_State * L) +{ + void *ud = luaL_checkudata (L, 1, "rspamd{hash_table}"); + luaL_argcheck (L, ud != NULL, 1, "'hash_table' expected"); + return **((GHashTable ***)ud); +} + /*** Config functions ***/ static int lua_config_get_module_opt (lua_State * L) @@ -205,6 +245,66 @@ lua_config_register_function (lua_State *L) return 0; } +static int +lua_config_add_radix_map (lua_State *L) +{ + struct config_file *cfg = lua_check_config (L); + const char *map_line; + radix_tree_t **r, ***ud; + + if (cfg) { + map_line = luaL_checkstring (L, 2); + r = g_malloc (sizeof (radix_tree_t *)); + *r = radix_tree_create (); + if (!add_map (map_line, read_radix_list, fin_radix_list, (void **)r)) { + msg_warn ("add_radix_map: invalid radix map %s", map_line); + radix_tree_free (*r); + g_free (r); + lua_pushnil (L); + return 1; + } + ud = lua_newuserdata (L, sizeof (radix_tree_t *)); + *ud = r; + lua_setclass (L, "rspamd{radix}", -1); + + return 1; + } + + lua_pushnil (L); + return 1; + +} + +static int +lua_config_add_hash_map (lua_State *L) +{ + struct config_file *cfg = lua_check_config (L); + const char *map_line; + GHashTable **r, ***ud; + + if (cfg) { + map_line = luaL_checkstring (L, 2); + r = g_malloc (sizeof (GHashTable *)); + *r = g_hash_table_new (rspamd_strcase_hash, rspamd_strcase_equal); + if (!add_map (map_line, read_host_list, fin_host_list, (void **)r)) { + msg_warn ("add_radix_map: invalid hash map %s", map_line); + g_hash_table_destroy (*r); + g_free (r); + lua_pushnil (L); + return 1; + } + ud = lua_newuserdata (L, sizeof (GHashTable *)); + *ud = r; + lua_setclass (L, "rspamd{hash_table}", -1); + + return 1; + } + + lua_pushnil (L); + return 1; + +} + /*** Metric functions ***/ @@ -240,9 +340,48 @@ lua_metric_register_symbol (lua_State * L) cd = g_malloc (sizeof (struct lua_callback_data)); cd->name = g_strdup (callback); cd->L = L; - register_symbol (&metric->cache, cd->name, weight, lua_metric_symbol_callback, cd); + register_symbol (&metric->cache, name, weight, lua_metric_symbol_callback, cd); + } + } + return 1; +} + +/* Radix and hash table functions */ +static int +lua_radix_get_key (lua_State * L) +{ + radix_tree_t *radix = lua_check_radix (L); + uint32_t key; + + if (radix) { + key = luaL_checkint (L, 2); + + if (radix32tree_find (radix, key) != RADIX_NO_VALUE) { + lua_pushboolean (L, 1); + return 1; + } + } + + lua_pushboolean (L, 0); + return 1; +} + +static int +lua_hash_table_get_key (lua_State * L) +{ + GHashTable *tbl = lua_check_hash_table (L); + const char *key; + + if (tbl) { + key = luaL_checkstring (L, 2); + + if (g_hash_table_lookup (tbl, key) != NULL) { + lua_pushboolean (L, 1); + return 1; } } + + lua_pushboolean (L, 0); return 1; } @@ -263,3 +402,21 @@ luaopen_metric (lua_State * L) return 1; } + +int +luaopen_radix (lua_State * L) +{ + lua_newclass (L, "rspamd{radix}", radixlib_m); + luaL_openlib (L, "rspamd_radix", null_reg, 0); + + return 1; +} + +int +luaopen_hash_table (lua_State * L) +{ + lua_newclass (L, "rspamd{hash_table}", hashlib_m); + luaL_openlib (L, "rspamd_hash_table", null_reg, 0); + + return 1; +} diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c index 7912b7de0..cfacf3f2f 100644 --- a/src/lua/lua_task.c +++ b/src/lua/lua_task.c @@ -41,6 +41,8 @@ LUA_FUNCTION_DEF (task, call_rspamd_function); LUA_FUNCTION_DEF (task, get_recipients); LUA_FUNCTION_DEF (task, get_from); LUA_FUNCTION_DEF (task, get_from_ip); +LUA_FUNCTION_DEF (task, get_from_ip_num); +LUA_FUNCTION_DEF (task, get_client_ip_num); LUA_FUNCTION_DEF (task, get_helo); static const struct luaL_reg tasklib_m[] = { @@ -56,6 +58,8 @@ static const struct luaL_reg tasklib_m[] = { LUA_INTERFACE_DEF (task, get_recipients), LUA_INTERFACE_DEF (task, get_from), LUA_INTERFACE_DEF (task, get_from_ip), + LUA_INTERFACE_DEF (task, get_from_ip_num), + LUA_INTERFACE_DEF (task, get_client_ip_num), LUA_INTERFACE_DEF (task, get_helo), {"__tostring", lua_class_tostring}, {NULL, NULL} @@ -436,6 +440,38 @@ lua_task_get_from_ip (lua_State *L) return 1; } +static int +lua_task_get_from_ip_num (lua_State *L) +{ + struct worker_task *task = lua_check_task (L); + + if (task) { + if (task->from_addr.s_addr != 0) { + lua_pushinteger (L, ntohl (task->from_addr.s_addr)); + return 1; + } + } + + lua_pushnil (L); + return 1; +} + +static int +lua_task_get_client_ip_num (lua_State *L) +{ + struct worker_task *task = lua_check_task (L); + + if (task) { + if (task->client_addr.s_addr != 0) { + lua_pushinteger (L, ntohl (task->client_addr.s_addr)); + return 1; + } + } + + lua_pushnil (L); + return 1; +} + static int lua_task_get_helo (lua_State *L) { diff --git a/src/main.c b/src/main.c index 4b811bb36..4a02fa159 100644 --- a/src/main.c +++ b/src/main.c @@ -674,24 +674,6 @@ main (int argc, char **argv, char **env) exit (-errno); } -#ifndef WITHOUT_PERL - /* Init perl interpreter */ - dTHXa (perl_interpreter); - PERL_SYS_INIT3 (&argc, &argv, &env); - perl_interpreter = perl_alloc (); - if (perl_interpreter == NULL) { - msg_err ("main: cannot allocate perl interpreter, %s", strerror (errno)); - exit (-errno); - } - - PERL_SET_CONTEXT (perl_interpreter); - perl_construct (perl_interpreter); - perl_parse (perl_interpreter, xs_init, 3, args, NULL); - init_perl_filters (cfg); -#elif defined(WITH_LUA) - init_lua_filters (cfg); -#endif - /* Block signals to use sigsuspend in future */ sigprocmask (SIG_BLOCK, &signals.sa_mask, NULL); @@ -726,6 +708,24 @@ main (int argc, char **argv, char **env) l = g_list_next (l); } +#ifndef WITHOUT_PERL + /* Init perl interpreter */ + dTHXa (perl_interpreter); + PERL_SYS_INIT3 (&argc, &argv, &env); + perl_interpreter = perl_alloc (); + if (perl_interpreter == NULL) { + msg_err ("main: cannot allocate perl interpreter, %s", strerror (errno)); + exit (-errno); + } + + PERL_SET_CONTEXT (perl_interpreter); + perl_construct (perl_interpreter); + perl_parse (perl_interpreter, xs_init, 3, args, NULL); + init_perl_filters (cfg); +#elif defined(WITH_LUA) + init_lua_filters (cfg); +#endif + rspamd->workers = g_hash_table_new (g_direct_hash, g_direct_equal); spawn_workers (rspamd, TRUE); diff --git a/src/plugins/lua/whitelist.lua b/src/plugins/lua/whitelist.lua new file mode 100644 index 000000000..e95aca6c7 --- /dev/null +++ b/src/plugins/lua/whitelist.lua @@ -0,0 +1,73 @@ +-- Module that add symbols to those hosts or from domains that are contained in whitelist + +local metric = 'default' +local symbol_ip = nil +local symbol_from = nil + +local r = nil +local h = nil -- radix tree and hash table + +function check_whitelist (task) + if symbol_ip then + -- check client's ip + local ipn = task:get_from_ip_num() + if ipn then + local key = r:get_key(ipn) + if key then + task:insert_result(metric, symbol_ip, 1) + end + end + end + + if symbol_from then + -- check client's from domain + local from = task:get_from() + if from then + local _,_,domain = string.find(from, '@(.+)>?$') + local key = h:get_key(domain) + if key then + task:insert_result(metric, symbol_from, 1) + end + end + end + +end + + +-- Configuration +local opts = rspamd_config:get_all_opt('whitelist') +if opts then + if opts['symbol_ip'] or opts['symbol_from'] then + symbol_ip = opts['symbol_ip'] + symbol_from = opts['symbol_from'] + + if symbol_ip then + if opts['ip_whitelist'] then + r = rspamd_config:add_radix_map (opts['ip_whitelist']) + else + -- No whitelist defined + symbol_ip = nil + end + end + if symbol_from then + if opts['from_whitelist'] then + h = rspamd_config:add_host_map (opts['from_whitelist']) + else + -- No whitelist defined + symbol_from = nil + end + end + + if opts['metric'] then + metric = opts['metric'] + end + + -- Register symbol's callback + local m = rspamd_config:get_metric(metric) + if symbol_ip then + m:register_symbol(symbol_ip, 1.0, 'check_whitelist') + elseif symbol_from then + m:register_symbol(symbol_from, 1.0, 'check_whitelist') + end + end +end -- 2.39.5