Mercurial > prosody-modules

annotate mod_limits/mod_limits.lua @ 4300:3f3b672b7616

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
mod_vcard_muc: Pass room object around instead of JID, hopefully fixing traceback More efficient to pass the object around instead of using the JID and looking up the object when needed. It seems in some (undetermined) cases get_room_from_jid(room.jid) is nil.
author Matthew Wild <mwild1@gmail.com>
date Tue, 15 Dec 2020 10:49:11 +0000
parents 1bb2a90398d3
children
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";
2777
55a7ef2fb628 mod_limits: Handle fractional outstanding balance (imported from prosody 25237002aba4)
Matthew Wild <mwild1@gmail.com>
parents: 2057
diff changeset
11 local ceil = math.ceil;
738
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
12
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
13 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
14 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
15
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
16 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
17 local default_burst = 2;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
18
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
19 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
20 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
21 local quantity, unit, exp;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
22 if rate then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
23 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
24 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
25 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
26 if not exp then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
27 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
28 return default_bytes_per_second;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
29 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
30 return quantity*(10^exp);
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
31 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
32
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
33 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
34 if type(burst) == "string" then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
35 burst = burst:match("^(%d+) ?s$");
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
36 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
37 local n_burst = tonumber(burst);
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
38 if not n_burst then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
39 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
40 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
41 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
42 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
43
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
44 -- 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
45 -- 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
46 local limits = {};
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
47
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
48 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
49 limits[sess_type] = {
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
50 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
51 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
52 };
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
53 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
54
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
55 local default_filter_set = {};
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
56
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
57 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
58 local throttle = session.throttle;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
59 if throttle then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
60 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
61 if not ok then
3542
1bb2a90398d3 mod_limits: log throttled JID
Georg Lukas <georg@op-co.de>
parents: 2885
diff changeset
62 session.log("debug", "Session %q over rate limit (%d) with %d (by %d), pausing", session.full_jid or session.from_host or session.to_host, throttle.max, #bytes, outstanding);
2777
55a7ef2fb628 mod_limits: Handle fractional outstanding balance (imported from prosody 25237002aba4)
Matthew Wild <mwild1@gmail.com>
parents: 2057
diff changeset
63 outstanding = ceil(outstanding);
738
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
64 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
65 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
66 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
67 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
68 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
69 if throttle:peek(#outstanding_data) then
2057
1c126c49f5c1 mod_limits: Add newline between statements on long line
Kim Alvefur <zash@zash.se>
parents: 738
diff changeset
70 session.log("debug", "Resuming paused session");
1c126c49f5c1 mod_limits: Add newline between statements on long line
Kim Alvefur <zash@zash.se>
parents: 738
diff changeset
71 session.conn:resume();
738
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
72 end
2885
88b16084eda7 mod_limits: Add debug logging just before we feed data into stream
Matthew Wild <mwild1@gmail.com>
parents: 2777
diff changeset
73 session.log("debug", "mod_limits feeding %d bytes of delayed data into stream", #outstanding_data);
738
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
74 -- 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
75 session.data(outstanding_data);
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 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
78 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
79 return bytes;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
80 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
81
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
82 local type_filters = {
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
83 c2s = default_filter_set;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
84 s2sin = default_filter_set;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
85 s2sout = default_filter_set;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
86 };
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
87
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
88 local function filter_hook(session)
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
89 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
90 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
91 if opts then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
92 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
93 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
94 end
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.load()
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
98 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
99 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
100
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
101 function module.unload()
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
102 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
103 end