Upgraded replies and known senders modules (#4895)

* Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* FIXED. Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* Added pre-test for replies set

* Update functional of replies_set

* Few changes to replies and added check for incoming mail

* Few changes in known_senders in check_known_incoming_mail_callback

* Few changes in known_senders and replies

* An attempt to write test(not tested)

* Clean up

* Clean up

* Clean up

* Added tests for replies and known_senders (all tests failed, debug required)

* Moved replies test to the 001_merged

* Cleared up code

* Few changes to replies

* Small changes in score of CHECK_INC_MAIL symbol

* Small debug in known_senders

* Plugin known_senders is fully working

* Troubleshooting replies module

* Changed symbol for check_known_incoming_mail_callback

* Added test for failed incoming mail check

* Little rework

* Rewritten test for more appropriate

* Rewritten tests for replies module. All test passed(debugging not adding to global set)

* Debugged replies module

* Replies module works and tested(needs performance improvements)

* Cleaned up code

* Improved readability and cleaned up code

* Connected auth back(Tests not working, needs user)

* Added test for incoming mail check in known senders module

* Debugged. Works normally(tested, needs to add user)

* Debug + clean up. Tested. Works. User auth required for tests

* Improved performance

* Small changes

* Changed adding to global replies set logic + improved logs messaging

* Added authenticated user to tests

* Cleaned up

* Made a few changes according to the comments on pull request

* [Rework] Added removal of extra senders and recipients in global and local replies sets

* [Minor] Small cleanup

* [Minor] Cleaned up code

* [Fix] Fixed call of incorrect function when making key

* [Rework] Reworked scripts. Added ZADD redis call for local and global replies set

* [Minor] Cleaned up code

* [Fix] Improved performance and eliminated unnecessary invocations

* [Minor] Reassigned script ids

* [Feature] Made a check for local set

* [Fix] Upgraded tests for known senders

* [Fix] Upgraded tests for known senders

* [Fix] Fixed performance of verification of local replies set

* [Minor] Cleaned up code

* [Feature] Added new test to the known_senders tests

* [Test] Ubuntu test

* [Fix] Fixing local replies test check

* [Fix] Fixed code for local replies set check(was not working in previous versions of redis)

* [Fix] Reorganized code to more convenient style and made better loading for scripts

* [Minor] Code has been rewritten in a more appropriate format

* [Minor] Fixed debug messaging

* [Fix] Reworked expiration of replies sets

* [Minor] Upgrade code style

* [Fix] Small fix

* [Feature] Change LFU logic of global replies set to LRU logic

* Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* FIXED. Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* Made the individual replies_set for senders and their recipients.
Made the global replies_set for verified recipients.

* Added pre-test for replies set

* Update functional of replies_set

* Few changes to replies and added check for incoming mail

* Few changes in known_senders in check_known_incoming_mail_callback

* Few changes in known_senders and replies

* An attempt to write test(not tested)

* Clean up

* Clean up

* Clean up

* Added tests for replies and known_senders (all tests failed, debug required)

* Moved replies test to the 001_merged

* Cleared up code

* Few changes to replies

* Small changes in score of CHECK_INC_MAIL symbol

* Small debug in known_senders

* Plugin known_senders is fully working

* Troubleshooting replies module

* Changed symbol for check_known_incoming_mail_callback

* Added test for failed incoming mail check

* Little rework

* Rewritten test for more appropriate

* Rewritten tests for replies module. All test passed(debugging not adding to global set)

* Debugged replies module

* Replies module works and tested(needs performance improvements)

* Cleaned up code

* Improved readability and cleaned up code

* Connected auth back(Tests not working, needs user)

* Added test for incoming mail check in known senders module

* Debugged. Works normally(tested, needs to add user)

* Debug + clean up. Tested. Works. User auth required for tests

* Improved performance

* Small changes

* Changed adding to global replies set logic + improved logs messaging

* Added authenticated user to tests

* Cleaned up

* Made a few changes according to the comments on pull request

* [Rework] Added removal of extra senders and recipients in global and local replies sets

* [Minor] Small cleanup

* [Minor] Cleaned up code

* [Fix] Fixed call of incorrect function when making key

* [Rework] Reworked scripts. Added ZADD redis call for local and global replies set

* [Minor] Cleaned up code

* [Fix] Improved performance and eliminated unnecessary invocations

* [Minor] Reassigned script ids

* [Feature] Made a check for local set

* [Fix] Upgraded tests for known senders

* [Fix] Upgraded tests for known senders

* [Fix] Fixed performance of verification of local replies set

* [Minor] Cleaned up code

* [Feature] Added new test to the known_senders tests

* [Test] Ubuntu test

* [Fix] Fixing local replies test check

* [Fix] Fixed code for local replies set check(was not working in previous versions of redis)

* [Fix] Reorganized code to more convenient style and made better loading for scripts

* [Minor] Code has been rewritten in a more appropriate format

* [Minor] Fixed debug messaging

* [Fix] Reworked expiration of replies sets

* [Minor] Upgrade code style

* [Fix] Small fix

* [Feature] Change LFU logic of global replies set to LRU logic

* [Fix] Fix test conflict

* [Minor] Revert rename

* [Minor] Clean up code

* [Fix] Fix commit history

2 files changed, 273 insertions, 17 deletions

diff --git a/src/plugins/lua/known_senders.lua b/src/plugins/lua/known_senders.lua
index d26a1df3b..64a28059e 100644
--- a/src/plugins/lua/known_senders.lua
+++ b/src/plugins/lua/known_senders.lua
@@ -52,7 +52,17 @@ local settings = {
   use_bloom = false,
   symbol = 'KNOWN_SENDER',
   symbol_unknown = 'UNKNOWN_SENDER',
+  symbol_check_mail_global = 'INC_MAIL_KNOWN_GLOBALLY',
+  symbol_check_mail_local = 'INC_MAIL_KNOWN_LOCALLY',
+  max_recipients = 15,
   redis_key = 'rs_known_senders',
+  sender_prefix = 'rsrk',
+  sender_key_global = 'verified_senders',
+  sender_key_size = 20,
+  reply_sender_privacy = false,
+  reply_sender_privacy_alg = 'blake2',
+  reply_sender_privacy_prefix = 'obf',
+  reply_sender_privacy_length = 16,
 }
 
 local settings_schema = lua_redis.enrich_schema({
@@ -72,6 +82,40 @@ local function make_key(input)
   return hash:hex()
 end
 
+local function make_key_replies(goop, sz, prefix)
+  local h = rspamd_cryptobox_hash.create()
+  h:update(goop)
+  local key = (prefix or '') .. h:base32():sub(1, sz)
+  return key
+end
+
+local zscore_script_id
+
+local function configure_scripts(_, _, _)
+  -- script checks if given recipients are in the local replies set of the sender
+  local redis_zscore_script = [[
+    local replies_recipients_addrs = ARGV
+    if replies_recipients_addrs then
+      for _, rcpt in ipairs(replies_recipients_addrs) do
+        local score = redis.call('ZSCORE', KEYS[1], rcpt)
+        -- check if score is nil (for some reason redis script does not see if score is a nil value)
+        if type(score) == 'boolean' then
+          score = nil
+          -- 0 is stand for failure code
+          return 0
+        end
+      end
+      -- first number in return statement is stands for the success/failure code
+      -- where success code is 1 and failure code is 0
+      return 1
+    else
+    -- 0 is a failure code
+      return 0
+    end
+  ]]
+  zscore_script_id = lua_redis.add_redis_script(redis_zscore_script, redis_params)
+end
+
 local function check_redis_key(task, key, key_ty)
   lua_util.debugm(N, task, 'check key %s, type: %s', key, key_ty)
   local function redis_zset_callback(err, data)
@@ -197,6 +241,89 @@ local function known_senders_callback(task)
   end
 end
 
+local function verify_local_replies_set(task)
+  local replies_sender = task:get_reply_sender()
+  if not replies_sender then
+    lua_util.debugm(N, task, 'Could not get sender')
+    return nil
+  end
+
+  local replies_recipients = task:get_recipients('mime')
+
+  local replies_sender_string = lua_util.maybe_obfuscate_string(tostring(replies_sender), settings, settings.sender_prefix)
+  local replies_sender_key = make_key_replies(replies_sender_string:lower(), 8)
+
+  local function redis_zscore_script_cb(err, data)
+    if err ~= nil then
+      rspamd_logger.errx(task, 'Could not verify %s local replies set %s', replies_sender_key, err)
+    end
+    if data ~= 1 then
+      lua_util.debugm(N, task, 'Recipients were not verified')
+      return
+    end
+    lua_util.debugm(N, task, 'Recipients were verified')
+    task:insert_result(settings.symbol_check_mail_local, 1.0, replies_sender_key)
+  end
+
+  local replies_recipients_addrs = {}
+  -- assigning addresses of recipients for params and limiting number of recipients to be checked
+  local max_rcpts = math.min(settings.max_recipients, #replies_recipients)
+  for i = 1, max_rcpts do
+    table.insert(replies_recipients_addrs, replies_recipients[i].addr)
+  end
+
+  lua_util.debugm(N, task, 'Making redis request to local replies set')
+  lua_redis.exec_redis_script(zscore_script_id,
+          {task = task, is_write = true},
+          redis_zscore_script_cb,
+          { replies_sender_key },
+          replies_recipients_addrs  )
+end
+
+local function check_known_incoming_mail_callback(task)
+  local replies_sender = task:get_reply_sender()
+  if not replies_sender then
+    lua_util.debugm(N, task, 'Could not get sender')
+    return nil
+  end
+
+  -- making sender key
+  lua_util.debugm(N, task, 'Sender: %s', replies_sender)
+  local replies_sender_string = lua_util.maybe_obfuscate_string(tostring(replies_sender), settings, settings.sender_prefix)
+  local replies_sender_key = make_key_replies(replies_sender_string:lower(), 8)
+
+  lua_util.debugm(N, task, 'Sender key: %s', replies_sender_key)
+
+  local function redis_zscore_global_cb(err, data)
+    if err ~= nil then
+      rspamd_logger.errx(task, 'Couldn\'t find sender %s in global replies set. Ended with error: %s', replies_sender, err)
+      return
+    end
+
+    --checking if zcore have not found score of a sender
+    if type(data) ~= 'userdata' then
+      lua_util.debugm(N, task, 'Sender: %s verified. Output: %s', replies_sender, data)
+      task:insert_result(settings.symbol_check_mail_global, 1.0, replies_sender)
+      return
+    end
+    lua_util.debugm(N, task, 'Sender: %s was not verified', replies_sender)
+  end
+
+  -- key for global replies set
+  local replies_global_key = make_key_replies(settings.sender_key_global,
+          settings.sender_key_size, settings.sender_prefix)
+  -- using zscore to find sender in global set
+  lua_util.debugm(N, task, 'Making redis request to global replies set')
+  lua_redis.redis_make_request(task,
+          redis_params, -- connect params
+          replies_sender_key, -- hash key
+          false, -- is write
+          redis_zscore_global_cb, --callback
+          'ZSCORE', -- command
+          { replies_global_key, replies_sender } -- arguments
+  )
+end
+
 local opts = rspamd_config:get_all_opt('known_senders')
 if opts then
   settings = lua_util.override_defaults(settings, opts)
@@ -210,7 +337,8 @@ if opts then
 
   if redis_params then
     local map_conf = settings.domains
-    settings.domains = lua_maps.map_add_from_ucl(settings.domains, 'set', 'domains to track senders from')
+    settings.domains = lua_maps.map_add_from_ucl(settings.domains,
+            'set', 'domains to track senders from')
     if not settings.domains then
       rspamd_logger.errx(rspamd_config, "couldn't add map %s, disable module",
           map_conf)
@@ -221,6 +349,8 @@ if opts then
         'Known elements redis key', {
           type = 'zset/bloom filter',
         })
+    lua_redis.register_prefix(settings.sender_prefix, N,
+        'Prefix to identify replies sets')
     local id = rspamd_config:register_symbol({
       name = settings.symbol,
       type = 'normal',
@@ -230,6 +360,20 @@ if opts then
       augmentations = { string.format("timeout=%f", redis_params.timeout or 0.0) }
     })
 
+    rspamd_config:register_symbol({
+      name = settings.symbol_check_mail_local,
+      type = 'normal',
+      callback = verify_local_replies_set,
+      score = 1.0
+    })
+
+    rspamd_config:register_symbol({
+      name = settings.symbol_check_mail_global,
+      type = 'normal',
+      callback = check_known_incoming_mail_callback,
+      score = 1.0
+    })
+
     if settings.symbol_unknown and #settings.symbol_unknown > 0 then
       rspamd_config:register_symbol({
         name = settings.symbol_unknown,
@@ -243,3 +387,7 @@ if opts then
     lua_util.disable_module(N, "redis")
   end
 end
+
+rspamd_config:add_post_init(function(cfg, ev_base, worker)
+  configure_scripts(cfg, ev_base, worker)
+end)
diff --git a/src/plugins/lua/replies.lua b/src/plugins/lua/replies.lua
index c4df9c97e..08fb68bc7 100644
--- a/src/plugins/lua/replies.lua
+++ b/src/plugins/lua/replies.lua
@@ -34,9 +34,14 @@ local settings = {
   expire = 86400, -- 1 day by default
   key_prefix = 'rr',
   key_size = 20,
+  sender_prefix = 'rsrk',
+  sender_key_global = 'verified_senders',
+  sender_key_size = 20,
   message = 'Message is reply to one we originated',
   symbol = 'REPLY',
   score = -4, -- Default score
+  max_local_size = 20,
+  max_global_size = 30,
   use_auth = true,
   use_local = true,
   cookie = nil,
@@ -44,6 +49,10 @@ local settings = {
   cookie_is_pattern = false,
   cookie_valid_time = '2w', -- 2 weeks by default
   min_message_id = 2, -- minimum length of the message-id header
+  reply_sender_privacy = false,
+  reply_sender_privacy_alg = 'blake2',
+  reply_sender_privacy_prefix = 'obf',
+  reply_sender_privacy_length = 16,
 }
 
 local N = "replies"
@@ -51,25 +60,58 @@ local N = "replies"
 local function make_key(goop, sz, prefix)
   local h = hash.create()
   h:update(goop)
-  local key
-  if sz then
-    key = h:base32():sub(1, sz)
-  else
-    key = h:base32()
-  end
-
-  if prefix then
-    key = prefix .. key
-  end
-
+  local key = (prefix or '') .. h:base32():sub(1, sz)
   return key
 end
 
+local global_replies_set_script
+local local_replies_set_script
+
+local function configure_redis_scripts(_, _)
+  local redis_script_zadd_global = [[
+      redis.call('ZREMRANGEBYRANK', KEYS[1], 0, -({= max_global_size =} + 1)) -- keeping size of global replies set
+      local recipients_addrs = ARGV
+      if recipients_addrs ~= nil then
+        for _, rcpt in ipairs(recipients_addrs) do
+          -- adding recipients to the global replies set
+          redis.call('ZINCRBY', KEYS[1], 1, tostring(rcpt))
+        end
+      end
+      ]]
+  local set_script_zadd_global = lua_util.jinja_template(redis_script_zadd_global,
+          { max_global_size = settings.max_global_size })
+  global_replies_set_script =  lua_redis.add_redis_script(set_script_zadd_global, redis_params)
+
+  local redis_script_zadd_local = [[
+      redis.call('ZREMRANGEBYRANK', KEYS[1], 0, -({= max_local_size =} + 1)) -- keeping size of local replies set
+      local given_params = ARGV
+      if given_params ~= nil then
+          local task_time = given_params[1]
+          table.remove(given_params, 1)
+          -- passing_params is a table that will be passed to the redis
+          local passing_params = {}
+          for _, rcpt in ipairs(given_params) do
+            -- adding recipients for the local replies set
+            table.insert(passing_params, task_time)
+            table.insert(passing_params, rcpt)
+          end
+          redis.call('ZADD', KEYS[1], unpack(passing_params))
+
+          -- setting expire for local replies set
+          redis.call('EXPIRE', KEYS[1], tostring(math.floor('{= expire_time =}')))
+      end
+    ]]
+  local set_script_zadd_local = lua_util.jinja_template(redis_script_zadd_local,
+          { expire_time = settings.expire, max_local_size = settings.max_local_size })
+  local_replies_set_script = lua_redis.add_redis_script(set_script_zadd_local, redis_params)
+end
+
 local function replies_check(task)
   local in_reply_to
+
   local function check_recipient(stored_rcpt)
     local rcpts = task:get_recipients('mime')
-
+    lua_util.debugm(N, task, 'recipients: %s', rcpts)
     if rcpts then
       local filter_predicate = function(input_rcpt)
         local real_rcpt_h = make_key(input_rcpt:lower(), 8)
@@ -81,7 +123,12 @@ local function replies_check(task)
         return rcpt.addr or ''
       end, rcpts)) then
         lua_util.debugm(N, task, 'reply to %s validated', in_reply_to)
-        return true
+
+        --storing only addr of rcpt
+        for i = 1, #rcpts do
+          rcpts[i] = rcpts[i].addr
+        end
+        return rcpts
       end
 
       rspamd_logger.infox(task, 'ignoring reply to %s as no recipients are matching hash %s',
@@ -91,7 +138,60 @@ local function replies_check(task)
           in_reply_to, stored_rcpt)
     end
 
-    return false
+    return nil
+  end
+
+  local function add_to_global_replies_set(params)
+    local global_key = make_key(settings.sender_key_global, settings.sender_key_size, settings.sender_prefix)
+
+    lua_util.debugm(N, task, 'Adding recipients %s to global replies set', params)
+
+    local function zadd_global_set_cb(err, _)
+      if err ~= nil then
+        rspamd_logger.errx(task, 'failed to add recipients %s to global replies set with error: %s', params, err)
+        return
+      end
+      lua_util.debugm(N, task, 'added recipients %s to global replies set', params)
+    end
+
+    lua_redis.exec_redis_script(global_replies_set_script,
+            { task = task, is_write = true },
+            zadd_global_set_cb,
+            { global_key }, params)
+  end
+
+  local function add_to_replies_set(recipients)
+    local sender = task:get_reply_sender()
+
+    local task_time = task:get_timeval(true)
+
+    -- making params out of recipients list for replies set
+    local task_time_str = tostring(task_time)
+
+    local sender_string = lua_util.maybe_obfuscate_string(tostring(sender), settings, settings.sender_prefix)
+    local sender_key = make_key(sender_string:lower(), 8)
+
+    local params = recipients
+    lua_util.debugm(N, task,
+    'Adding recipients %s to sender %s local replies set', recipients, sender_key)
+
+    local function zadd_cb(err, _)
+      if err ~= nil then
+        rspamd_logger.errx(task, 'adding to %s failed with error: %s', sender_key, err)
+        return
+      end
+
+      lua_util.debugm(N, task, 'added data: %s to sender: %s', recipients, sender_key)
+
+      table.remove(params, 1) -- removing task_time_str from params
+      add_to_global_replies_set(params)
+    end
+
+    table.insert(params, 1, task_time_str)
+    lua_redis.exec_redis_script(local_replies_set_script,
+            { task = task, is_write = true },
+            zadd_cb,
+            { sender_key }, params)
   end
 
   local function redis_get_cb(err, data, addr)
@@ -99,8 +199,10 @@ local function replies_check(task)
       rspamd_logger.errx(task, 'redis_get_cb error when reading data from %s: %s', addr:get_addr(), err)
       return
     end
-    if data and type(data) == 'string' and check_recipient(data) then
+    local recipients = check_recipient(data)
+    if type(data) == 'string' and recipients then
       -- Hash was found
+      add_to_replies_set(recipients)
       task:insert_result(settings['symbol'], 1.0)
       if settings['action'] ~= nil then
         local ip_addr = task:get_ip()
@@ -225,7 +327,6 @@ local function replies_check_cookie(task)
   if irt == nil then
     return
   end
-
   local cr = require "rspamd_cryptobox"
   -- Extract user part if needed
   local extracted_cookie = irt:match('^%<?([^@]+)@.*$')
@@ -285,6 +386,9 @@ if opts then
         settings.cookie_valid_time = lua_util.parse_time_interval(settings.cookie_valid_time)
       end
 
+      lua_redis.register_prefix(settings.sender_prefix, N,
+              'Prefix to identify replies sets')
+
       local id = rspamd_config:register_symbol({
         name = 'REPLIES_CHECK',
         type = 'prefilter',
@@ -326,3 +430,7 @@ if opts then
     })
   end
 end
+
+rspamd_config:add_on_load(function(cfg, ev_base, _)
+  configure_redis_scripts(cfg, ev_base)
+end)