From 4b1dbdccc7067ff71565b48c1e9722c152a83a4a Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Fri, 16 Dec 2011 18:37:50 +0300 Subject: [PATCH] Add expire and whitelist options to ip_score plugin. Some fixes to expiration of keys (still need to be reworked however). --- src/kvstorage.c | 12 ++++++++++-- src/kvstorage.h | 3 ++- src/lua/lua_redis.c | 3 +++ src/plugins/lua/ip_score.lua | 30 ++++++++++++++++++++++++++++-- 4 files changed, 43 insertions(+), 5 deletions(-) diff --git a/src/kvstorage.c b/src/kvstorage.c index 76e58adba..76def504e 100644 --- a/src/kvstorage.c +++ b/src/kvstorage.c @@ -341,8 +341,14 @@ rspamd_kv_storage_increment (struct rspamd_kv_storage *storage, gpointer key, gu } if (elt && (elt->flags & KV_ELT_INTEGER) != 0) { lp = &ELT_LONG (elt); - *lp += *value; - *value = *lp; + /* Handle need expire here */ + if (elt->flags & KV_ELT_NEED_EXPIRE) { + *lp = *value; + } + else { + *lp += *value; + *value = *lp; + } elt->age = time (NULL); if (storage->backend) { if (storage->backend->replace_func (storage->backend, key, keylen, elt)) { @@ -394,6 +400,8 @@ rspamd_kv_storage_lookup (struct rspamd_kv_storage *storage, gpointer key, guint if (elt && (elt->flags & KV_ELT_PERSISTENT) == 0 && elt->expire > 0) { /* Check expiration */ if (now - elt->age > elt->expire) { + /* Set need expire as we have no write lock here */ + elt->flags |= KV_ELT_NEED_EXPIRE; elt = NULL; } } diff --git a/src/kvstorage.h b/src/kvstorage.h index 67be01fcf..d32db2ee4 100644 --- a/src/kvstorage.h +++ b/src/kvstorage.h @@ -73,7 +73,8 @@ enum rspamd_kv_flags { KV_ELT_OUSTED = 1 << 3, KV_ELT_NEED_FREE = 1 << 4, KV_ELT_INTEGER = 1 << 5, - KV_ELT_NEED_INSERT = 1 << 6 + KV_ELT_NEED_INSERT = 1 << 6, + KV_ELT_NEED_EXPIRE = 1 << 7 }; #define ELT_DATA(elt) (gchar *)(elt)->data + (elt)->keylen + 1 diff --git a/src/lua/lua_redis.c b/src/lua/lua_redis.c index c427e60ed..4d2469e5f 100644 --- a/src/lua/lua_redis.c +++ b/src/lua/lua_redis.c @@ -143,6 +143,9 @@ lua_redis_push_data (const redisReply *r, struct lua_redis_userdata *ud) else if (r->type == REDIS_REPLY_STATUS) { lua_pushlstring (ud->L, r->str, r->len); } + else if (r->type == REDIS_REPLY_NIL) { + lua_pushnil (ud->L); + } else { msg_info ("bad type is passed: %d", r->type); lua_pushnil (ud->L); diff --git a/src/plugins/lua/ip_score.lua b/src/plugins/lua/ip_score.lua index 1fd7cc0bf..4c822c8c7 100644 --- a/src/plugins/lua/ip_score.lua +++ b/src/plugins/lua/ip_score.lua @@ -10,6 +10,8 @@ local no_action_score = -2 local symbol = 'IP_SCORE' -- This score is used for normalization of scores from keystorage local normalize_score = 100 +local whitelist = nil +local expire = 240 -- Set score based on metric's action local ip_score_set = function(task) @@ -26,7 +28,7 @@ local ip_score_set = function(task) end end end - rspamd_redis.make_request(task, 'localhost', 6050, cb_set, 'SET %b %b %b', ip, '100', '0') + rspamd_redis.make_request(task, keystorage_host, keystorage_port, cb_set, 'SET %b %b %b', ip, expire, '0') else rspamd_logger.info('got error while incrementing: ' .. err) end @@ -36,6 +38,15 @@ local ip_score_set = function(task) end local action = task:get_metric_action(metric) if action then + -- Check whitelist + if whitelist then + local ipnum = task:get_from_ip_num() + if ipnum and whitelist:get_key(ipnum) then + -- Address is whitelisted + return + end + end + -- Now check action if action == 'reject' then local ip = task:get_from_ip() if ip then @@ -74,7 +85,7 @@ end local ip_score_check = function(task) local cb = function(task, err, data) if err then - -- Key is not found + -- Key is not found or error occured return elseif data then local score = tonumber(data) @@ -91,6 +102,13 @@ local ip_score_check = function(task) end local ip = task:get_from_ip() if ip then + if whitelist then + local ipnum = task:get_from_ip_num() + if whitelist:get_key(ipnum) then + -- Address is whitelisted + return + end + end rspamd_redis.make_request(task, keystorage_host, keystorage_port, cb, 'GET %b', ip) end end @@ -124,6 +142,12 @@ local configure_ip_score_module = function() if opts['normalize_score'] then normalize_score = opts['normalize_score'] end + if opts['whitelist'] then + whitelist = rspamd_config:add_radix_map(opts['whitelist']) + end + if opts['expire'] then + expire = opts['expire'] + end end end @@ -137,6 +161,8 @@ if rspamd_config:get_api_version() >= 9 then rspamd_config:register_module_option('ip_score', 'no_action_score', 'int') rspamd_config:register_module_option('ip_score', 'symbol', 'string') rspamd_config:register_module_option('ip_score', 'normalize_score', 'uint') + rspamd_config:register_module_option('ip_score', 'whitelist', 'map') + rspamd_config:register_module_option('ip_score', 'expire', 'uint') configure_ip_score_module() if keystorage_host and keystorage_port and normalize_score > 0 then -- 2.39.5