Mercurial > prosody-modules
annotate mod_limits/mod_limits.lua @ 1114:6c0e1f9926f6
mod_mam: Swap first and last first for a simpler argument to RSM in reverse queries
author | Kim Alvefur <zash@zash.se> |
---|---|
date | Sat, 13 Jul 2013 17:00:13 +0200 |
parents | 92db76641b3f |
children | 1c126c49f5c1 |
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 |