mod_onions.lua

module:set_global();

local wrapclient = require "net.server".wrapclient;
local s2s_new_outgoing = require "core.s2smanager".new_outgoing;
local initialize_filters = require "util.filters".initialize;
local bit = require "bit32";
local st = require "util.stanza";
local portmanager = require "core.portmanager";

-- Configuration, todo: actual configuration
local proxy_ip = "127.0.0.1";
local proxy_port = "9050";

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

-- The socks5listener handles connection while still connecting to the proxy,
-- then it hands them over to the normal listener (in mod_s2s)
local socks5listener = { default_port = 9050, default_mode = "*a", default_interface = "*" };

local function socks5_connect_sent(conn, data)
	
	local session = sessions[conn];

	if #data < 5 then
		module:log("debug", "Did not receive a full reply, waiting.");
		session.socks5_buffer = data
		return
	end

	request_status = string.byte(data, 2);

	if not request_status == 0x00 then
		module:log("debug", "Failed to connect to the SOCKS5 proxy. :(");
		return;
	end

	module:log("debug", "Succesfully connected over SOCKS5");
	
	local response = string.byte(data, 4);

	module:log("debug", "Response (2): "..response);

	-- see if we should connect somewhere else
	if response == 0x03 then

		-- this means the server tells us to connect on a hostname
		local len = string.byte(data, 5);

		if #data < 6+len+2 then
			-- let's try again when we have enough
			module:log("debug", "Did not receive a full hostname reply, waiting.");
			session.socks5_buffer = data
			return
		end

		local hostname = string.byte(data, 6, 6+len);
		local port = bit.band(string.byte(data, 6+len+1), bit.lshift(string.byte(data, 6+len+2), 8));
		module:log("debug", "Should connect to: "..hostname..":"..port);

		-- TODO: connect on the other host instead
	elseif response == 0x01 then

		if #data < 10 then
			-- let's try again when we have enough
			module:log("debug", "Did not receive a full IP address reply, waiting.");
			session.socks5_buffer = data
			return
		end

		-- this means the server tells us to connect on an IPv4 address
		local ip1 = string.byte(data, 5);
		local ip2 = string.byte(data, 6);
		local ip3 = string.byte(data, 7);
		local ip4 = string.byte(data, 8);
		local port = bit.band(string.byte(data, 9), bit.lshift(string.byte(data, 10), 8));
		module:log("debug", "Should connect to: "..ip1.."."..ip2.."."..ip3.."."..ip4..":"..port);

		if not (ip1 == 0 and ip2 == 0 and ip3 == 0 and ip4 == 0 and port == 0) then
			module:log("debug", "The SOCKS5 proxy tells us to connect to a different IP, don't know how. :(");
			return
		end

		-- Now the real s2s listener can take over the connection.
		local listener = portmanager.get_service("s2s").listener;

		module:log("debug", "SOCKS5 done, handing over listening to "..tostring(listener));

		conn.setlistener(conn, listener);

		local w, log = conn.send, session.log;

		local filter = initialize_filters(session);

		session.sends2s = function (t)
			log("debug", "sending (socks5): %s", (t.top_tag and t:top_tag()) or t:match("^[^>]*>?"));
			if t.name then
				t = filter("stanzas/out", t);
			end
			if t then
				t = filter("bytes/out", tostring(t));
				if t then
					return w(conn, tostring(t));
				end
			end
		end

		listener.register_outgoing(conn, session);

		session.sends2s(st.stanza("stream:stream", {
			xmlns='jabber:server', ["xmlns:db"]='jabber:server:dialback',
			["xmlns:stream"]='http://etherx.jabber.org/streams',
			from=session.from_host, to=session.to_host, version='1.0', ["xml:lang"]='en'}):top_tag());

		session.socks5_handler = nil;
		session.socks5_buffer = nil;

		listener.onconnect(conn);
	end
end

local function socks5_handshake_sent(conn, data)

	local session = sessions[conn];

	if #data < 2 then
		module:log("debug", "Did not receive a full reply, waiting.");
		session.socks5_buffer = data
		return
	end

	-- version, method
	local request_status = string.byte(data, 2);

	module:log("debug", "SOCKS version: "..string.byte(data, 1));
	module:log("debug", "Response: "..request_status);

	if not request_status == 0x00 then
		module:log("debug", "Failed to connect to the SOCKS5 proxy. :( It seems to require authentication.");
		return;
	end

	module:log("debug", "Sending connect message: 05 01 00 03 " .. #session.socks5_to .. " " .. session.socks5_port);

	-- version 5, connect, (reserved), type: domainname, (length, hostname), port
	query = "\05\01\00\03"..string.char(#session.socks5_to)..session.socks5_to..string.char(bit.rshift(session.socks5_port, 8))..string.char(bit.band(session.socks5_port, 0xff));

	conn:send(query);

	session.socks5_handler = socks5_connect_sent;
end

function socks5listener.onconnect(conn)
	module:log("debug", "Connected to SOCKS5 proxy.");

	module:log("debug", "Sending SOCKS5 handshake.");

	-- Socks version 5, 1 method, no auth
	local query = '\05\01\00';

	conn:send(query);
	sessions[conn].socks5_handler = socks5_handshake_sent;
end

function socks5listener.register_outgoing(conn, session)
	session.direction = "outgoing";
	sessions[conn] = session;
	sessions[conn].socks5status = "new";
end

function socks5listener.ondisconnect(conn, err)
	module:log("debug", "Closing connection to SOCKS5 proxy.")
	local session = sessions[conn];
	sessions[conn]  = nil;
end

function socks5listener.onincoming(conn, data)
	module:log("debug", "Received something from SOCKS5 proxy, "..sessions[conn].socks5status);

	local session = sessions[conn];

	if session.socks5_buffer then
		data = session.socks5_buffer .. data
	end

	if session.socks5_handler then
		session.socks5_handler(conn, data);
	end
end

local function connect_socks5(host_session, connect_host, connect_port)
	
	local conn, handler = socket.tcp();

	module:log("debug", "Connecting to " .. connect_host .. ":" .. connect_port);

	-- this is not necessarily the same as .to_host (it can be that this is a SRV record)
	host_session.socks5_to = connect_host;
	host_session.socks5_port = connect_port;

	conn:settimeout(0);

	local success, err = conn:connect(proxy_ip, proxy_port);

	conn = wrapclient(conn, connect_host, connect_port, socks5listener, "*a");

	socks5listener.register_outgoing(conn, host_session);

	host_session.conn = conn;
end

-- There's two signals that are handled: pre-try-connect (which I added) and route/remote.
-- route/remote gets called when routing anything to a remote server, so if that already matches *.onion, intercept it
-- if the server has a hostname, but a SRV record with a *.onion address, pre-try-connect will intercept it.

module:hook("pre-try-connect", function (event)
	module:log("debug", "Wrapping connection attempt to " .. event.connect_host .. ":" .. event.connect_port);

	local connect_host = event.connect_host;

	if connect_host:find(".onion(.?)$") then

		if string.sub(connect_host, -1) == "." then
			connect_host = string.sub(connect_host, 1, -2);
		end

		module:log("debug", "It's an onion, intercepting it.");

		connect_socks5(event.host_session, connect_host, event.connect_port);

		return true;
	end

	return false;
end, 100);

module:hook("route/remote", function(event)

	if not event.to_host:find(".onion(.?)$") then
		module:log("debug", event.to_host .. " is not an onion. Not doing anything.");
		return;
	end

	module:log("debug", "Onion routing something to ".. event.to_host);

	if hosts[event.from_host].s2sout[event.to_host] then
		module:log("debug", "Host session exists, no need to set up SOCKS5.")
		return;
	end

	local host_session = s2s_new_outgoing(event.from_host, event.to_host);

	hosts[event.from_host].s2sout[event.to_host] = host_session;

	connect_socks5(host_session, event.to_host, 5269);

	return false;
end, 250);