aboutsummaryrefslogtreecommitdiffstats
path: root/plugins/mod_limits.lua
diff options
context:
space:
mode:
Diffstat (limited to 'plugins/mod_limits.lua')
-rw-r--r--plugins/mod_limits.lua96
1 files changed, 96 insertions, 0 deletions
diff --git a/plugins/mod_limits.lua b/plugins/mod_limits.lua
new file mode 100644
index 00000000..2a6ee8a2
--- /dev/null
+++ b/plugins/mod_limits.lua
@@ -0,0 +1,96 @@
+-- Because we deal we pre-authed sessions and streams we can't be host-specific
+module:set_global();
+
+local filters = require "util.filters";
+local throttle = require "util.throttle";
+local timer = require "util.timer";
+
+local limits_cfg = module:get_option("limits", {});
+local limits_resolution = module:get_option_number("limits_resolution", 1);
+
+local default_bytes_per_second = 3000;
+local default_burst = 2;
+
+local rate_units = { b = 1, k = 3, m = 6, g = 9, t = 12 } -- Plan for the future.
+local function parse_rate(rate, sess_type)
+ local quantity, unit, exp;
+ if rate then
+ quantity, unit = rate:match("^(%d+) ?([^/]+)/s$");
+ exp = quantity and rate_units[unit:sub(1,1):lower()];
+ end
+ if not exp then
+ module:log("error", "Error parsing rate for %s: %q, using default rate (%d bytes/s)", sess_type, rate, default_bytes_per_second);
+ return default_bytes_per_second;
+ end
+ return quantity*(10^exp);
+end
+
+local function parse_burst(burst, sess_type)
+ if type(burst) == "string" then
+ burst = burst:match("^(%d+) ?s$");
+ end
+ local n_burst = tonumber(burst);
+ if not n_burst then
+ module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, tostring(burst), default_burst);
+ end
+ return n_burst or default_burst;
+end
+
+-- Process config option into limits table:
+-- limits = { c2s = { bytes_per_second = X, burst_seconds = Y } }
+local limits = {};
+
+for sess_type, sess_limits in pairs(limits_cfg) do
+ limits[sess_type] = {
+ bytes_per_second = parse_rate(sess_limits.rate, sess_type);
+ burst_seconds = parse_burst(sess_limits.burst, sess_type);
+ };
+end
+
+local default_filter_set = {};
+
+function default_filter_set.bytes_in(bytes, session)
+ local throttle = session.throttle;
+ if throttle then
+ local ok, balance, outstanding = throttle:poll(#bytes, true);
+ if not ok then
+ session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", throttle.max, #bytes, outstanding);
+ session.conn:pause(); -- Read no more data from the connection until there is no outstanding data
+ local outstanding_data = bytes:sub(-outstanding);
+ bytes = bytes:sub(1, #bytes-outstanding);
+ timer.add_task(limits_resolution, function ()
+ if not session.conn then return; end
+ if throttle:peek(#outstanding_data) then
+ session.log("debug", "Resuming paused session");
+ session.conn:resume();
+ end
+ -- Handle what we can of the outstanding data
+ session.data(outstanding_data);
+ end);
+ end
+ end
+ return bytes;
+end
+
+local type_filters = {
+ c2s = default_filter_set;
+ s2sin = default_filter_set;
+ s2sout = default_filter_set;
+};
+
+local function filter_hook(session)
+ local session_type = session.type:match("^[^_]+");
+ local filter_set, opts = type_filters[session_type], limits[session_type];
+ if opts then
+ session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
+ filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
+ end
+end
+
+function module.load()
+ filters.add_filter_hook(filter_hook);
+end
+
+function module.unload()
+ filters.remove_filter_hook(filter_hook);
+end