diff options
Diffstat (limited to 'lualib')
| -rw-r--r-- | lualib/rspamadm/ratelimit.lua | 248 |
1 files changed, 133 insertions, 115 deletions
diff --git a/lualib/rspamadm/ratelimit.lua b/lualib/rspamadm/ratelimit.lua index 5258f59e3..1dd42472f 100644 --- a/lualib/rspamadm/ratelimit.lua +++ b/lualib/rspamadm/ratelimit.lua @@ -3,11 +3,15 @@ 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) + :name "ratelimit" + :description "Manage ratelimit functional" + :help_description_margin(32) + :command_target('command') + :require_command(true) +parser:option "-c --config" + :description "Path to config file" + :argname("<cfg>") + :default(rspamd_paths["CONFDIR"] .. "/" .. "rspamd.conf") local track_limits = parser:command 'track' :description 'Track last limits of ratelimit module' @@ -17,7 +21,6 @@ track_limits:option "-q --quantity" :argname("<quantity>") :default(1) - local upgrade_bucket = parser:command 'upgrade' :description 'Upgrade certain bucket or top limit bucket' upgrade_bucket:argument "prefix" @@ -56,153 +59,168 @@ parser:mutex( :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, + 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 i = 1, args.quantity do - local res, prefix = redis.request(redis_params, redis_attrs, - { 'ZRANGE', lfb_cache_prefix, -i, -i }) - if res ~= 1 then - 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 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)) + for i = 1, args.quantity do + local res, prefix = redis.request(redis_params, redis_attrs, + { 'ZRANGE', lfb_cache_prefix, -i, -i }) + if not res then + logger.errx('Redis request error: %s', prefix) + os.exit(1) + end + + if #prefix == 1 then + 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 end local function upgrade_bucket_handler(args) - 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 + 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 - - if args.burst then - local res = redis.request(redis_params, redis_attrs, - { '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) - end + end + + if args.burst then + local res, err = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(prefix), 'b', tostring(args.burst) }) + if not res then + logger.errx('Incorrect arguments for redis request for burst or prefix is incorrect: %s', err) + os.exit(1) end - - if args.rate then - local res = redis.request(redis_params, redis_attrs, - { '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) - end + end + + if args.rate then + local res, err = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(prefix), 'r', tostring(args.rate) }) + if not res then + logger.errx('Incorrect arguments for redis request for rate or prefix is incorrect: %s', err) + os.exit(1) end - - if args.symbol then - local res = redis.request(redis_params, redis_attrs, - { '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) - end + end + + if args.symbol then + local res, err = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(prefix), 's', tostring(args.symbol) }) + if not res then + logger.errx('Incorrect arguments for redis request for symbol or prefix is incorrect: %s', err) + os.exit(1) end - - if args.dynb then - local res = redis.request(redis_params, redis_attrs, - { '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) - end + end + + if args.dynb then + local res, err = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(prefix), 'db', tostring(args.dynb) }) + if not res then + logger.errx('Incorrect arguments for redis request for dynamic burst or prefix is incorrect: %s', err) + os.exit(1) end - - if args.dynr then - local res = redis.request(redis_params, redis_attrs, - { '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) - end + end + + if args.dynr then + local res, err = redis.request(redis_params, redis_attrs, + { 'HSET', tostring(prefix), 'dr', tostring(args.dynr) }) + if not res then + logger.errx('Incorrect arguments for redis request for dynamic rate or prefix is incorrect: %s', err) + os.exit(1) 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 - logger.errx('Failed to unblock bucket') - os.exit(1) - end + local res, err = redis.request(redis_params, redis_attrs, { 'HSET', tostring(prefix), 'b', 0 }) + if not res then + logger.errx('Failed to unblock bucket: %s', err) + os.exit(1) + end end local function unblock_buckets_handler(args) - for i = 1, args.quantity do - local res, prefix = redis.request(redis_params, redis_attrs, - { 'ZRANGE', lfb_cache_prefix, -i, -i }) - if res ~= 1 then - logger.errx('Redis request parameters are wrong') - os.exit(1) - end - unblock_bucket_helper(prefix) + for i = 1, args.quantity do + local res, prefix = redis.request(redis_params, redis_attrs, + { 'ZRANGE', lfb_cache_prefix, -i, -i }) + if not res then + logger.errx('Redis request parameters are wrong: %s', prefix) + os.exit(1) end + unblock_bucket_helper(prefix) + 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) + 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, - unblock_bucket = unblock_bucket_handler + track = track_limits_handler, + upgrade = upgrade_bucket_handler, + unblock = unblock_bucket_handler } local function handler(args) - local cmd_opts = parser:parse(args) + local function load_config(opts) + local _r, err = rspamd_config:load_ucl(opts['config']) - 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) + if not _r then + logger.errx('cannot parse %s: %s', opts['config'], err) + 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)) + _r, err = rspamd_config:parse_rcl({ 'logging', 'worker' }) + if not _r then + logger.errx('cannot process %s: %s', opts['config'], err) + os.exit(1) end - f(cmd_opts) + end + local cmd_opts = parser:parse(args) + load_config(cmd_opts) + + redis_params = redis.parse_redis_server('ratelimit') + if not redis_params then + logger.errx(rspamd_config, 'no servers are specified, cannot work with rate limits') + 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 + name = 'ratelimit', + aliases = { 'ratelimit' }, + handler = handler, + description = parser._description }
\ No newline at end of file |