# HG changeset patch
# User Kim Alvefur <zash@zash.se>
# Date 1475590686 -7200
# Node ID 1424aa8877f015e3e09dcc46c84dafe0b6528422
# Parent  35ae59a8196d300f635a4ca8f43add25e9088875
mod_throttle_unsolicited: Add support for throttling unsolicited messages on incoming s2s connections

diff -r 35ae59a8196d -r 1424aa8877f0 mod_throttle_unsolicited/README.markdown
--- a/mod_throttle_unsolicited/README.markdown	Sat Oct 01 20:34:43 2016 +0200
+++ b/mod_throttle_unsolicited/README.markdown	Tue Oct 04 16:18:06 2016 +0200
@@ -2,14 +2,15 @@
 depends:
 - 'mod\_track\_muc\_joins'
 summary: Limit rate of outgoing unsolicited messages
-...
+---
 
 Introduction
 ============
 
-This module limits the rate of outgoing unsolicited messages. A message
-counts as "unsolicited" if the receiving user hasn't added the sending
-user to their roster.
+This module limits the rate of outgoing unsolicited messages from local
+clients. Optionally, unsolicited messages coming in from remote servers
+may be limited per s2s conneciton. A message counts as "unsolicited" if
+the receiving user hasn't added the sending user to their roster.
 
 The module depends on [mod\_track\_muc\_joins] in order to allow sent
 messages to joined MUC rooms.
@@ -17,6 +18,14 @@
 Configuration
 =============
 
+To set a limit on messages from local sessions:
+
 ``` {.lua}
 unsolicited_messages_per_minute = 10
 ```
+
+To enable limits on unsolicited messages from s2s connections:
+
+``` {.lua}
+unsolicited_s2s_messages_per_minute = 100
+```
diff -r 35ae59a8196d -r 1424aa8877f0 mod_throttle_unsolicited/mod_throttle_unsolicited.lua
--- a/mod_throttle_unsolicited/mod_throttle_unsolicited.lua	Sat Oct 01 20:34:43 2016 +0200
+++ b/mod_throttle_unsolicited/mod_throttle_unsolicited.lua	Tue Oct 04 16:18:06 2016 +0200
@@ -6,6 +6,7 @@
 local gettime = require "socket".gettime;
 
 local max = module:get_option_number("unsolicited_messages_per_minute", 10);
+local s2s_max = module:get_option_number("unsolicited_s2s_messages_per_minute");
 local multiplier = module:get_option_number("throttle_unsolicited_burst", 1);
 
 function check_subscribed(event)
@@ -51,4 +52,42 @@
 module:hook("pre-message/bare", check_subscribed, 200);
 module:hook("pre-message/full", check_subscribed, 200);
 
+local full_sessions = prosody.full_sessions;
+
+-- Rooms and throttle creation will differ for s2s
+function check_subscribed_s2s(event)
+	local stanza, origin = event.stanza, event.origin;
+	local log = origin.log or module._log;
+
+	if origin.type ~= "s2sin" then return end
+
+	local to_orig = stanza.attr.to;
+	local from_orig = stanza.attr.from;
+	local from_bare = jid_bare(from_orig);
+
+	local target = full_sessions[to_orig];
+	if target then
+		local rooms = target.rooms_joined;
+		if rooms and rooms[from_bare] then
+			log("debug", "Message to joined room, no limit");
+			return
+		end
+	end
+
+	-- Retrieve or create throttle object
+	local lim = origin.throttle_unsolicited;
+	if not lim then
+		log("debug", "New s2s throttle");
+		lim = throttle.create(s2s_max * multiplier, 60 * multiplier);
+		origin.throttle_unsolicited = lim;
+	end
+
+	return check_subscribed(event);
+end
+
+if s2s_max then
+	module:hook("message/bare", check_subscribed_s2s, 200);
+	module:hook("message/full", check_subscribed_s2s, 200);
+end
+
 module:depends("track_muc_joins");