commit 7ba0e69: [Feature] Ratelimit: Add per bucket configurations

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2021-10-01 14:35:43 +0100
URL: https://github.com/rspamd/rspamd/commit/7ba0e69b2f31454f21b80e113737df968cc9862a (HEAD -> master)

[Feature] Ratelimit: Add per bucket configurations
This change add custom per bucket configuration
- symbol: for custom symbol
- message: for custom message
- skip_soft_reject: to disable soft rejection logic for this bucket

---
 src/plugins/lua/ratelimit.lua | 54 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 40 insertions(+), 14 deletions(-)

diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua
index 8e51b4e6b..e31026366 100644
--- a/src/plugins/lua/ratelimit.lua
+++ b/src/plugins/lua/ratelimit.lua
@@ -288,6 +288,9 @@ local bucket_schema = ts.shape{
   burst = ts.number + ts.string / lua_util.dehumanize_number,
   rate = ts.number + ts.string / str_to_rate,
   skip_recipients = ts.boolean:is_optional(),
+  symbol = ts.string:is_optional(),
+  message = ts.string:is_optional(),
+  skip_soft_reject = ts.boolean:is_optional(),
 }
 
 local function parse_limit(name, data)
@@ -598,27 +601,33 @@ local function ratelimit_cb(task)
 
         if data[1] == 1 then
           -- set symbol only and do NOT soft reject
-          if settings.symbol then
-            task:insert_result(settings.symbol, 1.0,
-                string.format('%s(%s)', lim_name, lim_key))
-            rspamd_logger.infox(task,
-                'set_symbol_only: ratelimit "%s(%s)" exceeded, (%s / %s): %s (%s:%s dyn); redis key: %s',
-                lim_name, prefix,
-                bucket.burst, bucket.rate,
-                data[2], data[3], data[4], lim_key)
-            return
-            -- set INFO symbol and soft reject
-          elseif settings.info_symbol then
-            task:insert_result(settings.info_symbol, 1.0,
+          if bucket.symbol then
+            -- Per bucket symbol
+            task:insert_result(bucket.symbol, 1.0,
                 string.format('%s(%s)', lim_name, lim_key))
+          else
+            if settings.symbol then
+              task:insert_result(settings.symbol, 1.0,
+                  string.format('%s(%s)', lim_name, lim_key))
+            elseif settings.info_symbol then
+              task:insert_result(settings.info_symbol, 1.0,
+                  string.format('%s(%s)', lim_name, lim_key))
+            end
           end
           rspamd_logger.infox(task,
               'ratelimit "%s(%s)" exceeded, (%s / %s): %s (%s:%s dyn); redis key: %s',
               lim_name, prefix,
               bucket.burst, bucket.rate,
               data[2], data[3], data[4], lim_key)
-          task:set_pre_result('soft reject',
-              message_func(task, lim_name, prefix, bucket, lim_key), N)
+
+          if not settings.symbol and not bucket.skip_soft_reject then
+            if not bucket.message then
+              task:set_pre_result('soft reject',
+                  message_func(task, lim_name, prefix, bucket, lim_key), N)
+            else
+              task:set_pre_result('soft reject', bucket.message)
+            end
+          end
         end
       end
     end
@@ -866,6 +875,23 @@ if opts then
 
     local id = rspamd_config:register_symbol(s)
 
+    -- Register per bucket symbols
+    -- Display what's enabled
+    fun.each(function(set)
+      if set.buckets then
+        for _,b in ipairs(set.buckets) do
+          if b.symbol then
+            rspamd_config:register_symbol{
+              type = 'virtual',
+              name = b.symbol,
+              score = 0.0,
+              parent = id
+            }
+          end
+        end
+      end
+    end, settings.limits)
+
     if settings.info_symbol then
       rspamd_config:register_symbol{
         type = 'virtual',