From 99488d5a2a601f4a4ae30036b07189ff9c17ed3c Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Wed, 2 Nov 2016 13:15:05 +0000
Subject: [PATCH] [Feature] Allow custom functions for ratelimits

---
 src/plugins/lua/greylist.lua  |  7 ++++++-
 src/plugins/lua/ratelimit.lua | 10 +++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/plugins/lua/greylist.lua b/src/plugins/lua/greylist.lua
index 249d52b53..68af9b933 100644
--- a/src/plugins/lua/greylist.lua
+++ b/src/plugins/lua/greylist.lua
@@ -379,8 +379,13 @@ end
 
 local opts =  rspamd_config:get_all_opt('greylist')
 if opts then
+  if opts['message_func'] then
+    settings.message_func = assert(loadstring(opts['message_func']))()
+  end
   for k,v in pairs(opts) do
-    settings[k] = v
+    if k ~= 'message_func' then
+      settings[k] = v
+    end
   end
   if settings['whitelisted_ip'] then
     whitelisted_ip = rspamd_config:add_radix_map(settings['whitelisted_ip'],
diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua
index 7d6086c3b..571cfda97 100644
--- a/src/plugins/lua/ratelimit.lua
+++ b/src/plugins/lua/ratelimit.lua
@@ -36,6 +36,10 @@ local ip_score_lower_bound = 10
 local ip_score_ham_multiplier = 1.1
 local ip_score_spam_divisor = 1.1
 
+local message_func = function(task, limit_type, bucket, threshold)
+  return string.format('Ratelimit "%s" exceeded', limit_type)
+end
+
 local rspamd_logger = require "rspamd_logger"
 local rspamd_redis = require "rspamd_redis"
 local upstream_list = require "rspamd_upstream_list"
@@ -285,7 +289,7 @@ local function check_limits(task, args)
                 'ratelimit "%s" exceeded: %s elements with %s limit',
                 rtype, bucket, threshold)
               task:set_pre_result('soft reject',
-                string.format('Ratelimit "%s" exceeded', rtype))
+                message_func(task, rtype, bucket, threshold))
             end
           end
         end
@@ -566,6 +570,10 @@ if opts then
     user_keywords = opts['user_keywords']
   end
 
+  if opts['message_func'] then
+    message_func = assert(loadstring(opts['message_func']))()
+  end
+
   redis_params = rspamd_parse_redis_server('ratelimit')
   if not redis_params then
     rspamd_logger.infox(rspamd_config, 'no servers are specified, disabling module')
-- 
2.39.5