mod_net_proxy.lua

module:set_global();

local portmanager = require "core.portmanager";
local set = require "util.set";
local ip = require "util.ip";

local listener = {};
local sessions = {};
local mappings = {};

local protocol_max_header_length = 256;
local protocol_handlers = {
	PROXYv1 = { signature = "\x50\x52\x4F\x58\x59" },
	PROXYv2 = { signature = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A" }
};
local protocol_handler_responses = { success = 0, postpone = 1, failure = 2};

local proxy_data_mt = {}; proxy_data_mt.__index = proxy_data_mt;
function proxy_data_mt:describe()
	return string.format("AFI=%s Source=%s:%d Destination=%s:%d",
		self.address_family, self.source_addr, self.source_port, self.dest_addr, self.dest_port);
end

protocol_handlers["PROXYv1"].callback = function(conn, session)
	local valid_address_families = set.new({"TCP4", "TCP6", "UNKNOWN"});

	-- Abort with failure if CRLF (PROXYv1 header terminator) does not exist within buffer
	if session.buffer:find("\r\n") == nil then
		return protocol_handler_responses.postpone, nil;
	end

	-- Declare header pattern and match current buffer against pattern
	local header_pattern = "^PROXY (%S+) (%S+) (%S+) (%d+) (%d+)\r\n";
	local address_family, source_addr, dest_addr, source_port, dest_port = session.buffer:match(header_pattern);
	source_port, dest_port = tonumber(source_port), tonumber(dest_port);

	-- Ensure that header was successfully parsed and contains a valid address family
	if address_family == nil or source_addr == nil or dest_addr == nil or source_port == nil or dest_port == nil then
		module:log("warn", "Received unparseable PROXYv1 header from %s", conn:ip());
		return protocol_handler_responses.failure, nil;
	end
	if not valid_address_families:contains(address_family) then
		module:log("warn", "Received invalid PROXYv1 address family from %s: %s", conn:ip(), address_family);
		return protocol_handler_responses.failure, nil;
	end
	if address_family == "UNKNOWN" then
		module:log("warn", "Received forbidden PROXYv1 address family 'UNKNOWN' from %s", conn:ip());
		return protocol_handler_responses.failure, nil;
	end

	-- Ensure that received source and destination ports are within 1 and 65535 (0xFFFF)
	if source_port <= 0 or source_port >= 0xFFFF then
		module:log("warn", "Received invalid PROXYv1 source port from %s: %d", conn:ip(), source_port);
		return protocol_handler_responses.failure, nil;
	end
	if dest_port <= 0 or dest_port >= 0xFFFF then
		module:log("warn", "Received invalid PROXYv1 destination port from %s: %d", conn:ip(), dest_port);
		return protocol_handler_responses.failure, nil;
	end

	-- Ensure that received source and destination address can be parsed
	local _, err = ip.new_ip(source_addr);
	if err ~= nil then
		module:log("warn", "Received unparseable PROXYv1 source address from %s: %s", conn:ip(), source_addr);
	end
	_, err = ip.new_ip(dest_addr);
	if err ~= nil then
		module:log("warn", "Received unparseable PROXYv1 destination address from %s: %s", conn:ip(), dest_addr);
		return protocol_handler_responses.failure, nil;
	end

	-- Strip parsed header from session buffer and build proxy data
	session.buffer = session.buffer:gsub(header_pattern, "");

	local proxy_data = {
		["version"] = 1,
		["address_family"] = address_family,
		["source_addr"] = source_addr, ["source_port"] = source_port,
		["dest_addr"] = dest_addr, ["dest_port"] = dest_port
	};
	setmetatable(proxy_data, proxy_data_mt);

	-- Return successful response with gathered proxy data
	return protocol_handler_responses.success, proxy_data;
end

protocol_handlers["PROXYv2"].callback = function(conn, session)
	module:log("error", "Received unsupported PROXYv2 header from %s", conn:ip());
	return protocol_handler_responses.failure, nil;
end

local function wrap_proxy_connection(conn, session, proxy_data)
	-- Override and add functions of 'conn' object
	conn.proxyip, conn.proxyport = conn.ip, conn.port;
	conn.ip = function() return proxy_data.source_addr; end
	conn.port = function() return proxy_data.source_port; end
	conn.clientport = conn.port;

	-- Attempt to find service by processing port<>service mappings
	local mapping = mappings[conn:serverport()];
	if mapping == nil then
		conn:close();
		module:log("error", "Connection %s-%s terminated: Could not find mapping for port %d", conn:proxyip(), conn:ip(), conn:serverport());
		return;
	end

	if mapping.service == nil then
		local services = portmanager.get_registered_services();
		local service = services[mapping.service_name];
		if service ~= nil then
			mapping.service = service;
		else
			conn:close();
			module:log("error", "Connection %s-%s terminated: Could not process mapping for unknown service %s",
				conn:proxyip(), conn:ip(), mapping.service_name);
			return;
		end
	end

	-- Pass connection to actual service listener and simulate onconnect/onincoming callbacks
	local service_listener = mapping.service[1].listener;

	module:log("debug", "Passing connection %s-%s to service %s", conn:proxyip(), conn:ip(), mapping.service_name);
	conn:setlistener(service_listener);
	if service_listener.onconnect then
		service_listener.onconnect(conn);
	end
	return service_listener.onincoming(conn, session.buffer);
end

function listener.onconnect(conn)
	sessions[conn] = {
		handler = nil;
		buffer = nil;
	};
end

function listener.onincoming(conn, data)
	-- Abort processing if no data has been received
	if not data then
		return;
	end

	-- Lookup session for connection and append received data to buffer
	local session = sessions[conn];
	session.buffer = session.buffer and session.buffer .. data or data;

	-- Attempt to determine protocol handler if not done previously
	if session.handler == nil then
		-- Match current session buffer against all known protocol signatures to determine protocol handler
		for handler_name, handler in pairs(protocol_handlers) do
			if session.buffer:find("^" .. handler.signature) ~= nil then
				session.handler = handler.callback;
				module:log("debug", "Detected %s connection from %s:%d", handler_name, conn:ip(), conn:port());
				break;
			end
		end

		-- Decide between waiting for a complete header signature or terminating the connection when no handler has been found
		if session.handler == nil then
			-- Terminate connection if buffer size has exceeded tolerable maximum size
			if #session.buffer > protocol_max_header_length then
				conn:close();
				module:log("warn", "Connection %s terminated: No valid header signature within %d bytes", conn:ip(), protocol_max_header_length);
			end

			-- Skip further processing without a valid protocol handler
			return;
		end
	end

	-- Execute proxy protocol handler and process response
	local response, proxy_data = session.handler(conn, session);
	if response == protocol_handler_responses.success then
		module:log("debug", "Received PROXY header from %s: %s", conn:ip(), proxy_data:describe());
		return wrap_proxy_connection(conn, session, proxy_data);
	elseif response == protocol_handler_responses.postpone then
		module:log("debug", "Postponed parsing of incomplete PROXY header received from %s", conn:ip());
		return;
	elseif response == protocol_handler_responses.failure then
		conn:close();
		module:log("warn", "Connection %s terminated: Could not process PROXY header from client, see previous log messages.", conn:ip());
		return;
	else
		-- This code should be never reached, but is included for completeness
		conn:close();
		module:log("error", "Connection terminated: Received invalid protocol handler response with code %d", response);
		return;
	end
end

function listener.ondisconnect(conn)
	sessions[conn] = nil;
end

-- Process all configured port mappings
local config_ports = module:get_option_set("proxy_ports", {});
local config_mappings = module:get_option("proxy_port_mappings", {});
for port in config_ports do
	if config_mappings[port] ~= nil then
		mappings[port] = {
			service_name = config_mappings[port],
			service = nil
		};
	else
		module:log("warn", "No port<>service mapping found for port: %d", port);
	end
end

-- Register the proxy network listener
module:provides("net", {
	name = "proxy";
	listener = listener;
});