prosody
cdffe33efae4
plugins/mod_limits.lua

Find changesets by keywords by author, files, the commit message, revision number or hash, or revset expression.

Software / code / prosody

Comparison

plugins/mod_limits.lua @ 8256:cdffe33efae4

mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129)
author Matthew Wild <mwild1@gmail.com>
date Sat, 23 Sep 2017 13:29:54 +0100
child 8269:25237002aba4
comparison
equal deleted inserted replaced
8255:d70d4c1ac17a 8256:cdffe33efae4
1 -- Because we deal we pre-authed sessions and streams we can't be host-specific
2 module:set_global();
3
4 local filters = require "util.filters";
5 local throttle = require "util.throttle";
6 local timer = require "util.timer";
7
8 local limits_cfg = module:get_option("limits", {});
9 local limits_resolution = module:get_option_number("limits_resolution", 1);
10
11 local default_bytes_per_second = 3000;
12 local default_burst = 2;
13
14 local rate_units = { b = 1, k = 3, m = 6, g = 9, t = 12 } -- Plan for the future.
15 local function parse_rate(rate, sess_type)
16 local quantity, unit, exp;
17 if rate then
18 quantity, unit = rate:match("^(%d+) ?([^/]+)/s$");
19 exp = quantity and rate_units[unit:sub(1,1):lower()];
20 end
21 if not exp then
22 module:log("error", "Error parsing rate for %s: %q, using default rate (%d bytes/s)", sess_type, rate, default_bytes_per_second);
23 return default_bytes_per_second;
24 end
25 return quantity*(10^exp);
26 end
27
28 local function parse_burst(burst, sess_type)
29 if type(burst) == "string" then
30 burst = burst:match("^(%d+) ?s$");
31 end
32 local n_burst = tonumber(burst);
33 if not n_burst then
34 module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, tostring(burst), default_burst);
35 end
36 return n_burst or default_burst;
37 end
38
39 -- Process config option into limits table:
40 -- limits = { c2s = { bytes_per_second = X, burst_seconds = Y } }
41 local limits = {};
42
43 for sess_type, sess_limits in pairs(limits_cfg) do
44 limits[sess_type] = {
45 bytes_per_second = parse_rate(sess_limits.rate, sess_type);
46 burst_seconds = parse_burst(sess_limits.burst, sess_type);
47 };
48 end
49
50 local default_filter_set = {};
51
52 function default_filter_set.bytes_in(bytes, session)
53 local throttle = session.throttle;
54 if throttle then
55 local ok, balance, outstanding = throttle:poll(#bytes, true);
56 if not ok then
57 session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", throttle.max, #bytes, outstanding);
58 session.conn:pause(); -- Read no more data from the connection until there is no outstanding data
59 local outstanding_data = bytes:sub(-outstanding);
60 bytes = bytes:sub(1, #bytes-outstanding);
61 timer.add_task(limits_resolution, function ()
62 if not session.conn then return; end
63 if throttle:peek(#outstanding_data) then
64 session.log("debug", "Resuming paused session");
65 session.conn:resume();
66 end
67 -- Handle what we can of the outstanding data
68 session.data(outstanding_data);
69 end);
70 end
71 end
72 return bytes;
73 end
74
75 local type_filters = {
76 c2s = default_filter_set;
77 s2sin = default_filter_set;
78 s2sout = default_filter_set;
79 };
80
81 local function filter_hook(session)
82 local session_type = session.type:match("^[^_]+");
83 local filter_set, opts = type_filters[session_type], limits[session_type];
84 if opts then
85 session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
86 filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
87 end
88 end
89
90 function module.load()
91 filters.add_filter_hook(filter_hook);
92 end
93
94 function module.unload()
95 filters.remove_filter_hook(filter_hook);
96 end