Merge branch 'master' into lua-coroutine-model

2 files changed, 82 insertions, 21 deletions

diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index 4514bc4f7..909260681 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -29,6 +29,7 @@ local rspamd_expression = require "rspamd_expression"
 local rspamd_ip = require "rspamd_ip"
 local lua_util = require "lua_util"
 local rspamd_dns = require "rspamd_dns"
+local lua_selectors = require "lua_selectors"
 local redis_params
 local fun = require "fun"
 local N = 'multimap'
@@ -81,6 +82,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,
   },
@@ -368,6 +372,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 :(
 }
@@ -833,6 +838,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]
@@ -875,6 +886,21 @@ local function add_multimap_rule(key, newrule)
     rspamd_logger.errx(rspamd_config, 'mempool map requires variable')
     return nil
   end
+  if newrule['type'] == 'selector' then
+    if 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
+  end
   -- Check cdb flag
   if type(newrule['map']) == 'string' and string.find(newrule['map'], '^cdb://.*$') then
     newrule['cdb'] = cdb.create(newrule['map'])
@@ -987,17 +1013,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'],
diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua
index 59f3c0522..8c0c48387 100644
--- a/src/plugins/lua/ratelimit.lua
+++ b/src/plugins/lua/ratelimit.lua
@@ -27,6 +27,7 @@ local fun = require "fun"
 local lua_maps = require "lua_maps"
 local lua_util = require "lua_util"
 local rspamd_hash = require "rspamd_cryptobox_hash"
+local lua_selectors = require "lua_selectors"
 
 -- A plugin that implements ratelimits using redis
 
@@ -431,13 +432,37 @@ end
 
 local function limit_to_prefixes(task, k, v, prefixes)
   local n = 0
-  for _,bucket in ipairs(v) do
-    local prefix = gen_rate_key(task, k, bucket)
-
-    if prefix then
-      prefixes[prefix] = make_prefix(prefix, k, bucket)
-      n = n + 1
+  for _,bucket in ipairs(v.buckets) do
+    if v.selector then
+      local selectors = lua_selectors.process_selectors(task, v.selector)
+      if selectors then
+        local combined = lua_selectors.combine_selectors(task, selectors, ':')
+        if type(combined) == 'string' then
+          prefixes[combined] = make_prefix(combined, k, bucket)
+          n = n + 1
+        else
+          fun.each(function(p)
+            prefixes[p] = make_prefix(p, k, bucket)
+            n = n + 1
+          end, combined)
+        end
+      end
+    else
+      local prefix = gen_rate_key(task, k, bucket)
+      if prefix then
+        if type(prefix) == 'string' then
+          prefixes[prefix] = make_prefix(prefix, k, bucket)
+          n = n + 1
+        else
+          fun.each(function(p)
+            prefixes[p] = make_prefix(p, k, bucket)
+            n = n + 1
+          end, prefix)
+        end
+      end
     end
+
+
   end
 
   return n
@@ -620,10 +645,19 @@ if opts then
   if opts['rates'] and type(opts['rates']) == 'table' then
     -- new way of setting limits
     fun.each(function(t, lim)
-      local buckets = parse_limit(t, lim)
-
-      if buckets and #buckets > 0 then
-        settings.limits[t] = buckets
+      local buckets
+      if type(lim) == 'table' and lim.selector and lim.bucket then
+        settings.limits[t] = {
+          selector = lua_selectors.parse_selector(rspamd_config, lim.selector),
+          buckets = parse_limit(t, lim)
+        }
+      else
+        buckets = parse_limit(t, lim)
+        if buckets and #buckets > 0 then
+          settings.limits[t] = {
+            buckets = buckets
+          }
+        end
       end
     end, opts['rates'])
   end