-- mod_s2s_auth_dane

-- Copyright (C) 2013-2014 Kim Alvefur

--

-- This file is MIT/X11 licensed.

--

-- Could be done much cleaner if mod_s2s was using util.async

module:set_global();

local dns_lookup = require"net.adns".lookup;

local hashes = require"util.hashes";

local base64 = require"util.encodings".base64;

local s2sout = module:depends"s2s".route_to_new_session.s2sout;

local pat = "%-%-%-%-%-BEGIN ([A-Z ]+)%-%-%-%-%-\r?\n"..

"([0-9A-Za-z=+/\r\n]*)\r?\n%-%-%-%-%-END %1%-%-%-%-%-";

local function pem2der(pem)

	local typ, data = pem:match(pat);

	if typ and data then

		return base64.decode(data), typ;

	end

end

-- TODO Things to test/handle:

-- Negative or bogus answers

-- No SRV records

-- No encryption offered

-- This function is called when a new SRV target has been picked

-- the original function does A/AAAA resolution before continuing

local _try_connect = s2sout.try_connect;

function s2sout.try_connect(host_session, connect_host, connect_port, err)

	local srv_hosts = host_session.srv_hosts;

	local srv_choice = host_session.srv_choice;

	if srv_hosts and srv_hosts.answer.secure and not srv_hosts[srv_choice].dane then

		srv_hosts[srv_choice].dane = dns_lookup(function(answer)

			if answer and ( #answer > 0 or answer.bogus ) then

				srv_hosts[srv_choice].dane = answer;

				for i, tlsa in ipairs(answer) do

					module:log("debug", "TLSA %s", tostring(tlsa));

				end

			else

				srv_hosts[srv_choice].dane = false;

			end

			-- "blocking" until TLSA reply, but no race condition

			return _try_connect(host_session, connect_host, connect_port, err);

		end, ("_%d._tcp.%s"):format(connect_port, connect_host), "TLSA");

		return true

	end

	return _try_connect(host_session, connect_host, connect_port, err);

end

module:hook("s2s-check-certificate", function(event)

	local session, cert = event.session, event.cert;

	local srv_hosts = session.srv_hosts;

	local srv_choice = session.srv_choice;

	local choosen = srv_hosts and srv_hosts[srv_choice] or session;

	if choosen.dane then

		local use, select, match, tlsa, certdata, match_found;

		for i, rr in ipairs(choosen.dane) do

			tlsa = rr.tlsa;

			module:log("debug", "TLSA %s", tostring(tlsa));

			use, select, match, certdata = tlsa.use, tlsa.select, tlsa.match;

			-- PKIX-EE or DANE-EE

			if use == 1 or use == 3 then

				if select == 0 then

					certdata = pem2der(cert:pem());

				elseif select == 1 and cert.pubkey then

				else

					module:log("warn", "DANE selector %d is unsupported", select);

				end

				if match == 1 then

					certdata = hashes.sha256(certdata);

				elseif match == 2 then

					certdata = hashes.sha512(certdata);

				elseif match ~= 0 then

					module:log("warn", "DANE match rule %d is unsupported", match);

					certdata = nil;

				end

				-- Should we check if the cert subject matches?

				if certdata and certdata == tlsa.data then

					(session.log or module._log)("info", "DANE validation successful");

					session.cert_identity_status = "valid";

					if use == 3 then -- DANE-EE, chain status equals DNSSEC chain status

						session.cert_chain_status = "valid";

						-- for usage 1, PKIX-EE, the chain has to be valid already

					end

					match_found = true;

					break;

				end

			else

				module:log("warn", "DANE %s is unsupported", tlsa:getUsage() or ("usage "..tostring(use)));

			end

		end

		if not match_found then

			-- No TLSA matched or response was bogus

			(session.log or module._log)("warn", "DANE validation failed");

			session.cert_identity_status = "invalid";

			session.cert_chain_status = "invalid";

		end

	end

end);

function module.add_host(module)

	module:hook("s2s-authenticated", function(event)

		local session = event.session;

		local srv_hosts = session.srv_hosts;

		local srv_choice = session.srv_choice;

		if (session.dane or srv_hosts and srv_hosts[srv_choice].dane) and not session.secure then

			-- TLSA record but no TLS, not ok.

			-- TODO Optional?

			session:close({

				condition = "policy-violation",

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

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

			});

			return false;

		end

	end);

	-- DANE for s2sin

	-- Looks for TLSA at the same QNAME as the SRV record

	module:hook("stanza/urn:ietf:params:xml:ns:xmpp-tls:starttls", function(event)

		local origin = event.origin;

		if not origin.from_host then return end

		origin.dane = dns_lookup(function(answer)

			if answer and ( #answer > 0 or answer.bogus ) then

				origin.dane = answer;

				for i, tlsa in ipairs(answer) do

					module:log("debug", "TLSA %s", tostring(tlsa));

				end

			else

				origin.dane = false;

			end

			-- "blocking" until TLSA reply, but no race condition

		end, ("_xmpp-server._tcp.%s"):format(origin.from_host), "TLSA");

	end, 1);

end

function module.unload()

	-- Restore the original try_connect function

	s2sout.try_connect = _try_connect;

end