--% requires: s2sout-pre-connect-event
-- mod_s2sout_yggdrasil - Removes non-yggdrasil addresses from DNS response.
--
-- Usage:
-- modules_enabled = {
--     "s2sout_yggdrasil";
-- }
-- allow_yggdrasil = true
-- allow_clearnet = false
--
-- Based on mod_s2sout_filter by MattJ [1]
-- [1] https://matthewwild.co.uk/uploads/mod_s2sout_filter.lua

local filtering_resolver = {}
local filtering_resolver_mt = { __index = filtering_resolver };
local allow_yggdrasil = module:get_option("allow_yggdrasil", true);
local allow_clearnet = module:get_option("allow_clearnet", false);

local function is_yggdrasil(ip)
	local f = string.match(ip, "(%x+):")
	if (f) then
		local addr = tonumber(f, 16);
		return addr >= 512 and addr <= 1023
	end
	return false
end

function filtering_resolver:next(callback)
	local function filtering_callback(conn_type, ip, port, extra, more_targets_available)
		if conn_type == nil then
			-- No more targets, nothing to filter
			callback(conn_type, ip, port, extra, more_targets_available);
			return;
		end

		local is_ygg = is_yggdrasil(ip);
		local is_clearnet = not is_ygg;

		if (is_ygg and allow_yggdrasil) or (is_clearnet and allow_clearnet) then
			-- We are not filtering out this target, so pass it to
			-- Prosody's original callback
			callback(conn_type, ip, port, extra, more_targets_available);
		else
			self.original_resolver:next(filtering_callback);
		end
	end

	self.original_resolver:next(filtering_callback);
end

local function new_filtering_resolver(original_resolver)
	return setmetatable({
		original_resolver = original_resolver;
	}, filtering_resolver_mt);
end


module:hook("s2sout-pre-connect", function(event)
	-- Replace the resolver for this connection, with one that filters the targets
	event.resolver = new_filtering_resolver(event.resolver);
end);