Mercurial > prosody-hg

comparison plugins/mod_limits.lua @ 14133:4087c5b600b1 13.0

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
mod_limits: Support limiting total number of connections by IP address/block
author Matthew Wild <mwild1@gmail.com>
date Thu, 16 Apr 2026 18:54:56 +0100
parents f65302ea37b0
children e38916e22f59
comparison
equal deleted inserted replaced
14132:1e005ba71f0d 14133:4087c5b600b1
1 -- Because we deal with pre-authed sessions and streams we can't be host-specific 1 -- Because we deal with pre-authed sessions and streams we can't be host-specific
2 module:set_global(); 2 module:set_global();
3 3
4 local it = require "prosody.util.iterators";
4 local filters = require "prosody.util.filters"; 5 local filters = require "prosody.util.filters";
5 local throttle = require "prosody.util.throttle"; 6 local throttle = require "prosody.util.throttle";
6 local timer = require "prosody.util.timer"; 7 local timer = require "prosody.util.timer";
8
9 local pton = require "prosody.util.net".pton;
10
7 local ceil = math.ceil; 11 local ceil = math.ceil;
8 12
9 local limits_cfg = module:get_option("limits", {}); 13 local limits_cfg = module:get_option("limits", {});
10 local limits_resolution = module:get_option_period("limits_resolution", 1); 14 local limits_resolution = module:get_option_period("limits_resolution", 1);
11 15
48 bytes_per_second = 30 * 1024; 52 bytes_per_second = 30 * 1024;
49 burst_seconds = 2; 53 burst_seconds = 2;
50 }; 54 };
51 }; 55 };
52 56
57 local function get_truncated_ip(conn)
58 local session_ip = pton(conn:ip());
59 if #session_ip == 128 then
60 -- Use /64 for IPv6
61 session_ip = session_ip:sub(1, 8);
62 end
63 return session_ip;
64 end
65
66 local function limit_ip_sessions(session_type, max_sessions)
67 -- [packed_truncated_ip] = count
68 local current_ips = module:shared("ips."..session_type);
69 module:hook(session_type.."-connected", function (event)
70 local session_ip = get_truncated_ip(event.conn);
71 local new_count = (current_ips[session_ip] or 0) + 1;
72 current_ips[session_ip] = new_count;
73 if new_count > max_sessions then
74 module:log("warn", "Rejecting stream from %s (%d of %d current connections)", event.conn:ip(), new_count, max_sessions);
75 event.session:close({ condition = "policy-violation", text = "too many connections" });
76 end
77 end);
78
79 module:hook(session_type.."-closed", function (event)
80 local session_ip = get_truncated_ip(event.conn);
81 local new_count = (current_ips[session_ip] or 0) - 1;
82 if new_count < 0 then
83 module:log("warn", "IP counting failure");
84 new_count = 0;
85 end
86 if new_count == 0 then
87 current_ips[session_ip] = nil;
88 else
89 current_ips[session_ip] = new_count;
90 end
91 end);
92 end
93
53 for sess_type, sess_limits in pairs(limits_cfg) do 94 for sess_type, sess_limits in pairs(limits_cfg) do
54 limits[sess_type] = { 95 limits[sess_type] = {
55 bytes_per_second = parse_rate(sess_limits.rate, sess_type); 96 bytes_per_second = parse_rate(sess_limits.rate, sess_type);
56 burst_seconds = parse_burst(sess_limits.burst, sess_type); 97 burst_seconds = parse_burst(sess_limits.burst, sess_type);
57 }; 98 };
99 if sess_limits.max_connections_per_ip then
100 limit_ip_sessions(sess_type, sess_limits.max_connections_per_ip);
101 end
58 end 102 end
59 103
60 if not limits.s2sin and limits.s2s then 104 if not limits.s2sin and limits.s2s then
61 limits.s2sin = limits.s2s; 105 limits.s2sin = limits.s2s;
62 end 106 end
156 end 200 end
157 end); 201 end);
158 202
159 end 203 end
160 end 204 end
205
206 module:add_item("shell-command", {
207 section = "limits";
208 section_desc = "Commands to inspect rate-limiting state";
209 name = "connections";
210 desc = "View per-IP connection limit status";
211 args = {
212 { name = "ip", type = "string" };
213 };
214 handler = function(self, ip_str) -- luacheck: ignore 212/self
215 local check_ip = pton(ip_str);
216 if not check_ip then
217 return false, "Unable to parse IP address: "..ip_str;
218 end
219 if #check_ip == 128 then
220 check_ip = check_ip:sub(1, 8);
221 end
222
223 local found, any_limits = 0;
224
225 for sess_type, sess_limits in it.sorted_pairs(limits_cfg) do
226 if sess_limits.max_connections_per_ip then
227 any_limits = true;
228 local current_ips = module:shared("ips."..sess_type);
229 local count = current_ips[check_ip] or 0;
230 self.session.print(("%s: %d (limit %d)"):format(sess_type, count or 0, sess_limits.max_connections_per_ip));
231 found = found + count;
232 end
233 end
234
235 if not any_limits then
236 return false, "No per-IP limits are configured";
237 end
238
239 return true, ("Total of %d tracked connections with this IP"):format(found);
240 end;
241 });
242
243