Add expire and whitelist options to ip_score plugin.

Some fixes to expiration of keys (still need to be reworked however).

diff --git a/src/kvstorage.c b/src/kvstorage.c
index 76e58adbaebe84791a2e90505b45a8cac3ba9a20..76def504e9f1a77011d09d6cc4f51b117662f0ca 100644 (file)
--- 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 67be01fcf022edbd2f0c6a5d686cc3fc400183c9..d32db2ee4231e6a943833c6c63c52e750bfe66b2 100644 (file)
--- 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 c427e60ed68a014bbd1e5690ea2e992d32239b11..4d2469e5f943779e0dd3f4e0fdfe858fd883cb6c 100644 (file)
--- 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 1fd7cc0bf719c91b3866d76f70a942e369665071..4c822c8c7b01a14043566831fcf83968fcd2fcbe 100644 (file)
--- 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