prosody
bf8f2da84007
plugins/mod_limits.lua

Find changesets by keywords by author, files, the commit message, revision number or hash, or revset expression.

Software / code / prosody

Comparison

plugins/mod_limits.lua @ 11200:bf8f2da84007

Merge 0.11->trunk
author Kim Alvefur <zash@zash.se>
date Thu, 05 Nov 2020 22:31:25 +0100
parent 10551:27b275633156
child 11560:3bbb1af92514
comparison
equal deleted inserted replaced
11199:6c7c50a4de32 11200:bf8f2da84007
30 if type(burst) == "string" then 30 if type(burst) == "string" then
31 burst = burst:match("^(%d+) ?s$"); 31 burst = burst:match("^(%d+) ?s$");
32 end 32 end
33 local n_burst = tonumber(burst); 33 local n_burst = tonumber(burst);
34 if not n_burst then 34 if not n_burst then
35 module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, tostring(burst), default_burst); 35 module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, burst, default_burst);
36 end 36 end
37 return n_burst or default_burst; 37 return n_burst or default_burst;
38 end 38 end
39 39
40 -- Process config option into limits table: 40 -- Process config option into limits table:
49 end 49 end
50 50
51 local default_filter_set = {}; 51 local default_filter_set = {};
52 52
53 function default_filter_set.bytes_in(bytes, session) 53 function default_filter_set.bytes_in(bytes, session)
54 local sess_throttle = session.throttle; 54 local sess_throttle = session.throttle;
55 if sess_throttle then 55 if sess_throttle then
56 local ok, balance, outstanding = sess_throttle:poll(#bytes, true); 56 local ok, _, outstanding = sess_throttle:poll(#bytes, true);
57 if not ok then 57 if not ok then
58 session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding); 58 session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding);
59 outstanding = ceil(outstanding); 59 outstanding = ceil(outstanding);
60 session.conn:pause(); -- Read no more data from the connection until there is no outstanding data 60 session.conn:pause(); -- Read no more data from the connection until there is no outstanding data
61 local outstanding_data = bytes:sub(-outstanding); 61 local outstanding_data = bytes:sub(-outstanding);
62 bytes = bytes:sub(1, #bytes-outstanding); 62 bytes = bytes:sub(1, #bytes-outstanding);
63 timer.add_task(limits_resolution, function () 63 timer.add_task(limits_resolution, function ()
64 if not session.conn then return; end 64 if not session.conn then return; end
65 if sess_throttle:peek(#outstanding_data) then 65 if sess_throttle:peek(#outstanding_data) then
66 session.log("debug", "Resuming paused session"); 66 session.log("debug", "Resuming paused session");
67 session.conn:resume(); 67 session.conn:resume();
68 end 68 end
69 -- Handle what we can of the outstanding data 69 -- Handle what we can of the outstanding data
70 session.data(outstanding_data); 70 session.data(outstanding_data);
82 82
83 local function filter_hook(session) 83 local function filter_hook(session)
84 local session_type = session.type:match("^[^_]+"); 84 local session_type = session.type:match("^[^_]+");
85 local filter_set, opts = type_filters[session_type], limits[session_type]; 85 local filter_set, opts = type_filters[session_type], limits[session_type];
86 if opts then 86 if opts then
87 session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds); 87 if session.conn and session.conn.setlimit then
88 filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000); 88 session.conn:setlimit(opts.bytes_per_second);
89 -- Currently no burst support
90 else
91 session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
92 filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
93 end
89 end 94 end
90 end 95 end
91 96
92 function module.load() 97 function module.load()
93 filters.add_filter_hook(filter_hook); 98 filters.add_filter_hook(filter_hook);
94 end 99 end
95 100
96 function module.unload() 101 function module.unload()
97 filters.remove_filter_hook(filter_hook); 102 filters.remove_filter_hook(filter_hook);
98 end 103 end
104
105 function module.add_host(module)
106 local unlimited_jids = module:get_option_inherited_set("unlimited_jids", {});
107
108 if not unlimited_jids:empty() then
109 module:hook("authentication-success", function (event)
110 local session = event.session;
111 local session_type = session.type:match("^[^_]+");
112 local jid = session.username .. "@" .. session.host;
113 if unlimited_jids:contains(jid) then
114 if session.conn and session.conn.setlimit then
115 session.conn:setlimit(0);
116 -- Currently no burst support
117 else
118 local filter_set = type_filters[session_type];
119 filters.remove_filter(session, "bytes/in", filter_set.bytes_in);
120 session.throttle = nil;
121 end
122 end
123 end);
124 end
125 end