Add expire and whitelist options to ip_score plugin.

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

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