Mercurial > prosody-modules

comparison mod_limits/mod_limits.lua @ 738:92db76641b3f

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