Mercurial > prosody-modules

annotate mod_limits/mod_limits.lua @ 1268:854a3933cfcd

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