--- a/mod_s2s_keepalive/mod_s2s_keepalive.lua	Sat Dec 21 17:47:52 2019 +0100
+++ b/mod_s2s_keepalive/mod_s2s_keepalive.lua	Sat Dec 21 17:52:32 2019 +0100
@@ -1,7 +1,9 @@
 local st = require "util.stanza";
+local watchdog = require "util.watchdog";
 
 local keepalive_servers = module:get_option_set("keepalive_servers");
 local keepalive_interval = module:get_option_number("keepalive_interval", 60);
+local keepalive_timeout = module:get_option_number("keepalive_timeout", 593);
 
 local host = module.host;
 local s2sout = prosody.hosts[host].s2sout;
@@ -15,7 +17,6 @@
 			session.sends2s(st.iq({ to = remote_domain, type = "get", from = host, id = "keepalive" })
 				:tag("ping", { xmlns = "urn:xmpp:ping" })
 			);
-			-- Note: We don't actually check if this comes back.
 		end
 	end
 
@@ -38,4 +39,43 @@
 	return keepalive_interval;
 end
 
+module:hook("s2sin-established", function (event)
+	local session = event.session;
+	if session.watchdog_keepalive then return end -- in case mod_bidi fires this twice
+	session.watchdog_keepalive = watchdog.new(keepalive_timeout, function ()
+		session.log("info", "Keepalive ping timed out, closing connection");
+		session:close("connection-timeout");
+	end);
+end);
+
+module:hook("s2sout-established", function (event)
+	local session = event.session;
+	if session.watchdog_keepalive then return end -- in case mod_bidi fires this twice
+	session.watchdog_keepalive = watchdog.new(keepalive_timeout, function ()
+		session.log("info", "Keepalive ping timed out, closing connection");
+		session:close("connection-timeout");
+	end);
+end);
+
+module:hook("iq-result/host/keepalive", function (event)
+	local origin = event.origin;
+	if origin.watchdog_keepalive then
+		origin.watchdog_keepalive:reset();
+	end
+	if s2sout[origin.from_host] then
+		s2sout[origin.from_host]:reset();
+	end
+end);
+
+module:hook("iq-error/host/keepalive", function (event)
+	local origin = event.origin;
+	if origin.dummy then return end -- Probably a sendq bounce
+
+	if origin.type == "s2sin" or origin.type == "s2sout" then
+		-- An error from the remote means connectivity is ok,
+		-- so treat it the same as a result
+		return module:fire_event("iq-result/host/keepalive");
+	end
+end);
+
 module:add_timer(keepalive_interval, send_pings);