diff options
author | Vsevolod Stakhov <vsevolod@rspamd.com> | 2024-07-09 14:22:31 +0100 |
---|---|---|
committer | Vsevolod Stakhov <vsevolod@rspamd.com> | 2024-07-09 14:22:31 +0100 |
commit | 87702745e128977d238284b2a69415bcf94974e6 (patch) | |
tree | eb1d91f5b100eb39ad0a8a14e2dbf8abe32fe878 /lualib | |
parent | 69621218bc47eb54a0dffda7ce819c9309e42c67 (diff) | |
download | rspamd-87702745e128977d238284b2a69415bcf94974e6.tar.gz rspamd-87702745e128977d238284b2a69415bcf94974e6.zip |
[Project] Enable compatibility with the existing buckets
Diffstat (limited to 'lualib')
-rw-r--r-- | lualib/redis_scripts/ratelimit_check.lua | 22 | ||||
-rw-r--r-- | lualib/redis_scripts/ratelimit_update.lua | 74 |
2 files changed, 56 insertions, 40 deletions
diff --git a/lualib/redis_scripts/ratelimit_check.lua b/lualib/redis_scripts/ratelimit_check.lua index d39cdf148..f24e0daf0 100644 --- a/lualib/redis_scripts/ratelimit_check.lua +++ b/lualib/redis_scripts/ratelimit_check.lua @@ -7,6 +7,7 @@ -- KEYS[4]: The maximum allowed burst -- KEYS[5]: The expiration time for a bucket -- KEYS[6]: The number of recipients for the message +-- KEYS[7]: Enable dynamic ratelimits -- Redis keys used: -- l: Last hit (time in milliseconds) @@ -29,6 +30,7 @@ local nrcpt = tonumber(KEYS[6]) local leak_rate = tonumber(KEYS[3]) local max_burst = tonumber(KEYS[4]) local prefix = KEYS[1] +local enable_dynamic = KEYS[7] == 'true' local dynr, dynb, leaked = 0, 0, 0 if not last then -- New bucket @@ -52,9 +54,13 @@ pending = pending + nrcpt -- this message if burst + pending > 0 then -- If we have any time passed if burst > 0 and last < now then - dynr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000.0 - if dynr == 0 then - dynr = 0.0001 + if enable_dynamic then + dynr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000.0 + if dynr == 0 then + dynr = 0.0001 + end + else + dynr = 1.0 end leak_rate = leak_rate * dynr leaked = ((now - last) * leak_rate) @@ -66,9 +72,13 @@ if burst + pending > 0 then redis.call('HSET', prefix, 'l', tostring(now)) end - dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0 - if dynb == 0 then - dynb = 0.0001 + if enable_dynamic then + dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0 + if dynb == 0 then + dynb = 0.0001 + end + else + dynb = 1.0 end burst = burst + pending diff --git a/lualib/redis_scripts/ratelimit_update.lua b/lualib/redis_scripts/ratelimit_update.lua index caee8fb31..8b7a934dc 100644 --- a/lualib/redis_scripts/ratelimit_update.lua +++ b/lualib/redis_scripts/ratelimit_update.lua @@ -9,12 +9,14 @@ -- KEYS[6] - max_burst_rate: The maximum allowed value for the dynamic burst multiplier. -- KEYS[7] - expire: The expiration time for the Redis key storing the bucket information, in seconds. -- KEYS[8] - number_of_recipients: The number of requests to be allowed (or the increase rate). +-- KEYS[9] - Enable dynamic ratelimits -- 1. Retrieve the last hit time and initialize variables local prefix = KEYS[1] local last = redis.call('HGET', prefix, 'l') local now = tonumber(KEYS[2]) local nrcpt = tonumber(KEYS[8]) +local enable_dynamic = KEYS[9] == 'true' if not last then -- 2. Initialize a new bucket if the last hit time is not found (must not happen) redis.call('HMSET', prefix, 'l', tostring(now), 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0') @@ -25,48 +27,52 @@ end -- 3. Update the dynamic rate multiplier based on input parameters local dr, db = 1.0, 1.0 -local max_dr = tonumber(KEYS[5]) +if enable_dynamic then + local max_dr = tonumber(KEYS[5]) -if max_dr > 1 then - local rate_mult = tonumber(KEYS[3]) - dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000 + if max_dr > 1 then + local rate_mult = tonumber(KEYS[3]) + dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000 - if rate_mult > 1.0 and dr < max_dr then - dr = dr * rate_mult - if dr > 0.0001 then - redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) - else - redis.call('HSET', prefix, 'dr', '1') - end - elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then - dr = dr * rate_mult - if dr > 0.0001 then - redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) - else - redis.call('HSET', prefix, 'dr', '1') + if rate_mult > 1.0 and dr < max_dr then + dr = dr * rate_mult + if dr > 0.0001 then + redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) + else + redis.call('HSET', prefix, 'dr', '1') + end + elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then + dr = dr * rate_mult + if dr > 0.0001 then + redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) + else + redis.call('HSET', prefix, 'dr', '1') + end end end end -- 4. Update the dynamic burst multiplier based on input parameters -local max_db = tonumber(KEYS[6]) -if max_db > 1 then - local rate_mult = tonumber(KEYS[4]) - db = tonumber(redis.call('HGET', prefix, 'db')) / 10000 +if enable_dynamic then + local max_db = tonumber(KEYS[6]) + if max_db > 1 then + local rate_mult = tonumber(KEYS[4]) + db = tonumber(redis.call('HGET', prefix, 'db')) / 10000 - if rate_mult > 1.0 and db < max_db then - db = db * rate_mult - if db > 0.0001 then - redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) - else - redis.call('HSET', prefix, 'db', '1') - end - elseif rate_mult < 1.0 and db > (1.0 / max_db) then - db = db * rate_mult - if db > 0.0001 then - redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) - else - redis.call('HSET', prefix, 'db', '1') + if rate_mult > 1.0 and db < max_db then + db = db * rate_mult + if db > 0.0001 then + redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) + else + redis.call('HSET', prefix, 'db', '1') + end + elseif rate_mult < 1.0 and db > (1.0 / max_db) then + db = db * rate_mult + if db > 0.0001 then + redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) + else + redis.call('HSET', prefix, 'db', '1') + end end end end |