Mercurial > prosody-modules

annotate mod_limits/mod_limits.lua @ 2608:362ca94192ee

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
mod_smacks: Add resumed session to event "smacks-hibernation-end" Older versions of this event only have the "intermediate" session in event.session (the one used to resume the existing session), but not the resumed one. This adds event.resumed which contains the resumed one alongside to event.session.
author tmolitor <thilo@eightysoft.de>
date Sat, 11 Mar 2017 01:37:28 +0100
parents 1c126c49f5c1
children 55a7ef2fb628
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
2057
1c126c49f5c1 mod_limits: Add newline between statements on long line
Kim Alvefur <zash@zash.se>
parents: 738
diff changeset
68 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
69 session.conn:resume();
738
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
70 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
71 -- 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
72 session.data(outstanding_data);
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 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
76 return bytes;
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
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
79 local type_filters = {
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
80 c2s = default_filter_set;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
81 s2sin = default_filter_set;
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
82 s2sout = default_filter_set;
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
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
85 local function filter_hook(session)
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
86 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
87 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
88 if opts then
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
89 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
90 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
91 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
92 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
93
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
94 function module.load()
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
95 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
96 end
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
97
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
98 function module.unload()
92db76641b3f mod_limits: Import to prosody-modules, connection-level rate limiting
Matthew Wild <mwild1@gmail.com>
parents:
diff changeset
99 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
100 end