Software /
code /
prosody-modules
Changeset
2128:21bc4d7cddae
mod_firewall: Add support for throttling based on user-defined properties (experimental)
author | Matthew Wild <mwild1@gmail.com> |
---|---|
date | Fri, 18 Mar 2016 09:47:52 +0000 |
parents | 2127:59023dffbdd4 |
children | 2129:26334f4a8eb9 |
files | mod_firewall/conditions.lib.lua mod_firewall/definitions.lib.lua mod_firewall/mod_firewall.lua |
diffstat | 3 files changed, 64 insertions(+), 5 deletions(-) [+] |
line wrap: on
line diff
--- a/mod_firewall/conditions.lib.lua Fri Mar 18 09:45:02 2016 +0000 +++ b/mod_firewall/conditions.lib.lua Fri Mar 18 09:47:52 2016 +0000 @@ -182,8 +182,22 @@ return table.concat(conditions, " or "), { "time:hour,min" }; end -function condition_handlers.LIMIT(name) - return ("not throttle_%s:poll(1)"):format(name), { "throttle:"..name }; +function condition_handlers.LIMIT(spec) + local name, param = spec:match("^(%w+) on (.+)$"); + + if not name then + name = spec:match("^%w+$"); + if not name then + error("Unable to parse LIMIT specification"); + end + else + param = meta(("%q"):format(param)); + end + + if not param then + return ("not global_throttle_%s:poll(1)"):format(name), { "globalthrottle:"..name }; + end + return ("not multi_throttle_%s:poll_on(%s, 1)"):format(name, param), { "multithrottle:"..name }; end function condition_handlers.ORIGIN_MARKED(name_and_time)
--- a/mod_firewall/definitions.lib.lua Fri Mar 18 09:45:02 2016 +0000 +++ b/mod_firewall/definitions.lib.lua Fri Mar 18 09:47:52 2016 +0000 @@ -6,6 +6,9 @@ local set = require"util.set"; local new_throttle = require "util.throttle".create; +local new_cache = require "util.cache".new; + +local multirate_cache_size = module:get_option_number("firewall_multirate_cache_limit", 1000); function definition_handlers.ZONE(zone_name, zone_members) local zone_member_list = {}; @@ -15,10 +18,45 @@ return set.new(zone_member_list)._items; end +-- Helper function used by RATE handler +local function evict_only_unthrottled(name, throttle) + throttle:update(); + -- Check whether the throttle is at max balance (i.e. totally safe to forget about it) + if throttle.balance < throttle.max then + -- Not safe to forget + return false; + end +end + function definition_handlers.RATE(name, line) local rate = assert(tonumber(line:match("([%d.]+)")), "Unable to parse rate"); local burst = tonumber(line:match("%(%s*burst%s+([%d.]+)%s*%)")) or 1; - return new_throttle(rate*burst, burst); + local max_throttles = tonumber(line:match("%(%s*entries%s+([%d]+)%s*%)")) or multirate_cache_size; + + local cache = new_cache(max_throttles, evict_only_unthrottled); + + return { + single = function () + return new_throttle(rate*burst, burst); + end; + + multi = function () + return { + poll_on = function (_, key, amount) + assert(key, "no key"); + local throttle = cache:get(key); + if not throttle then + throttle = new_throttle(rate*burst, burst); + if not cache:set(key, throttle) then + module:log("warn", "Multirate '%s' has hit its maximum number of active throttles (%d), denying new events", name, max_throttles); + return false; + end + end + return throttle:poll(amount); + end; + } + end; + }; end return definition_handlers;
--- a/mod_firewall/mod_firewall.lua Fri Mar 18 09:45:02 2016 +0000 +++ b/mod_firewall/mod_firewall.lua Fri Mar 18 09:47:52 2016 +0000 @@ -108,11 +108,18 @@ return table.concat(defs, " "); end, depends = { "date_time" }; }; timestamp = { global_code = [[local get_time = require "socket".gettime]]; local_code = [[local current_timestamp = get_time()]]; }; - throttle = { + globalthrottle = { global_code = function (throttle) assert(idsafe(throttle), "Invalid rate limit name: "..throttle); assert(active_definitions.RATE[throttle], "Unknown rate limit: "..throttle); - return ("local throttle_%s = rates.%s;"):format(throttle, throttle); + return ("local global_throttle_%s = rates.%s:single();"):format(throttle, throttle); + end; + }; + multithrottle = { + global_code = function (throttle) + assert(idsafe(throttle), "Invalid rate limit name: "..throttle); + assert(active_definitions.RATE[throttle], "Unknown rate limit: "..throttle); + return ("local multi_throttle_%s = rates.%s:multi();"):format(throttle, throttle); end; }; };