From bf23daa22687b96406f5a1b54e1eb8b74765a1a8 Mon Sep 17 00:00:00 2001 From: left-try Date: Thu, 7 Nov 2024 18:48:12 +0600 Subject: [Feature] Add ratelimit util to manage last filled ratelimit buckets --- lualib/rspamadm/ratelimit.lua | 189 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 lualib/rspamadm/ratelimit.lua diff --git a/lualib/rspamadm/ratelimit.lua b/lualib/rspamadm/ratelimit.lua new file mode 100644 index 000000000..9acffc26a --- /dev/null +++ b/lualib/rspamadm/ratelimit.lua @@ -0,0 +1,189 @@ +local argparse = require 'argparse' +local redis = require 'lua_redis' +local logger = require 'rspamd_logger' + +local parser = argparse() + :name "ratelimit" + :description "Manage ratelimit functional" + :help_description_margin(32) + :command_target('command') + :require_command(true) + +local track_limits = parser:command 'track_limits' + :description 'Track last limits of ratelimit module' + +track_limits:option "-q --quantity" + :description("Number of limits to track") + :argname("") + :default(1) + + +local upgrade_bucket = parser:command 'upgrade_bucket' + :description 'Upgrade certain bucket' + +upgrade_bucket:option "-p --prefix" + :description("Prefix of bucket to operate with") + :argname("") + :args(1) +upgrade_bucket:option "-b --burst" + :description("Burst to set") + :argname("") + :args("?") +upgrade_bucket:option "-r --rate" + :description("Rate to set") + :argname("") + :args("?") +upgrade_bucket:option "-s --symbol" + :description("Symbol to set") + :argname("") + :args("?") +upgrade_bucket:option "-B --dynamic_burst" + :description("Dynamic burst to set") + :argname("") + :args("?") +upgrade_bucket:option "-R --dynamic_rate" + :description("Dynamic rate to set") + :argname("") + :args("?") + +local unblock_bucket = parser:command 'unblock_bucket' + :description 'Unblock certain bucket' + +unblock_bucket:option "-p --prefix" + :description("Prefix of bucket to operate with") + :argname("") + +local unblock_buckets = parser:command 'unblock_buckets' + :description("Unblock provided number of buckets(default: 1)") +unblock_buckets:option "-q --quantity" + :description("Number of buckets to ublock") + :argname("") + :default(1) + + +local redis_params +local lfb_cache_prefix = 'RL_cache_prefix' +local redis_attrs = { + config = rspamd_config, + ev_base = rspamadm_ev_base, + session = rspamadm_session, + log_obj = rspamd_config, + resolver = rspamadm_dns_resolver, +} + + +local function track_limits_handler(args) + for _ = 1, args.quantity do + local res, prefix = redis.request(redis_params, redis_attrs, + { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) + if res ~= 1 then + print('Redis request parameters are wrong') + end + local _, bucket_params = redis.request(redis_params, redis_attrs, + { 'HMGET', tostring(prefix[1]), 'l', 'b', 'r', 'dr', 'db' }) + local last = tonumber(bucket_params[1]) + local burst = tonumber(bucket_params[2]) + local rate = tonumber(bucket_params[3]) + local dynr = tonumber(bucket_params[4]) / 10000.0 + local dynb = tonumber(bucket_params[5]) / 10000.0 + + print(string.format('prefix: %s, last: %s, burst: %s, rate: %s, dynamic_rate: %s, dynamic_burst: %s', + prefix[1], last, burst, rate, dynr, dynb)) + end +end + +local function upgrade_bucket_handler(args) + if args.burst then + local res = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(args.prefix), 'b', tostring(args.burst) }) + if res ~= 1 then + print('Incorrect arguments for redis request for burst') + end + end + + if args.rate then + local res = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(args.prefix), 'r', tostring(args.rate) }) + if res ~= 1 then + print('Incorrect arguments for redis request for rate') + end + end + + if args.symbol then + local res = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(args.prefix), 's', tostring(args.symbol) }) + if res ~= 1 then + print('Incorrect arguments for redis request for symbol') + end + end + + if args.dynb then + local res = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(args.prefix), 'db', tostring(args.dynb) }) + if res ~= 1 then + print('Incorrect arguments for redis request for dynamic burst') + end + end + + if args.dynr then + local res = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(args.prefix), 'dr', tostring(args.dynr) }) + if res ~= 1 then + print('Incorrect arguments for redis request for dynamic rate') + end + end + +end + +local function unblock_bucket_helper(prefix) + local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(prefix), 'b', 0 }) + if res ~= 1 then + print('Failed to unblock bucket') + end +end + +local function unblock_bucket_handler(args) + unblock_bucket_helper(args.prefix) +end + +local function unblock_buckets_handler(args) + for _ = 1, args.quantity do + local res, prefix = redis.request(redis_params, redis_attrs, + { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) + if res ~= 1 then + print('Redis request parameters are wrong') + end + unblock_bucket_helper(prefix) + end +end + +local command_handlers = { + track_limits = track_limits_handler, + upgrade_bucket = upgrade_bucket_handler, + unblock_bucket = unblock_bucket_handler, + unblock_buckets = unblock_buckets_handler +} + +local function handler(args) + local cmd_opts = parser:parse(args) + + redis_params = redis.parse_redis_server('ratelimit') + if not redis_params then + logger.errx(rspamd_config, 'no servers are specified, disabling module') + os.exit(1) + end + + local f = command_handlers[cmd_opts.command] + if not f then + parser:error(string.format("command isn't implemented: %s", + cmd_opts.command)) + end + f(cmd_opts) +end + +return { + name = 'ratelimit', + aliases = { 'ratelimit' }, + handler = handler, + description = parser._description +} \ No newline at end of file -- cgit v1.2.3 From 1f54d0dd714b2580820f480136fa2e64939374c5 Mon Sep 17 00:00:00 2001 From: left-try Date: Thu, 7 Nov 2024 18:48:46 +0600 Subject: [CI] Add ability to run GA on any push from any branch --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 07c4cc4d0..c6aa3f74b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,7 +6,7 @@ on: - master push: branches: - - master + - '**' jobs: linters: -- cgit v1.2.3 From 730e1288773636115d59e07a92c575b250ec0d52 Mon Sep 17 00:00:00 2001 From: left-try Date: Thu, 14 Nov 2024 09:39:20 +0600 Subject: [Minor] Small clean up for error handling --- lualib/rspamadm/ratelimit.lua | 38 ++++++++++++++++++++++++++------------ 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/lualib/rspamadm/ratelimit.lua b/lualib/rspamadm/ratelimit.lua index 9acffc26a..524f8853d 100644 --- a/lualib/rspamadm/ratelimit.lua +++ b/lualib/rspamadm/ratelimit.lua @@ -77,15 +77,18 @@ local function track_limits_handler(args) local res, prefix = redis.request(redis_params, redis_attrs, { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) if res ~= 1 then - print('Redis request parameters are wrong') + logger.errx('Redis request parameters are wrong') + os.exit(1) end + local _, bucket_params = redis.request(redis_params, redis_attrs, { 'HMGET', tostring(prefix[1]), 'l', 'b', 'r', 'dr', 'db' }) - local last = tonumber(bucket_params[1]) + + local last = tonumber(bucket_params[1]) local burst = tonumber(bucket_params[2]) - local rate = tonumber(bucket_params[3]) - local dynr = tonumber(bucket_params[4]) / 10000.0 - local dynb = tonumber(bucket_params[5]) / 10000.0 + local rate = tonumber(bucket_params[3]) + local dynr = tonumber(bucket_params[4]) / 10000.0 + local dynb = tonumber(bucket_params[5]) / 10000.0 print(string.format('prefix: %s, last: %s, burst: %s, rate: %s, dynamic_rate: %s, dynamic_burst: %s', prefix[1], last, burst, rate, dynr, dynb)) @@ -93,11 +96,16 @@ local function track_limits_handler(args) end local function upgrade_bucket_handler(args) + if args.prefix == nil or args.prefix == "" then + logger.errx('Prefix is empty or nil') + end + if args.burst then local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(args.prefix), 'b', tostring(args.burst) }) if res ~= 1 then - print('Incorrect arguments for redis request for burst') + logger.errx('Incorrect arguments for redis request for burst or prefix is incorrect') + os.exit(1) end end @@ -105,7 +113,8 @@ local function upgrade_bucket_handler(args) local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(args.prefix), 'r', tostring(args.rate) }) if res ~= 1 then - print('Incorrect arguments for redis request for rate') + logger.errx('Incorrect arguments for redis request for rate or prefix is incorrect') + os.exit(1) end end @@ -113,7 +122,8 @@ local function upgrade_bucket_handler(args) local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(args.prefix), 's', tostring(args.symbol) }) if res ~= 1 then - print('Incorrect arguments for redis request for symbol') + logger.errx('Incorrect arguments for redis request for symbol or prefix is incorrect') + os.exit(1) end end @@ -121,7 +131,8 @@ local function upgrade_bucket_handler(args) local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(args.prefix), 'db', tostring(args.dynb) }) if res ~= 1 then - print('Incorrect arguments for redis request for dynamic burst') + logger.errx('Incorrect arguments for redis request for dynamic burst or prefix is incorrect') + os.exit(1) end end @@ -129,7 +140,8 @@ local function upgrade_bucket_handler(args) local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(args.prefix), 'dr', tostring(args.dynr) }) if res ~= 1 then - print('Incorrect arguments for redis request for dynamic rate') + logger.errx('Incorrect arguments for redis request for dynamic rate or prefix is incorrect') + os.exit(1) end end @@ -138,7 +150,8 @@ end local function unblock_bucket_helper(prefix) local res = redis.request(redis_params, redis_attrs, { 'HSET', tostring(prefix), 'b', 0 }) if res ~= 1 then - print('Failed to unblock bucket') + logger.errx('Failed to unblock bucket') + os.exit(1) end end @@ -151,7 +164,8 @@ local function unblock_buckets_handler(args) local res, prefix = redis.request(redis_params, redis_attrs, { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) if res ~= 1 then - print('Redis request parameters are wrong') + logger.errx('Redis request parameters are wrong') + os.exit(1) end unblock_bucket_helper(prefix) end -- cgit v1.2.3 From 27cda90084bb85e04e53858219554355a3c40189 Mon Sep 17 00:00:00 2001 From: left-try Date: Thu, 14 Nov 2024 09:51:58 +0600 Subject: [Minor] Small clean up --- lualib/rspamadm/ratelimit.lua | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/lualib/rspamadm/ratelimit.lua b/lualib/rspamadm/ratelimit.lua index 524f8853d..c74a1a880 100644 --- a/lualib/rspamadm/ratelimit.lua +++ b/lualib/rspamadm/ratelimit.lua @@ -21,9 +21,8 @@ track_limits:option "-q --quantity" local upgrade_bucket = parser:command 'upgrade_bucket' :description 'Upgrade certain bucket' -upgrade_bucket:option "-p --prefix" +upgrade_bucket:argument "prefix" :description("Prefix of bucket to operate with") - :argname("") :args(1) upgrade_bucket:option "-b --burst" :description("Burst to set") @@ -49,9 +48,9 @@ upgrade_bucket:option "-R --dynamic_rate" local unblock_bucket = parser:command 'unblock_bucket' :description 'Unblock certain bucket' -unblock_bucket:option "-p --prefix" +unblock_bucket:argument "prefix" :description("Prefix of bucket to operate with") - :argname("") + :args(1) local unblock_buckets = parser:command 'unblock_buckets' :description("Unblock provided number of buckets(default: 1)") -- cgit v1.2.3 From f363b0643da1125c0ae614724c3281880422d054 Mon Sep 17 00:00:00 2001 From: left-try Date: Sat, 16 Nov 2024 05:06:52 +0600 Subject: [Minor] Clean up code and fix some mistakes --- lualib/rspamadm/ratelimit.lua | 66 +++++++++++++++++++++++-------------------- 1 file changed, 36 insertions(+), 30 deletions(-) diff --git a/lualib/rspamadm/ratelimit.lua b/lualib/rspamadm/ratelimit.lua index c74a1a880..558554b10 100644 --- a/lualib/rspamadm/ratelimit.lua +++ b/lualib/rspamadm/ratelimit.lua @@ -9,7 +9,7 @@ local parser = argparse() :command_target('command') :require_command(true) -local track_limits = parser:command 'track_limits' +local track_limits = parser:command 'track' :description 'Track last limits of ratelimit module' track_limits:option "-q --quantity" @@ -18,12 +18,11 @@ track_limits:option "-q --quantity" :default(1) -local upgrade_bucket = parser:command 'upgrade_bucket' - :description 'Upgrade certain bucket' - +local upgrade_bucket = parser:command 'upgrade' + :description 'Upgrade certain bucket or top limit bucket' upgrade_bucket:argument "prefix" :description("Prefix of bucket to operate with") - :args(1) + :args("?") upgrade_bucket:option "-b --burst" :description("Burst to set") :argname("") @@ -45,19 +44,17 @@ upgrade_bucket:option "-R --dynamic_rate" :argname("") :args("?") -local unblock_bucket = parser:command 'unblock_bucket' - :description 'Unblock certain bucket' - -unblock_bucket:argument "prefix" - :description("Prefix of bucket to operate with") - :args(1) +local unblock_bucket = parser:command 'unblock' + :description 'Unblock certain bucket or number of buckets(default: 1)' -local unblock_buckets = parser:command 'unblock_buckets' - :description("Unblock provided number of buckets(default: 1)") -unblock_buckets:option "-q --quantity" - :description("Number of buckets to ublock") - :argname("") - :default(1) +parser:mutex( + unblock_bucket:argument "prefix" + :description("Prefix of bucket to operate with") + :args(1), + unblock_bucket:argument "quantity" + :description("Number of buckets to ublock") + :default(1) +) local redis_params @@ -72,9 +69,9 @@ local redis_attrs = { local function track_limits_handler(args) - for _ = 1, args.quantity do + for i = 1, args.quantity do local res, prefix = redis.request(redis_params, redis_attrs, - { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) + { 'ZRANGE', lfb_cache_prefix, -i, -i }) if res ~= 1 then logger.errx('Redis request parameters are wrong') os.exit(1) @@ -95,13 +92,20 @@ local function track_limits_handler(args) end local function upgrade_bucket_handler(args) - if args.prefix == nil or args.prefix == "" then - logger.errx('Prefix is empty or nil') + local prefix = args.prefix + if prefix == nil or prefix == "" then + local res = nil + res, prefix = redis.request(redis_params, redis_attrs, + { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) + if res ~= 1 then + logger.errx('Redis request parameters are wrong') + os.exit(1) + end end if args.burst then local res = redis.request(redis_params, redis_attrs, - { 'HSET', tostring(args.prefix), 'b', tostring(args.burst) }) + { 'HSET', tostring(prefix), 'b', tostring(args.burst) }) if res ~= 1 then logger.errx('Incorrect arguments for redis request for burst or prefix is incorrect') os.exit(1) @@ -110,7 +114,7 @@ local function upgrade_bucket_handler(args) if args.rate then local res = redis.request(redis_params, redis_attrs, - { 'HSET', tostring(args.prefix), 'r', tostring(args.rate) }) + { 'HSET', tostring(prefix), 'r', tostring(args.rate) }) if res ~= 1 then logger.errx('Incorrect arguments for redis request for rate or prefix is incorrect') os.exit(1) @@ -119,7 +123,7 @@ local function upgrade_bucket_handler(args) if args.symbol then local res = redis.request(redis_params, redis_attrs, - { 'HSET', tostring(args.prefix), 's', tostring(args.symbol) }) + { 'HSET', tostring(prefix), 's', tostring(args.symbol) }) if res ~= 1 then logger.errx('Incorrect arguments for redis request for symbol or prefix is incorrect') os.exit(1) @@ -128,7 +132,7 @@ local function upgrade_bucket_handler(args) if args.dynb then local res = redis.request(redis_params, redis_attrs, - { 'HSET', tostring(args.prefix), 'db', tostring(args.dynb) }) + { 'HSET', tostring(prefix), 'db', tostring(args.dynb) }) if res ~= 1 then logger.errx('Incorrect arguments for redis request for dynamic burst or prefix is incorrect') os.exit(1) @@ -137,7 +141,7 @@ local function upgrade_bucket_handler(args) if args.dynr then local res = redis.request(redis_params, redis_attrs, - { 'HSET', tostring(args.prefix), 'dr', tostring(args.dynr) }) + { 'HSET', tostring(prefix), 'dr', tostring(args.dynr) }) if res ~= 1 then logger.errx('Incorrect arguments for redis request for dynamic rate or prefix is incorrect') os.exit(1) @@ -155,13 +159,16 @@ local function unblock_bucket_helper(prefix) end local function unblock_bucket_handler(args) + if (args.prefix == nil or args.prefix == "") then + unblock_bucket_handler(args) + end unblock_bucket_helper(args.prefix) end local function unblock_buckets_handler(args) - for _ = 1, args.quantity do + for i = 1, args.quantity do local res, prefix = redis.request(redis_params, redis_attrs, - { 'ZRANGE', lfb_cache_prefix, '-1', '-1' }) + { 'ZRANGE', lfb_cache_prefix, -i, -i }) if res ~= 1 then logger.errx('Redis request parameters are wrong') os.exit(1) @@ -173,8 +180,7 @@ end local command_handlers = { track_limits = track_limits_handler, upgrade_bucket = upgrade_bucket_handler, - unblock_bucket = unblock_bucket_handler, - unblock_buckets = unblock_buckets_handler + unblock_bucket = unblock_bucket_handler } local function handler(args) -- cgit v1.2.3 From 21ec2578ad499dcdeefc79caf23af481d23b48f2 Mon Sep 17 00:00:00 2001 From: left-try Date: Sat, 16 Nov 2024 05:18:45 +0600 Subject: [Minor] Make unused function used --- lualib/rspamadm/ratelimit.lua | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/lualib/rspamadm/ratelimit.lua b/lualib/rspamadm/ratelimit.lua index 558554b10..5258f59e3 100644 --- a/lualib/rspamadm/ratelimit.lua +++ b/lualib/rspamadm/ratelimit.lua @@ -158,13 +158,6 @@ local function unblock_bucket_helper(prefix) end end -local function unblock_bucket_handler(args) - if (args.prefix == nil or args.prefix == "") then - unblock_bucket_handler(args) - end - unblock_bucket_helper(args.prefix) -end - local function unblock_buckets_handler(args) for i = 1, args.quantity do local res, prefix = redis.request(redis_params, redis_attrs, @@ -177,6 +170,13 @@ local function unblock_buckets_handler(args) end end +local function unblock_bucket_handler(args) + if (args.prefix == nil or args.prefix == "") then + unblock_buckets_handler(args) + end + unblock_bucket_helper(args.prefix) +end + local command_handlers = { track_limits = track_limits_handler, upgrade_bucket = upgrade_bucket_handler, -- cgit v1.2.3