Mercurial > prosody-modules

annotate mod_limits/mod_limits.lua @ 3529:6ac98c4dbbd3

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
mod_pubsub_github: Get event type from payload (should work with gitlab)
author Kim Alvefur <zash@zash.se>
date Sun, 31 Mar 2019 19:07:22 +0200
parents 88b16084eda7
children 1bb2a90398d3
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
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
62 session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", 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