# HG changeset patch
# User Matthew Wild <mwild1@gmail.com>
# Date 1458294472 0
# Node ID 21bc4d7cddae2314f5e46b89c4e2f04e6bff952b
# Parent  59023dffbdd48ad3038204270404d368109ce40b
mod_firewall: Add support for throttling based on user-defined properties (experimental)

diff -r 59023dffbdd4 -r 21bc4d7cddae mod_firewall/conditions.lib.lua
--- 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)
diff -r 59023dffbdd4 -r 21bc4d7cddae mod_firewall/definitions.lib.lua
--- 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;
diff -r 59023dffbdd4 -r 21bc4d7cddae mod_firewall/mod_firewall.lua
--- 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;
 	};
 };