diff options
| -rw-r--r-- | .luacheckrc | 5 | ||||
| -rw-r--r-- | lualib/lua_redis.lua | 2 | ||||
| -rw-r--r-- | lualib/redis_scripts/ratelimit_check.lua | 62 | ||||
| -rw-r--r-- | lualib/redis_scripts/ratelimit_update.lua | 80 | ||||
| -rw-r--r-- | src/plugins/lua/ratelimit.lua | 154 |
5 files changed, 154 insertions, 149 deletions
diff --git a/.luacheckrc b/.luacheckrc index 099d9ab65..d5a18cc3b 100644 --- a/.luacheckrc +++ b/.luacheckrc @@ -60,6 +60,11 @@ files['/**/lualib/lua_redis.lua'].globals = { 'rspamadm_ev_base', } +files['/**/lualib/redis_scripts/**'].globals = { + 'redis', + 'KEYS', +} + files['/**/src/rspamadm/*'].globals = { 'ansicolors', 'getopt', diff --git a/lualib/lua_redis.lua b/lualib/lua_redis.lua index 76e8db972..62511451e 100644 --- a/lualib/lua_redis.lua +++ b/lualib/lua_redis.lua @@ -1293,7 +1293,7 @@ local function load_redis_script_from_file(filename, redis_params, dir) if not dir then dir = rspamd_paths.LUALIBDIR end if filename:sub(1, 1) ~= package.config:sub(1,1) then -- Relative path - filename = lua_util.join_path(dir, filename) + filename = lua_util.join_path(dir, "redis_scripts", filename) end -- Read file contents local file = io.open(filename, "r") diff --git a/lualib/redis_scripts/ratelimit_check.lua b/lualib/redis_scripts/ratelimit_check.lua new file mode 100644 index 000000000..2b2af11bf --- /dev/null +++ b/lualib/redis_scripts/ratelimit_check.lua @@ -0,0 +1,62 @@ +-- Checks bucket, updating it if needed +-- KEYS[1] - prefix to update, e.g. RL_<triplet>_<seconds> +-- KEYS[2] - current time in milliseconds +-- KEYS[3] - bucket leak rate (messages per millisecond) +-- KEYS[4] - bucket burst +-- KEYS[5] - expire for a bucket +-- KEYS[6] - number of recipients +-- return 1 if message should be ratelimited and 0 if not +-- Redis keys used: +-- l - last hit +-- b - current burst +-- p - pending messages (those that are currently processing) +-- dr - current dynamic rate multiplier (*10000) +-- db - current dynamic burst multiplier (*10000) + +local last = redis.call('HGET', KEYS[1], 'l') +local now = tonumber(KEYS[2]) +local nrcpt = tonumber(KEYS[6]) +local dynr, dynb, leaked = 0, 0, 0 +if not last then + -- New bucket + redis.call('HMSET', KEYS[1], 'l', KEYS[2], 'b', '0', 'dr', '10000', 'db', '10000', 'p', tostring(nrcpt)) + redis.call('EXPIRE', KEYS[1], KEYS[5]) + return {0, '0', '1', '1', '0'} +end + +last = tonumber(last) +local burst,pending = unpack(redis.call('HMGET', KEYS[1], 'b', 'p')) +burst,pending = tonumber(burst or '0'),tonumber(pending or '0') +-- Sanity to avoid races +if burst < 0 then burst = 0 end +if pending < 0 then pending = 0 end +pending = pending + nrcpt -- this message +-- Perform leak +if burst + pending > 0 then + if burst > 0 and last < tonumber(KEYS[2]) then + local rate = tonumber(KEYS[3]) + dynr = tonumber(redis.call('HGET', KEYS[1], 'dr')) / 10000.0 + if dynr == 0 then dynr = 0.0001 end + rate = rate * dynr + leaked = ((now - last) * rate) + if leaked > burst then leaked = burst end + burst = burst - leaked + redis.call('HINCRBYFLOAT', KEYS[1], 'b', -(leaked)) + redis.call('HSET', KEYS[1], 'l', KEYS[2]) + end + + dynb = tonumber(redis.call('HGET', KEYS[1], 'db')) / 10000.0 + if dynb == 0 then dynb = 0.0001 end + + burst = burst + pending + if burst > 0 and (burst + tonumber(KEYS[6])) > tonumber(KEYS[4]) * dynb then + return {1, tostring(burst - pending), tostring(dynr), tostring(dynb), tostring(leaked)} + end + -- Increase pending if we allow ratelimit + redis.call('HINCRBY', KEYS[1], 'p', nrcpt) +else + burst = 0 + redis.call('HMSET', KEYS[1], 'b', '0', 'p', tostring(nrcpt)) +end + +return {0, tostring(burst), tostring(dynr), tostring(dynb), tostring(leaked)}
\ No newline at end of file diff --git a/lualib/redis_scripts/ratelimit_update.lua b/lualib/redis_scripts/ratelimit_update.lua new file mode 100644 index 000000000..682ddd0c6 --- /dev/null +++ b/lualib/redis_scripts/ratelimit_update.lua @@ -0,0 +1,80 @@ +-- Updates a bucket +-- KEYS[1] - prefix to update, e.g. RL_<triplet>_<seconds> +-- KEYS[2] - current time in milliseconds +-- KEYS[3] - dynamic rate multiplier +-- KEYS[4] - dynamic burst multiplier +-- KEYS[5] - max dyn rate (min: 1/x) +-- KEYS[6] - max burst rate (min: 1/x) +-- KEYS[7] - expire for a bucket +-- KEYS[8] - number of recipients (or increase rate) +-- Redis keys used: +-- l - last hit +-- b - current burst +-- p - messages pending (must be decreased by 1) +-- dr - current dynamic rate multiplier +-- db - current dynamic burst multiplier + +local last = redis.call('HGET', KEYS[1], 'l') +local nrcpt = tonumber(KEYS[8]) +if not last then + -- New bucket (why??) + redis.call('HMSET', KEYS[1], 'l', KEYS[2], 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0') + redis.call('EXPIRE', KEYS[1], KEYS[7]) + return {1, 1, 1} +end + +local dr, db = 1.0, 1.0 + +if tonumber(KEYS[5]) > 1 then + local rate_mult = tonumber(KEYS[3]) + local rate_limit = tonumber(KEYS[5]) + dr = tonumber(redis.call('HGET', KEYS[1], 'dr')) / 10000 + + if rate_mult > 1.0 and dr < rate_limit then + dr = dr * rate_mult + if dr > 0.0001 then + redis.call('HSET', KEYS[1], 'dr', tostring(math.floor(dr * 10000))) + else + redis.call('HSET', KEYS[1], 'dr', '1') + end + elseif rate_mult < 1.0 and dr > (1.0 / rate_limit) then + dr = dr * rate_mult + if dr > 0.0001 then + redis.call('HSET', KEYS[1], 'dr', tostring(math.floor(dr * 10000))) + else + redis.call('HSET', KEYS[1], 'dr', '1') + end + end +end + +if tonumber(KEYS[6]) > 1 then + local rate_mult = tonumber(KEYS[4]) + local rate_limit = tonumber(KEYS[6]) + db = tonumber(redis.call('HGET', KEYS[1], 'db')) / 10000 + + if rate_mult > 1.0 and db < rate_limit then + db = db * rate_mult + if db > 0.0001 then + redis.call('HSET', KEYS[1], 'db', tostring(math.floor(db * 10000))) + else + redis.call('HSET', KEYS[1], 'db', '1') + end + elseif rate_mult < 1.0 and db > (1.0 / rate_limit) then + db = db * rate_mult + if db > 0.0001 then + redis.call('HSET', KEYS[1], 'db', tostring(math.floor(db * 10000))) + else + redis.call('HSET', KEYS[1], 'db', '1') + end + end +end + +local burst,pending = unpack(redis.call('HMGET', KEYS[1], 'b', 'p')) +burst,pending = tonumber(burst or '0'),tonumber(pending or '0') +if burst < 0 then burst = nrcpt else burst = burst + nrcpt end +if pending < nrcpt then pending = 0 else pending = pending - nrcpt end + +redis.call('HMSET', KEYS[1], 'b', tostring(burst), 'p', tostring(pending), 'l', KEYS[2]) +redis.call('EXPIRE', KEYS[1], KEYS[7]) + +return {tostring(burst), tostring(dr), tostring(db)}
\ No newline at end of file diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua index d61e3990f..520efc99e 100644 --- a/src/plugins/lua/ratelimit.lua +++ b/src/plugins/lua/ratelimit.lua @@ -54,154 +54,12 @@ local settings = { prefilter = true, } --- Checks bucket, updating it if needed --- KEYS[1] - prefix to update, e.g. RL_<triplet>_<seconds> --- KEYS[2] - current time in milliseconds --- KEYS[3] - bucket leak rate (messages per millisecond) --- KEYS[4] - bucket burst --- KEYS[5] - expire for a bucket --- KEYS[6] - number of recipients --- return 1 if message should be ratelimited and 0 if not --- Redis keys used: --- l - last hit --- b - current burst --- p - pending messages (those that are currently processing) --- dr - current dynamic rate multiplier (*10000) --- db - current dynamic burst multiplier (*10000) -local bucket_check_script = [[ - local last = redis.call('HGET', KEYS[1], 'l') - local now = tonumber(KEYS[2]) - local nrcpt = tonumber(KEYS[6]) - local dynr, dynb, leaked = 0, 0, 0 - if not last then - -- New bucket - redis.call('HMSET', KEYS[1], 'l', KEYS[2], 'b', '0', 'dr', '10000', 'db', '10000', 'p', tostring(nrcpt)) - redis.call('EXPIRE', KEYS[1], KEYS[5]) - return {0, '0', '1', '1', '0'} - end - - last = tonumber(last) - local burst,pending = unpack(redis.call('HMGET', KEYS[1], 'b', 'p')) - burst,pending = tonumber(burst or '0'),tonumber(pending or '0') - -- Sanity to avoid races - if burst < 0 then burst = 0 end - if pending < 0 then pending = 0 end - pending = pending + nrcpt -- this message - -- Perform leak - if burst + pending > 0 then - if burst > 0 and last < tonumber(KEYS[2]) then - local rate = tonumber(KEYS[3]) - dynr = tonumber(redis.call('HGET', KEYS[1], 'dr')) / 10000.0 - if dynr == 0 then dynr = 0.0001 end - rate = rate * dynr - leaked = ((now - last) * rate) - if leaked > burst then leaked = burst end - burst = burst - leaked - redis.call('HINCRBYFLOAT', KEYS[1], 'b', -(leaked)) - redis.call('HSET', KEYS[1], 'l', KEYS[2]) - end - - dynb = tonumber(redis.call('HGET', KEYS[1], 'db')) / 10000.0 - if dynb == 0 then dynb = 0.0001 end - - burst = burst + pending - if burst > 0 and (burst + tonumber(KEYS[6])) > tonumber(KEYS[4]) * dynb then - return {1, tostring(burst - pending), tostring(dynr), tostring(dynb), tostring(leaked)} - end - -- Increase pending if we allow ratelimit - redis.call('HINCRBY', KEYS[1], 'p', nrcpt) - else - burst = 0 - redis.call('HMSET', KEYS[1], 'b', '0', 'p', tostring(nrcpt)) - end - return {0, tostring(burst), tostring(dynr), tostring(dynb), tostring(leaked)} -]] +local bucket_check_script = "ratelimit_check.lua" local bucket_check_id --- Updates a bucket --- KEYS[1] - prefix to update, e.g. RL_<triplet>_<seconds> --- KEYS[2] - current time in milliseconds --- KEYS[3] - dynamic rate multiplier --- KEYS[4] - dynamic burst multiplier --- KEYS[5] - max dyn rate (min: 1/x) --- KEYS[6] - max burst rate (min: 1/x) --- KEYS[7] - expire for a bucket --- KEYS[8] - number of recipients (or increase rate) --- Redis keys used: --- l - last hit --- b - current burst --- p - messages pending (must be decreased by 1) --- dr - current dynamic rate multiplier --- db - current dynamic burst multiplier -local bucket_update_script = [[ - local last = redis.call('HGET', KEYS[1], 'l') - local now = tonumber(KEYS[2]) - local nrcpt = tonumber(KEYS[8]) - if not last then - -- New bucket (why??) - redis.call('HMSET', KEYS[1], 'l', KEYS[2], 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0') - redis.call('EXPIRE', KEYS[1], KEYS[7]) - return {1, 1, 1} - end - - local dr, db = 1.0, 1.0 - - if tonumber(KEYS[5]) > 1 then - local rate_mult = tonumber(KEYS[3]) - local rate_limit = tonumber(KEYS[5]) - dr = tonumber(redis.call('HGET', KEYS[1], 'dr')) / 10000 - - if rate_mult > 1.0 and dr < rate_limit then - dr = dr * rate_mult - if dr > 0.0001 then - redis.call('HSET', KEYS[1], 'dr', tostring(math.floor(dr * 10000))) - else - redis.call('HSET', KEYS[1], 'dr', '1') - end - elseif rate_mult < 1.0 and dr > (1.0 / rate_limit) then - dr = dr * rate_mult - if dr > 0.0001 then - redis.call('HSET', KEYS[1], 'dr', tostring(math.floor(dr * 10000))) - else - redis.call('HSET', KEYS[1], 'dr', '1') - end - end - end - - if tonumber(KEYS[6]) > 1 then - local rate_mult = tonumber(KEYS[4]) - local rate_limit = tonumber(KEYS[6]) - db = tonumber(redis.call('HGET', KEYS[1], 'db')) / 10000 - - if rate_mult > 1.0 and db < rate_limit then - db = db * rate_mult - if db > 0.0001 then - redis.call('HSET', KEYS[1], 'db', tostring(math.floor(db * 10000))) - else - redis.call('HSET', KEYS[1], 'db', '1') - end - elseif rate_mult < 1.0 and db > (1.0 / rate_limit) then - db = db * rate_mult - if db > 0.0001 then - redis.call('HSET', KEYS[1], 'db', tostring(math.floor(db * 10000))) - else - redis.call('HSET', KEYS[1], 'db', '1') - end - end - end - - local burst,pending = unpack(redis.call('HMGET', KEYS[1], 'b', 'p')) - burst,pending = tonumber(burst or '0'),tonumber(pending or '0') - if burst < 0 then burst = nrcpt else burst = burst + nrcpt end - if pending < nrcpt then pending = 0 else pending = pending - nrcpt end - - redis.call('HMSET', KEYS[1], 'b', tostring(burst), 'p', tostring(pending), 'l', KEYS[2]) - redis.call('EXPIRE', KEYS[1], KEYS[7]) - - return {tostring(burst), tostring(dr), tostring(db)} -]] +local bucket_update_script = "ratelimit_update.lua" local bucket_update_id -- message_func(task, limit_type, prefix, bucket, limit_key) @@ -210,9 +68,9 @@ local message_func = function(_, limit_type, _, _, _) end -local function load_scripts(cfg, ev_base) - bucket_check_id = lua_redis.add_redis_script(bucket_check_script, redis_params) - bucket_update_id = lua_redis.add_redis_script(bucket_update_script, redis_params) +local function load_scripts(_, _) + bucket_check_id = lua_redis.load_redis_script_from_file(bucket_check_script, redis_params) + bucket_update_id = lua_redis.load_redis_script_from_file(bucket_update_script, redis_params) end local limit_parser @@ -927,6 +785,6 @@ if opts then end end -rspamd_config:add_on_load(function(cfg, ev_base, worker) +rspamd_config:add_on_load(function(cfg, ev_base, _) load_scripts(cfg, ev_base) end) |