[Feature] Add mempool maps in multimap

1 files changed, 36 insertions, 1 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'],