commit fb789af: [Rework] Reputation: Use more flexible types in get/set functions

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2021-04-13 11:57:15 +0100
URL: https://github.com/rspamd/rspamd/commit/fb789af8bbb72a2260bc67782bac542c8c335d06

[Rework] Reputation: Use more flexible types in get/set functions

---
 src/plugins/lua/reputation.lua | 74 ++++++++++++++++++++++++++++--------------
 1 file changed, 49 insertions(+), 25 deletions(-)

diff --git a/src/plugins/lua/reputation.lua b/src/plugins/lua/reputation.lua
index d846ece9a..227654428 100644
--- a/src/plugins/lua/reputation.lua
+++ b/src/plugins/lua/reputation.lua
@@ -197,7 +197,7 @@ local function dkim_reputation_filter(task, rule)
   for dom,res in pairs(requests) do
     -- tld + "." + check_result, e.g. example.com.+ - reputation for valid sigs
     local query = string.format('%s.%s', dom, res)
-    rule.backend.get_token(task, rule, query, tokens_cb)
+    rule.backend.get_token(task, rule, nil, query, tokens_cb, 'string')
   end
 end
 
@@ -209,7 +209,7 @@ local function dkim_reputation_idempotent(task, rule)
     for dom,res in pairs(requests) do
       -- tld + "." + check_result, e.g. example.com.+ - reputation for valid sigs
       local query = string.format('%s.%s', dom, res)
-      rule.backend.set_token(task, rule, query, sc)
+      rule.backend.set_token(task, rule, nil, query, sc, 'string')
     end
   end
 end
@@ -325,7 +325,7 @@ local function url_reputation_filter(task, rule)
       indexed_tokens_cb(err, i, values)
     end
 
-    rule.backend.get_token(task, rule, req[1], tokens_cb)
+    rule.backend.get_token(task, rule, nil, req[1], tokens_cb, 'string')
   end
 end
 
@@ -335,7 +335,7 @@ local function url_reputation_idempotent(task, rule)
 
   if sc then
     for _,tld in ipairs(requests) do
-      rule.backend.set_token(task, rule, tld[1], sc)
+      rule.backend.set_token(task, rule, nil, tld[1], sc, 'string')
     end
   end
 end
@@ -465,13 +465,16 @@ local function ip_reputation_filter(task, rule)
   end
 
   if asn then
-    rule.backend.get_token(task, rule, cfg.asn_prefix .. asn, gen_token_callback('asn'))
+    rule.backend.get_token(task, rule, cfg.asn_prefix, asn,
+            gen_token_callback('asn'), 'string')
   end
   if country then
-    rule.backend.get_token(task, rule, cfg.country_prefix .. country, gen_token_callback('country'))
+    rule.backend.get_token(task, rule, cfg.country_prefix, country,
+            gen_token_callback('country'), 'string')
   end
 
-  rule.backend.get_token(task, rule, cfg.ip_prefix .. tostring(ip), gen_token_callback('ip'))
+  rule.backend.get_token(task, rule, cfg.ip_prefix, ip,
+          gen_token_callback('ip'), 'ip')
 end
 
 -- Used to set scores
@@ -505,13 +508,13 @@ local function ip_reputation_idempotent(task, rule)
   local sc = extract_task_score(task, rule)
   if sc then
     if asn then
-      rule.backend.set_token(task, rule, cfg.asn_prefix .. asn, sc)
+      rule.backend.set_token(task, rule, cfg.asn_prefix, asn, sc, 'string')
     end
     if country then
-      rule.backend.set_token(task, rule, cfg.country_prefix .. country, sc)
+      rule.backend.set_token(task, rule, cfg.country_prefix, country, sc, 'string')
     end
 
-    rule.backend.set_token(task, rule, cfg.ip_prefix .. tostring(ip), sc)
+    rule.backend.set_token(task, rule, cfg.ip_prefix, ip, sc, 'ip')
   end
 end
 
@@ -568,7 +571,7 @@ local function spf_reputation_filter(task, rule)
     end
   end
 
-  rule.backend.get_token(task, rule, hkey, tokens_cb)
+  rule.backend.get_token(task, rule, nil, hkey, tokens_cb, 'string')
 end
 
 local function spf_reputation_idempotent(task, rule)
@@ -583,7 +586,7 @@ local function spf_reputation_idempotent(task, rule)
 
   lua_util.debugm(N, task, 'set spf record %s -> %s = %s',
       spf_record, hkey, sc)
-  rule.backend.set_token(task, rule, hkey, sc)
+  rule.backend.set_token(task, rule, nil, hkey, sc, 'string')
 end
 
 
@@ -655,12 +658,12 @@ local function generic_reputation_filter(task, rule)
       fun.each(function(e)
         lua_util.debugm(N, task, 'check generic reputation (%s) %s',
           rule['symbol'], e)
-        rule.backend.get_token(task, rule, e, tokens_cb)
+        rule.backend.get_token(task, rule, nil, e, tokens_cb, 'string')
       end, selector_res)
     else
       lua_util.debugm(N, task, 'check generic reputation (%s) %s',
         rule['symbol'], selector_res)
-      rule.backend.get_token(task, rule, selector_res, tokens_cb)
+      rule.backend.get_token(task, rule, nil, selector_res, tokens_cb, 'string')
     end
   end
 end
@@ -677,12 +680,12 @@ local function generic_reputation_idempotent(task, rule)
       fun.each(function(e)
         lua_util.debugm(N, task, 'set generic selector (%s) %s = %s',
             rule['symbol'], e, sc)
-        rule.backend.set_token(task, rule, e, sc)
+        rule.backend.set_token(task, rule, nil, e, sc, 'string')
       end, selector_res)
     else
       lua_util.debugm(N, task, 'set generic selector (%s) %s = %s',
           rule['symbol'], selector_res, sc)
-      rule.backend.set_token(task, rule, selector_res, sc)
+      rule.backend.set_token(task, rule, nil, selector_res, sc, 'string')
     end
   end
 end
@@ -736,8 +739,11 @@ local function reputation_dns_init(rule, _, _, _)
 end
 
 
-local function gen_token_key(token, rule)
-  local res = token
+local function gen_token_key(prefix, token, rule)
+  if prefix then
+    token = prefix .. token
+  end
+  local res = prefix
   if rule.backend.config.hashed then
     local hash_alg = rule.backend.config.hash_alg or "blake2"
     local encoding = "base32"
@@ -769,7 +775,7 @@ end
 
 --[[
 -- Generic interface for get and set tokens functions:
--- get_token(task, rule, token, continuation), where `continuation` is the following function:
+-- get_token(task, rule, prefix, token, continuation, token_type), where `continuation` is the following function:
 --
 -- function(err, token, values) ... end
 -- `err`: string value for error (similar to redis or DNS callbacks)
@@ -788,11 +794,23 @@ end
 -- example of tokens: {'s': 0, 'h': 0, 'p': 1}
 --]]
 
-local function reputation_dns_get_token(task, rule, token, continuation_cb)
+local function reputation_dns_get_token(task, rule, prefix, token, continuation_cb, token_type)
   -- local r = task:get_resolver()
-  local key = gen_token_key(token, rule)
+  -- In DNS we never ever use prefix as prefix, we use if as a suffix!
+  if token_type == 'ip' then
+    token = table.concat(token:inversed_str_octets(), '.')
+  end
+
+  local key = gen_token_key(nil, token, rule)
   local dns_name = key .. '.' .. rule.backend.config.list
 
+  if prefix then
+    dns_name = string.format('%s.%s.%s', key, prefix,
+            rule.backend.config.list)
+  else
+    dns_name = string.format('%s.%s', key, rule.backend.config.list)
+  end
+
   local function dns_cb(_, _, results, err)
     if err and (err ~= 'requested record is not found' and
         err ~= 'no records with this name') then
@@ -921,8 +939,11 @@ local function reputation_redis_init(rule, cfg, ev_base, worker)
   return true
 end
 
-local function reputation_redis_get_token(task, rule, token, continuation_cb)
-  local key = gen_token_key(token, rule)
+local function reputation_redis_get_token(task, rule, prefix, token, continuation_cb, token_type)
+  if token_type == 'ip' then
+    token = tostring(token)
+  end
+  local key = gen_token_key(prefix, token, rule)
 
   local function redis_get_cb(err, data)
     if data then
@@ -956,8 +977,11 @@ local function reputation_redis_get_token(task, rule, token, continuation_cb)
   end
 end
 
-local function reputation_redis_set_token(task, rule, token, sc, continuation_cb)
-  local key = gen_token_key(token, rule)
+local function reputation_redis_set_token(task, rule, prefix, token, sc, continuation_cb, token_type)
+  if token_type == 'ip' then
+    token = tostring(token)
+  end
+  local key = gen_token_key(prefix, token, rule)
 
   local function redis_set_cb(err, data)
     if err then