From 94226d58fbdfb8fdacefbfa3b0f84b9e48578e34 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sat, 7 Oct 2017 18:20:47 +0100
Subject: [PATCH] [Minor] Add redis backend

---
 src/plugins/lua/reputation.lua | 96 ++++++++++++++++++++++++++++++----
 1 file changed, 87 insertions(+), 9 deletions(-)

diff --git a/src/plugins/lua/reputation.lua b/src/plugins/lua/reputation.lua
index e0ee5f268..d19d0da94 100644
--- a/src/plugins/lua/reputation.lua
+++ b/src/plugins/lua/reputation.lua
@@ -38,11 +38,15 @@ local default_expiry = 864000 -- 10 day by default
 -- Selectors are used to extract reputation tokens
 local ip_selector = {
   config = {
-    actions = { -- how each action is treated in scoring
-      ['reject'] = 1.0,
-      ['add header'] = 0.25,
-      ['rewrite subject'] = 0.25,
-      ['no action'] = 1.0
+    -- keys map between actions and hash elements in bucket,
+    -- h is for ham,
+    -- s is for spam,
+    -- p is for probable spam
+    keys_map = {
+      ['reject'] = 's',
+      ['add header'] = 'p',
+      ['rewrite subject'] = 'p',
+      ['no action'] = 'h'
     },
     scores = { -- how each component is evaluated
       ['asn'] = 0.4,
@@ -51,7 +55,6 @@ local ip_selector = {
       ['ip'] = 1.0
     },
     symbol = 'IP_SCORE', -- symbol to be inserted
-    hash = 'ip_score', -- hash table in redis used for storing scores
     asn_suffix = 'a:', -- prefix for ASN hashes
     country_suffix = 'c:', -- prefix for country hashes
     ipnet_suffix = 'n:', -- prefix for ipnet hashes
@@ -172,10 +175,84 @@ local function reputation_dns_get_token(task, rule, token, continuation_cb)
   })
 end
 
-local function reputation_redis_get_token(task, token, continuation_cb)
+local function reputation_redis_get_token(task, rule, token, continuation_cb)
+  local key = gen_token_key(token)
+
+  local function redis_get_cb(err, data)
+    if data then
+      if type(data) == 'table' then
+        local values = {}
+        for i=1,#data,2 do
+          local ndata = tonumber(data[i + 1])
+          if ndata then
+            values[data[i]] = ndata
+          end
+        end
+        continuation_cb(nil, key, values)
+      else
+        rspamd_logger.errx(task, 'invalid type while getting reputation keys %s: %s',
+          key, type(data))
+        continuation_cb("invalid type", key, nil)
+      end
+
+    elseif err then
+      rspamd_logger.errx(task, 'got error while getting reputation keys %s: %s',
+        key, err)
+      continuation_cb(err, key, nil)
+    else
+      continuation_cb("unknown error", key, nil)
+    end
+  end
+
+  local ret = rspamd_redis_make_request(task,
+    redis_params, -- connect params
+    key, -- hash key
+    false, -- is write
+    redis_get_cb, --callback
+    'HGETALL', -- command
+    {key} -- arguments
+  )
+  if not ret then
+    rspamd_logger.errx(task, 'cannot make redis request to check results')
+  end
 end
 
-local function reputation_redis_set_token(task, token, values, continuation_cb)
+local function reputation_redis_set_token(task, rule, token, values, continuation_cb)
+  local key = gen_token_key(token)
+
+  local ret,conn,upstream
+
+  local function redis_set_cb(err, data)
+    if err then
+      rspamd_logger.errx(task, 'got error while setting reputation keys %s: %s',
+        key, err)
+      continuation_cb(err, key)
+    else
+      continuation_cb(nil, key)
+    end
+  end
+
+  -- We start from expiry update
+  ret,conn,upstream = rspamd_redis_make_request(task,
+    redis_params, -- connect params
+    nil, -- hash key
+    true, -- is write
+    redis_set_cb, --callback
+    'EXPIRE', -- command
+    {key, tostring(rule.backend.config.expiry)} -- arguments
+  )
+  -- Update greylisting record expire
+  if ret then
+    -- Here, we increment all hash keys that are listed in values
+    -- E.g. {'s': 1.0} or {'h': -1.0}, floating point allows better flexibility
+    fun.each(function(k, v)
+      conn:add_cmd('HINCRBYFLOAT', {key, tostring(k), tostring(v)})
+    end, values)
+    -- Add last modification time (might be not very consistent between updates)
+    conn:add_cmd('HSET', {key, 'last', tostring(rspamd_util:get_calendar_ticks())})
+  else
+    rspamd_logger.errx(task, 'got error while connecting to redis')
+  end
 end
 
 --[[ Backends are responsible for getting reputation tokens
@@ -195,6 +272,7 @@ local backends = {
   },
   dns = {
     config = {
+      -- list = rep.example.com
     },
     get_token = reputation_dns_get_token,
     -- No set token for DNS
@@ -335,7 +413,7 @@ local function parse_rule(name, tbl)
 end
 
 redis_params = rspamd_parse_redis_server('reputation')
-local opts = rspamd_config:get_all_opt("fann_redis")
+local opts = rspamd_config:get_all_opt("reputation")
 
 -- Initialization part
 if not (opts and type(opts) == 'table') then
-- 
2.39.5