Diff

--- a/plugins/mod_limits.lua	Thu Nov 05 22:27:17 2020 +0100
+++ b/plugins/mod_limits.lua	Thu Nov 05 22:31:25 2020 +0100
@@ -32,7 +32,7 @@
 	end
 	local n_burst = tonumber(burst);
 	if not n_burst then
-		module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, tostring(burst), default_burst);
+		module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, burst, default_burst);
 	end
 	return n_burst or default_burst;
 end
@@ -51,18 +51,18 @@
 local default_filter_set = {};
 
 function default_filter_set.bytes_in(bytes, session)
-  local sess_throttle = session.throttle;
-  if sess_throttle then
-    local ok, balance, outstanding = sess_throttle:poll(#bytes, true);
+	local sess_throttle = session.throttle;
+	if sess_throttle then
+		local ok, _, outstanding = sess_throttle:poll(#bytes, true);
 		if not ok then
-      session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding);
+			session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding);
 			outstanding = ceil(outstanding);
 			session.conn:pause(); -- Read no more data from the connection until there is no outstanding data
 			local outstanding_data = bytes:sub(-outstanding);
 			bytes = bytes:sub(1, #bytes-outstanding);
 			timer.add_task(limits_resolution, function ()
 				if not session.conn then return; end
-        if sess_throttle:peek(#outstanding_data) then
+				if sess_throttle:peek(#outstanding_data) then
 					session.log("debug", "Resuming paused session");
 					session.conn:resume();
 				end
@@ -84,8 +84,13 @@
 	local session_type = session.type:match("^[^_]+");
 	local filter_set, opts = type_filters[session_type], limits[session_type];
 	if opts then
-		session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
-		filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
+		if session.conn and session.conn.setlimit then
+			session.conn:setlimit(opts.bytes_per_second);
+			-- Currently no burst support
+		else
+			session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
+			filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
+		end
 	end
 end
 
@@ -96,3 +101,25 @@
 function module.unload()
 	filters.remove_filter_hook(filter_hook);
 end
+
+function module.add_host(module)
+	local unlimited_jids = module:get_option_inherited_set("unlimited_jids", {});
+
+	if not unlimited_jids:empty() then
+		module:hook("authentication-success", function (event)
+			local session = event.session;
+			local session_type = session.type:match("^[^_]+");
+			local jid = session.username .. "@" .. session.host;
+			if unlimited_jids:contains(jid) then
+				if session.conn and session.conn.setlimit then
+					session.conn:setlimit(0);
+					-- Currently no burst support
+				else
+					local filter_set = type_filters[session_type];
+					filters.remove_filter(session, "bytes/in", filter_set.bytes_in);
+					session.throttle = nil;
+				end
+			end
+		end);
+	end
+end