-- Prosody IM

-- Copyright (C) 2008-2010 Matthew Wild

-- Copyright (C) 2008-2010 Waqas Hussain

--

-- This project is MIT/X11 licensed. Please see the

-- COPYING file in the source package for more information.

--

module:set_global();

local prosody = prosody;

local hosts = prosody.hosts;

local core_process_stanza = prosody.core_process_stanza;

local tostring, type = tostring, type;

local t_insert = table.insert;

local xpcall, traceback = xpcall, debug.traceback;

local add_task = require "util.timer".add_task;

local st = require "util.stanza";

local initialize_filters = require "util.filters".initialize;

local nameprep = require "util.encodings".stringprep.nameprep;

local new_xmpp_stream = require "util.xmppstream".new;

local s2s_new_incoming = require "core.s2smanager".new_incoming;

local s2s_new_outgoing = require "core.s2smanager".new_outgoing;

local s2s_destroy_session = require "core.s2smanager".destroy_session;

local uuid_gen = require "util.uuid".generate;

local fire_global_event = prosody.events.fire_event;

local s2sout = module:require("s2sout");

local connect_timeout = module:get_option_number("s2s_timeout", 90);

local stream_close_timeout = module:get_option_number("s2s_close_timeout", 5);

local opt_keepalives = module:get_option_boolean("s2s_tcp_keepalives", module:get_option_boolean("tcp_keepalives", true));

local secure_auth = module:get_option_boolean("s2s_secure_auth", false); -- One day...

local secure_domains, insecure_domains =

	module:get_option_set("s2s_secure_domains", {})._items, module:get_option_set("s2s_insecure_domains", {})._items;

local require_encryption = module:get_option_boolean("s2s_require_encryption", false);

local measure_connections = module:measure("connections", "amount");

local sessions = module:shared("sessions");

local log = module._log;

module:hook("stats-update", function ()

	local count = 0;

	for _ in pairs(sessions) do

		count = count + 1;

	end

	measure_connections(count);

end);

--- Handle stanzas to remote domains

local bouncy_stanzas = { message = true, presence = true, iq = true };

local function bounce_sendq(session, reason)

	local sendq = session.sendq;

	if not sendq then return; end

	session.log("info", "Sending error replies for "..#sendq.." queued stanzas because of failed outgoing connection to "..tostring(session.to_host));

	local dummy = {

		type = "s2sin";

		send = function(s)

			(session.log or log)("error", "Replying to to an s2s error reply, please report this! Traceback: %s", traceback());

		end;

		dummy = true;

	};

	for i, data in ipairs(sendq) do

		local reply = data[2];

		if reply and not(reply.attr.xmlns) and bouncy_stanzas[reply.name] then

			reply.attr.type = "error";

			reply:tag("error", {type = "cancel"})

				:tag("remote-server-not-found", {xmlns = "urn:ietf:params:xml:ns:xmpp-stanzas"}):up();

			if reason then

				reply:tag("text", {xmlns = "urn:ietf:params:xml:ns:xmpp-stanzas"})

					:text("Server-to-server connection failed: "..reason):up();

			end

			core_process_stanza(dummy, reply);

		end

		sendq[i] = nil;

	end

	session.sendq = nil;

end

-- Handles stanzas to existing s2s sessions

function route_to_existing_session(event)

	local from_host, to_host, stanza = event.from_host, event.to_host, event.stanza;

	if not hosts[from_host] then

		log("warn", "Attempt to send stanza from %s - a host we don't serve", from_host);

		return false;

	end

	if hosts[to_host] then

		log("warn", "Attempt to route stanza to a remote %s - a host we do serve?!", from_host);

		return false;

	end

	local host = hosts[from_host].s2sout[to_host];

	if host then

		-- We have a connection to this host already

		if host.type == "s2sout_unauthed" and (stanza.name ~= "db:verify" or not host.dialback_key) then

			(host.log or log)("debug", "trying to send over unauthed s2sout to "..to_host);

			-- Queue stanza until we are able to send it

			if host.sendq then t_insert(host.sendq, {tostring(stanza), stanza.attr.type ~= "error" and stanza.attr.type ~= "result" and st.reply(stanza)});

			else host.sendq = { {tostring(stanza), stanza.attr.type ~= "error" and stanza.attr.type ~= "result" and st.reply(stanza)} }; end

			host.log("debug", "stanza [%s] queued ", stanza.name);

			return true;

		elseif host.type == "local" or host.type == "component" then

			log("error", "Trying to send a stanza to ourselves??")

			log("error", "Traceback: %s", traceback());

			log("error", "Stanza: %s", tostring(stanza));

			return false;

		else

			-- FIXME

			if host.from_host ~= from_host then

				log("error", "WARNING! This might, possibly, be a bug, but it might not...");

				log("error", "We are going to send from %s instead of %s", tostring(host.from_host), tostring(from_host));

			end

			if host.sends2s(stanza) then

				return true;

			end

		end

	end

end

-- Create a new outgoing session for a stanza

function route_to_new_session(event)

	local from_host, to_host, stanza = event.from_host, event.to_host, event.stanza;

	log("debug", "opening a new outgoing connection for this stanza");

	local host_session = s2s_new_outgoing(from_host, to_host);

	-- Store in buffer

	host_session.bounce_sendq = bounce_sendq;

	host_session.sendq = { {tostring(stanza), stanza.attr.type ~= "error" and stanza.attr.type ~= "result" and st.reply(stanza)} };

	log("debug", "stanza [%s] queued until connection complete", tostring(stanza.name));

	s2sout.initiate_connection(host_session);

	if (not host_session.connecting) and (not host_session.conn) then

		log("warn", "Connection to %s failed already, destroying session...", to_host);

		s2s_destroy_session(host_session, "Connection failed");

		return false;

	end

	return true;

end

local function keepalive(event)

	return event.session.sends2s(' ');

end

module:hook("s2s-read-timeout", keepalive, -1);

function module.add_host(module)

	if module:get_option_boolean("disallow_s2s", false) then

		module:log("warn", "The 'disallow_s2s' config option is deprecated, please see http://prosody.im/doc/s2s#disabling");

		return nil, "This host has disallow_s2s set";

	end

	module:hook("route/remote", route_to_existing_session, -1);

	module:hook("route/remote", route_to_new_session, -10);

	module:hook("s2s-authenticated", make_authenticated, -1);

	module:hook("s2s-read-timeout", keepalive, -1);

	module:hook_stanza("http://etherx.jabber.org/streams", "features", function (session, stanza)

		if session.type == "s2sout" then

			-- Stream is authenticated and we are seem to be done with feature negotiation,

			-- so the stream is ready for stanzas.  RFC 6120 Section 4.3

			mark_connected(session);

			return true;

		elseif not session.dialback_verifying then

			session.log("warn", "No SASL EXTERNAL offer and Dialback doesn't seem to be enabled, giving up");

			session:close();

			return false;

		end

	end, -1);

end

-- Stream is authorised, and ready for normal stanzas

function mark_connected(session)

	local sendq = session.sendq;

	local from, to = session.from_host, session.to_host;

	session.log("info", "%s s2s connection %s->%s complete", session.direction:gsub("^.", string.upper), from, to);

	local event_data = { session = session };

	if session.type == "s2sout" then

		fire_global_event("s2sout-established", event_data);

		hosts[from].events.fire_event("s2sout-established", event_data);

	else

		local host_session = hosts[to];

		session.send = function(stanza)

			return host_session.events.fire_event("route/remote", { from_host = to, to_host = from, stanza = stanza });

		end;

		fire_global_event("s2sin-established", event_data);

		hosts[to].events.fire_event("s2sin-established", event_data);

	end

	if session.direction == "outgoing" then

		if sendq then

			session.log("debug", "sending %d queued stanzas across new outgoing connection to %s", #sendq, session.to_host);

			local send = session.sends2s;

			for i, data in ipairs(sendq) do

				send(data[1]);

				sendq[i] = nil;

			end

			session.sendq = nil;

		end

		if session.resolver then

			session.resolver._resolver:closeall()

		end

		session.resolver = nil;

		session.ip_hosts = nil;

		session.srv_hosts = nil;

	end

end

function make_authenticated(event)

	local session, host = event.session, event.host;

	if not session.secure then

		if require_encryption or (secure_auth and not(insecure_domains[host])) or secure_domains[host] then

			session:close({

				condition = "policy-violation",

				text = "Encrypted server-to-server communication is required but was not "

				       ..((session.direction == "outgoing" and "offered") or "used")

			});

		end

	end

	if hosts[host] then

		session:close({ condition = "undefined-condition", text = "Attempt to authenticate as a host we serve" });

	end

	if session.type == "s2sout_unauthed" then

		session.type = "s2sout";

	elseif session.type == "s2sin_unauthed" then

		session.type = "s2sin";

		if host then

			if not session.hosts[host] then session.hosts[host] = {}; end

			session.hosts[host].authed = true;

		end

	elseif session.type == "s2sin" and host then

		if not session.hosts[host] then session.hosts[host] = {}; end

		session.hosts[host].authed = true;

	else

		return false;

	end

	session.log("debug", "connection %s->%s is now authenticated for %s", session.from_host, session.to_host, host);

	if (session.type == "s2sout" and session.external_auth ~= "succeeded") or session.type == "s2sin" then

		-- Stream either used dialback for authentication or is an incoming stream.

		mark_connected(session);

	end

	return true;

end

--- Helper to check that a session peer's certificate is valid

function check_cert_status(session)

	local host = session.direction == "outgoing" and session.to_host or session.from_host

	local conn = session.conn:socket()

	local cert

	if conn.getpeercertificate then

		cert = conn:getpeercertificate()

	end

	return module:fire_event("s2s-check-certificate", { host = host, session = session, cert = cert });

end

--- XMPP stream event handlers

local stream_callbacks = { default_ns = "jabber:server", handlestanza =  core_process_stanza };

local xmlns_xmpp_streams = "urn:ietf:params:xml:ns:xmpp-streams";

function stream_callbacks.streamopened(session, attr)

	session.version = tonumber(attr.version) or 0;

	-- TODO: Rename session.secure to session.encrypted

	if session.secure == false then

		session.secure = true;

		session.encrypted = true;

		local sock = session.conn:socket();

		if sock.info then

			local info = sock:info();

			(session.log or log)("info", "Stream encrypted (%s with %s)", info.protocol, info.cipher);

			session.compressed = info.compression;

		else

			(session.log or log)("info", "Stream encrypted");

			session.compressed = sock.compression and sock:compression(); --COMPAT mw/luasec-hg

		end

	end

	if session.direction == "incoming" then

		-- Send a reply stream header

		-- Validate to/from

		local to, from = nameprep(attr.to), nameprep(attr.from);

		if not to and attr.to then -- COMPAT: Some servers do not reliably set 'to' (especially on stream restarts)

			session:close({ condition = "improper-addressing", text = "Invalid 'to' address" });

			return;

		end

		if not from and attr.from then -- COMPAT: Some servers do not reliably set 'from' (especially on stream restarts)

			session:close({ condition = "improper-addressing", text = "Invalid 'from' address" });

			return;

		end

		-- Set session.[from/to]_host if they have not been set already and if

		-- this session isn't already authenticated

		if session.type == "s2sin_unauthed" and from and not session.from_host then

			session.from_host = from;

		elseif from ~= session.from_host then

			session:close({ condition = "improper-addressing", text = "New stream 'from' attribute does not match original" });

			return;

		end

		if session.type == "s2sin_unauthed" and to and not session.to_host then

			session.to_host = to;

		elseif to ~= session.to_host then

			session:close({ condition = "improper-addressing", text = "New stream 'to' attribute does not match original" });

			return;

		end

		-- For convenience we'll put the sanitised values into these variables

		to, from = session.to_host, session.from_host;

		session.streamid = uuid_gen();

		(session.log or log)("debug", "Incoming s2s received %s", st.stanza("stream:stream", attr):top_tag());

		if to then

			if not hosts[to] then

				-- Attempting to connect to a host we don't serve

				session:close({

					condition = "host-unknown";

					text = "This host does not serve "..to

				});

				return;

			elseif not hosts[to].modules.s2s then

				-- Attempting to connect to a host that disallows s2s

				session:close({

					condition = "policy-violation";

					text = "Server-to-server communication is disabled for this host";

				});

				return;

			end

		end

		if hosts[from] then

			session:close({ condition = "undefined-condition", text = "Attempt to connect from a host we serve" });

			return;

		end

		if session.secure and not session.cert_chain_status then

			if check_cert_status(session) == false then

				return;

			end

		end

		session:open_stream(session.to_host, session.from_host)

		session.notopen = nil;

		if session.version >= 1.0 then

			local features = st.stanza("stream:features");

			if to then

				hosts[to].events.fire_event("s2s-stream-features", { origin = session, features = features });

			else

				(session.log or log)("warn", "No 'to' on stream header from %s means we can't offer any features", from or session.ip or "unknown host");