From a7833e45d85c7f416ada70ad2a35d1172595f41a Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Fri, 17 Aug 2018 17:27:57 +0100
Subject: [PATCH] [Feature] Integrate selectors framework to multimap

---
 src/plugins/lua/multimap.lua | 47 +++++++++++++++++++++++++++---------
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index 2898bfe4a..8ba0a164d 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -28,6 +28,7 @@ local regexp = require "rspamd_regexp"
 local rspamd_expression = require "rspamd_expression"
 local rspamd_ip = require "rspamd_ip"
 local lua_util = require "lua_util"
+local lua_selectors = require "lua_selectors"
 local redis_params
 local fun = require "fun"
 local N = 'multimap'
@@ -80,6 +81,9 @@ local value_types = {
   mempool = {
     get_value = function(val) return val end,
   },
+  selector = {
+    get_value = function(val) return val end,
+  },
   symbol_options = {
     get_value = function(val) return val end,
   },
@@ -367,6 +371,7 @@ local multimap_filters = {
   url = apply_url_filter,
   filename = apply_filename_filter,
   mempool = apply_regexp_filter,
+  selector = apply_regexp_filter,
   hostname = apply_hostname_filter,
   --content = apply_content_filter, -- Content filters are special :(
 }
@@ -831,6 +836,12 @@ local function multimap_callback(task, rule)
         end
       end
     end,
+    selector = function()
+      local elts = lua_selectors.process_selectors(task, rule.selector) or {}
+      for _,elt in ipairs(elts) do
+        match_rule(rule, elt)
+      end
+    end,
   }
   process_rule_funcs.ip = process_rule_funcs.dnsbl
   local f = process_rule_funcs[rt]
@@ -873,6 +884,19 @@ local function add_multimap_rule(key, newrule)
     rspamd_logger.errx(rspamd_config, 'mempool map requires variable')
     return nil
   end
+  if newrule['type'] == 'selector' and not newrule['selector'] then
+    rspamd_logger.errx(rspamd_config, 'selector map requires selector definition')
+    return nil
+  else
+    local selector = lua_selectors.parse_selector(rspamd_config, newrule['selector'])
+
+    if not selector then
+      rspamd_logger.errx(rspamd_config, 'selector map has invalid selector')
+      return nil
+    end
+
+    newrule.selector = selector
+  end
   -- Check cdb flag
   if type(newrule['map']) == 'string' and string.find(newrule['map'], '^cdb://.*$') then
     newrule['cdb'] = cdb.create(newrule['map'])
@@ -985,17 +1009,18 @@ local function add_multimap_rule(key, newrule)
           end
         end
       elseif newrule['type'] == 'header'
-        or newrule['type'] == 'rcpt'
-        or newrule['type'] == 'from'
-        or newrule['type'] == 'helo'
-        or newrule['type'] == 'symbol_options'
-        or newrule['type'] == 'filename'
-        or newrule['type'] == 'url'
-        or newrule['type'] == 'content'
-        or newrule['type'] == 'hostname'
-        or newrule['type'] == 'asn'
-        or newrule['type'] == 'country'
-        or newrule['type'] == 'mempool' then
+          or newrule['type'] == 'rcpt'
+          or newrule['type'] == 'from'
+          or newrule['type'] == 'helo'
+          or newrule['type'] == 'symbol_options'
+          or newrule['type'] == 'filename'
+          or newrule['type'] == 'url'
+          or newrule['type'] == 'content'
+          or newrule['type'] == 'hostname'
+          or newrule['type'] == 'asn'
+          or newrule['type'] == 'country'
+          or newrule['type'] == 'mempool'
+          or newrule['type'] == 'selector'then
         if newrule['regexp'] then
           newrule['hash'] = rspamd_config:add_map ({
             url = newrule['map'],
-- 
2.39.5