[Feature] Add mempool maps in multimap

3 files changed, 52 insertions, 2 deletions

diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index a4577676a..b36fe9f79 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -201,6 +201,30 @@ local function apply_filename_filter(task, filter, fn, r)
   return fn
 end
 
+local function apply_regexp_filter(task, filter, fn, r)
+  if string.find(filter, 'regexp:') then
+    if not r['re_filter'] then
+      local type,pat = string.match(filter, '(regexp:)(.+)')
+      if type and pat then
+        r['re_filter'] = regexp.create(pat)
+      end
+    end
+
+    if not r['re_filter'] then
+      rspamd_logger.errx(task, 'bad search filter: %s', filter)
+    else
+      local results = r['re_filter']:search(fn)
+      if results then
+        return results[1]
+      else
+        return nil
+      end
+    end
+  end
+
+  return fn
+end
+
 local function apply_content_filter(task, filter, r)
   if filter == 'body' then
     return {task:get_rawbody()}
@@ -239,6 +263,7 @@ local multimap_filters = {
   header = apply_addr_filter,
   url = apply_url_filter,
   filename = apply_filename_filter,
+  mempool = apply_regex_filter,
   --content = apply_content_filter, -- Content filters are special :(
 }
 
@@ -519,6 +544,11 @@ local function multimap_callback(task, rule)
     if country then
       match_rule(rule, country)
     end
+  elseif rt == 'mempool' then
+    local var = task:get_mempool():get_variable(rule['variable'])
+    if var then
+      match_rule(rule, var)
+    end
   end
 end
 
@@ -547,6 +577,10 @@ local function add_multimap_rule(key, newrule)
     newrule['description'] = string.format('multimap, type %s: %s', newrule['type'],
       newrule['symbol'])
   end
+  if newrule['type'] == 'mempool' and not newrule['variable'] then
+    rspamd_logger.errx(rspamd_config, 'mempool map requires variable')
+    return nil
+  end
   -- Check cdb flag
   if string.find(newrule['map'], '^cdb://.*$') then
     local test = cdb.create(newrule['map'])
@@ -607,7 +641,8 @@ local function add_multimap_rule(key, newrule)
         or newrule['type'] == 'content'
         or newrule['type'] == 'hostname'
         or newrule['type'] == 'asn'
-        or newrule['type'] == 'country' then
+        or newrule['type'] == 'country'
+        or newrule['type'] == 'mempool' then
         if newrule['regexp'] then
           newrule['hash'] = rspamd_config:add_map ({
             url = newrule['map'],
diff --git a/test/functional/cases/102_multimap.robot b/test/functional/cases/102_multimap.robot
index de3595378..abe79f512 100644
--- a/test/functional/cases/102_multimap.robot
+++ b/test/functional/cases/102_multimap.robot
@@ -224,10 +224,19 @@ MAP - REDIS - CC - HIT
   ${result} =  Scan Message With Rspamc  ${MESSAGE}  -i  8.8.8.8
   Check Rspamc  ${result}  REDIS_COUNTRY
 
-MAP - REDIS - ASN - MISS
+MAP - REDIS - CC - MISS
   ${result} =  Scan Message With Rspamc  ${MESSAGE}  -i  46.228.47.114
   Check Rspamc  ${result}  REDIS_COUNTRY  inverse=1
 
+MAP - REDIS - ASN FILTERED - HIT
+  Redis HSET  asn  1  ${EMPTY}
+  ${result} =  Scan Message With Rspamc  ${MESSAGE}  -i  8.8.8.8
+  Check Rspamc  ${result}  REDIS_ASN_FILTERED
+
+MAP - REDIS - ASN FILTERED - MISS
+  ${result} =  Scan Message With Rspamc  ${MESSAGE}  -i  46.228.47.114
+  Check Rspamc  ${result}  REDIS_ASN_FILTERED  inverse=1
+
 *** Keywords ***
 Multimap Setup
   ${PLUGIN_CONFIG} =  Get File  ${TESTDIR}/configs/multimap.conf
diff --git a/test/functional/configs/multimap.conf b/test/functional/configs/multimap.conf
index 3823914cf..fc20fa14c 100644
--- a/test/functional/configs/multimap.conf
+++ b/test/functional/configs/multimap.conf
@@ -106,4 +106,10 @@ multimap {
     type = "asn";
     map = "redis://asn";
   }
+  REDIS_ASN_FILTERED {
+    type = "mempool";
+    variable = "asn";
+    map = "redis://asn";
+    filter = "regexp:/^(\d).*/";
+  }
 }