prosody-hg: plugins/mod_limits.lua annotate

annotate plugins/mod_limits.lua @ 14144:e38916e22f59 13.0

mod_limits: Only count connections with associated sessions I suspect some (all??) network backends and some edge cases cause ondisconnect to be called more than once. mod_c2s and mod_s2s already handle this, skipping cleanup if there is no session associated. Sessions are assigned very early on in onconnect, so it seems like a reliable way of counting.

author	Matthew Wild <mwild1@gmail.com>
date	Fri, 17 Apr 2026 19:42:44 +0100
parents	4087c5b600b1
children	4b8cf9420cf1

rev	line source
8453 6b3e7fddd723 mod_limits: Fix typo in comment Kim Alvefur <zash@zash.se> parents: 8269 diff changeset	1 -- Because we deal with pre-authed sessions and streams we can't be host-specific
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	2 module:set_global();
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	3
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	4 local it = require "prosody.util.iterators";
12977 74b9e05af71e plugins: Prefix module imports with prosody namespace Kim Alvefur <zash@zash.se> parents: 11735 diff changeset	5 local filters = require "prosody.util.filters";
74b9e05af71e plugins: Prefix module imports with prosody namespace Kim Alvefur <zash@zash.se> parents: 11735 diff changeset	6 local throttle = require "prosody.util.throttle";
74b9e05af71e plugins: Prefix module imports with prosody namespace Kim Alvefur <zash@zash.se> parents: 11735 diff changeset	7 local timer = require "prosody.util.timer";
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	8
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	9 local pton = require "prosody.util.net".pton;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	10
8269 25237002aba4 mod_limits: Handle fractional outstanding balance values (caused by e3f7b6fa46ba) Matthew Wild <mwild1@gmail.com> parents: 8256 diff changeset	11 local ceil = math.ceil;
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	12
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	13 local limits_cfg = module:get_option("limits", {});
13209 c8d949cf6b09 plugins: Switch to :get_option_period() for time range options Kim Alvefur <zash@zash.se> parents: 12977 diff changeset	14 local limits_resolution = module:get_option_period("limits_resolution", 1);
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	15
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	16 local default_bytes_per_second = 3000;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	17 local default_burst = 2;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	18
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) 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.
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	20 local function parse_rate(rate, sess_type)
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	21 local quantity, unit, exp;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	22 if rate then
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	23 quantity, unit = rate:match("^(%d+) ?([^/]+)/s$");
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	24 exp = quantity and rate_units[unit:sub(1,1):lower()];
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	25 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	26 if not exp then
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) 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);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	28 return default_bytes_per_second;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	29 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	30 return quantity*(10^exp);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	31 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	32
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	33 local function parse_burst(burst, sess_type)
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	34 if type(burst) == "string" then
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	35 burst = burst:match("^(%d+) ?s$");
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	36 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	37 local n_burst = tonumber(burst);
11550 929de6ade6b6 mod_limits: Don't emit error when no burst period is configured Matthew Wild <mwild1@gmail.com> parents: 8803 diff changeset	38 if burst and not n_burst then
10111 0f335815244f plugins: Remove tostring call from logging Kim Alvefur <zash@zash.se> parents: 10099 diff changeset	39 module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, burst, default_burst);
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	40 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	41 return n_burst or default_burst;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	42 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	43
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	44 -- Process config option into limits table:
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	45 -- limits = { c2s = { bytes_per_second = X, burst_seconds = Y } }
11554 db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	46 local limits = {
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	47 c2s = {
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	48 bytes_per_second = 10 * 1024;
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	49 burst_seconds = 2;
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	50 };
14024 f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	51 s2s = {
11554 db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	52 bytes_per_second = 30 * 1024;
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	53 burst_seconds = 2;
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	54 };
db8e41eb6eff mod_limits: Use default limits if none configured Matthew Wild <mwild1@gmail.com> parents: 11550 diff changeset	55 };
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	56
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	57 local function get_truncated_ip(conn)
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	58 local session_ip = pton(conn:ip());
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	59 if #session_ip == 128 then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	60 -- Use /64 for IPv6
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	61 session_ip = session_ip:sub(1, 8);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	62 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	63 return session_ip;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	64 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	65
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	66 local function limit_ip_sessions(session_type, max_sessions)
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	67 -- [packed_truncated_ip] = count
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	68 local current_ips = module:shared("ips."..session_type);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	69 module:hook(session_type.."-connected", function (event)
14144 e38916e22f59 mod_limits: Only count connections with associated sessions Matthew Wild <mwild1@gmail.com> parents: 14133 diff changeset	70 if not event.session then return; end
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	71 local session_ip = get_truncated_ip(event.conn);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	72 local new_count = (current_ips[session_ip] or 0) + 1;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	73 current_ips[session_ip] = new_count;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	74 if new_count > max_sessions then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	75 module:log("warn", "Rejecting stream from %s (%d of %d current connections)", event.conn:ip(), new_count, max_sessions);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	76 event.session:close({ condition = "policy-violation", text = "too many connections" });
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	77 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	78 end);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	79
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	80 module:hook(session_type.."-closed", function (event)
14144 e38916e22f59 mod_limits: Only count connections with associated sessions Matthew Wild <mwild1@gmail.com> parents: 14133 diff changeset	81 if not event.session then return; end
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	82 local session_ip = get_truncated_ip(event.conn);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	83 local new_count = (current_ips[session_ip] or 0) - 1;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	84 if new_count < 0 then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	85 module:log("warn", "IP counting failure");
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	86 new_count = 0;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	87 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	88 if new_count == 0 then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	89 current_ips[session_ip] = nil;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	90 else
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	91 current_ips[session_ip] = new_count;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	92 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	93 end);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	94 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	95
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	96 for sess_type, sess_limits in pairs(limits_cfg) do
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	97 limits[sess_type] = {
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	98 bytes_per_second = parse_rate(sess_limits.rate, sess_type);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	99 burst_seconds = parse_burst(sess_limits.burst, sess_type);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	100 };
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	101 if sess_limits.max_connections_per_ip then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	102 limit_ip_sessions(sess_type, sess_limits.max_connections_per_ip);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	103 end
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	104 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	105
14024 f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	106 if not limits.s2sin and limits.s2s then
f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	107 limits.s2sin = limits.s2s;
f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	108 end
f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	109 if not limits.s2sout then
f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	110 limits.s2sout = limits.s2s or limits.s2sin;
f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	111 end
f65302ea37b0 mod_limits: Allow configuration of general 's2s' limit, and have s2sout inherit from s2sin Matthew Wild <mwild1@gmail.com> parents: 13209 diff changeset	112
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	113 local default_filter_set = {};
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	114
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	115 function default_filter_set.bytes_in(bytes, session)
9941 a2f8d54dd445 mod_limits: Fix indentation Kim Alvefur <zash@zash.se> parents: 8803 diff changeset	116 local sess_throttle = session.throttle;
a2f8d54dd445 mod_limits: Fix indentation Kim Alvefur <zash@zash.se> parents: 8803 diff changeset	117 if sess_throttle then
10551 27b275633156 mod_limits: Remove an unused variable Kim Alvefur <zash@zash.se> parents: 10111 diff changeset	118 local ok, _, outstanding = sess_throttle:poll(#bytes, true);
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	119 if not ok then
9941 a2f8d54dd445 mod_limits: Fix indentation Kim Alvefur <zash@zash.se> parents: 8803 diff changeset	120 session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding);
8269 25237002aba4 mod_limits: Handle fractional outstanding balance values (caused by e3f7b6fa46ba) Matthew Wild <mwild1@gmail.com> parents: 8256 diff changeset	121 outstanding = ceil(outstanding);
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	122 session.conn:pause(); -- Read no more data from the connection until there is no outstanding data
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	123 local outstanding_data = bytes:sub(-outstanding);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	124 bytes = bytes:sub(1, #bytes-outstanding);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	125 timer.add_task(limits_resolution, function ()
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	126 if not session.conn then return; end
9941 a2f8d54dd445 mod_limits: Fix indentation Kim Alvefur <zash@zash.se> parents: 8803 diff changeset	127 if sess_throttle:peek(#outstanding_data) then
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	128 session.log("debug", "Resuming paused session");
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	129 session.conn:resume();
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	130 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	131 -- Handle what we can of the outstanding data
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	132 session.data(outstanding_data);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	133 end);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	134 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	135 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	136 return bytes;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	137 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	138
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	139 local type_filters = {
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	140 c2s = default_filter_set;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	141 s2sin = default_filter_set;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	142 s2sout = default_filter_set;
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	143 };
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	144
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	145 local function filter_hook(session)
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	146 local session_type = session.type:match("^[^_]+");
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	147 local filter_set, opts = type_filters[session_type], limits[session_type];
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	148 if opts then
10099 7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	149 if session.conn and session.conn.setlimit then
7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	150 session.conn:setlimit(opts.bytes_per_second);
7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	151 -- Currently no burst support
7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	152 else
7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	153 session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	154 filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
7e3196e0263e mod_limits: Use rate limiting in net.server if provided Kim Alvefur <zash@zash.se> parents: 9943 diff changeset	155 end
8256 cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	156 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	157 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	158
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	159 function module.load()
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	160 filters.add_filter_hook(filter_hook);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	161 end
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	162
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	163 function module.unload()
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	164 filters.remove_filter_hook(filter_hook);
cdffe33efae4 mod_limits: Import from prosody-modules 2c59f2f0c37d (fixes #129) Matthew Wild <mwild1@gmail.com> parents: diff changeset	165 end
9942 b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	166
11734 c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	167 function unlimited(session)
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	168 local session_type = session.type:match("^[^_]+");
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	169 if session.conn and session.conn.setlimit then
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	170 session.conn:setlimit(0);
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	171 -- Currently no burst support
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	172 else
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	173 local filter_set = type_filters[session_type];
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	174 filters.remove_filter(session, "bytes/in", filter_set.bytes_in);
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	175 session.throttle = nil;
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	176 end
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	177 end
c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	178
9942 b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	179 function module.add_host(module)
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	180 local unlimited_jids = module:get_option_inherited_set("unlimited_jids", {});
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	181
9943 46773fe2be45 mod_limits: Fix typo Kim Alvefur <zash@zash.se> parents: 9942 diff changeset	182 if not unlimited_jids:empty() then
9942 b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	183 module:hook("authentication-success", function (event)
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	184 local session = event.session;
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	185 local jid = session.username .. "@" .. session.host;
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	186 if unlimited_jids:contains(jid) then
11734 c0fc4ca74046 mod_limits: Factor out function for disabling limits allowing use from shell Kim Alvefur <zash@zash.se> parents: 11560 diff changeset	187 unlimited(session);
9942 b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	188 end
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	189 end);
11735 7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	190
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	191 module:hook("s2sout-established", function (event)
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	192 local session = event.session;
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	193 if unlimited_jids:contains(session.to_host) then
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	194 unlimited(session);
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	195 end
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	196 end);
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	197
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	198 module:hook("s2sin-established", function (event)
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	199 local session = event.session;
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	200 if session.from_host and unlimited_jids:contains(session.from_host) then
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	201 unlimited(session);
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	202 end
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	203 end);
7d29167bfcc3 mod_limits: Extend unlimited_jids to s2s sessions (for Ge0rG) Kim Alvefur <zash@zash.se> parents: 11734 diff changeset	204
9942 b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	205 end
b0d5f4ae92b7 mod_limits: Allow configuring a list of unrestricted JIDs (fixes #1323) Kim Alvefur <zash@zash.se> parents: 9941 diff changeset	206 end
14133 4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	207
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	208 module:add_item("shell-command", {
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	209 section = "limits";
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	210 section_desc = "Commands to inspect rate-limiting state";
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	211 name = "connections";
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	212 desc = "View per-IP connection limit status";
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	213 args = {
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	214 { name = "ip", type = "string" };
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	215 };
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	216 handler = function(self, ip_str) -- luacheck: ignore 212/self
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	217 local check_ip = pton(ip_str);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	218 if not check_ip then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	219 return false, "Unable to parse IP address: "..ip_str;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	220 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	221 if #check_ip == 128 then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	222 check_ip = check_ip:sub(1, 8);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	223 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	224
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	225 local found, any_limits = 0;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	226
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	227 for sess_type, sess_limits in it.sorted_pairs(limits_cfg) do
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	228 if sess_limits.max_connections_per_ip then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	229 any_limits = true;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	230 local current_ips = module:shared("ips."..sess_type);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	231 local count = current_ips[check_ip] or 0;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	232 self.session.print(("%s: %d (limit %d)"):format(sess_type, count or 0, sess_limits.max_connections_per_ip));
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	233 found = found + count;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	234 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	235 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	236
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	237 if not any_limits then
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	238 return false, "No per-IP limits are configured";
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	239 end
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	240
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	241 return true, ("Total of %d tracked connections with this IP"):format(found);
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	242 end;
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	243 });
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	244
4087c5b600b1 mod_limits: Support limiting total number of connections by IP address/block Matthew Wild <mwild1@gmail.com> parents: 14024 diff changeset	245

Mercurial > prosody-hg

annotate plugins/mod_limits.lua @ 14144:e38916e22f59 13.0