File

local set = require "util.set";

insulate("net.resolvers.service", function ()
	local adns = {
		resolver = function ()
			return {
				lookup = function (_, cb, qname, qtype, qclass)
					if qname == "_xmpp-server._tcp.example.com"
					   and (qtype or "SRV") == "SRV"
					   and (qclass or "IN") == "IN" then
						cb({
							{ -- 60+35+60
								srv = { target = "xmpp0-a.example.com", port = 5228, priority = 0, weight = 60 };
							};
							{
								srv = { target = "xmpp0-b.example.com", port = 5216, priority = 0, weight = 35 };
							};
							{
								srv = { target = "xmpp0-c.example.com", port = 5200, priority = 0, weight = 0 };
							};
							{
								srv = { target = "xmpp0-d.example.com", port = 5256, priority = 0, weight = 120 };
							};

							{
								srv = { target = "xmpp1-a.example.com", port = 5273, priority = 1, weight = 30 };
							};
							{
								srv = { target = "xmpp1-b.example.com", port = 5274, priority = 1, weight = 30 };
							};

							{
								srv = { target = "xmpp2.example.com", port = 5275, priority = 2, weight = 0 };
							};
						});
					elseif qname == "_xmpp-server._tcp.single.example.com"
					   and (qtype or "SRV") == "SRV"
					   and (qclass or "IN") == "IN" then
						cb({
							{
								srv = { target = "xmpp0-a.example.com", port = 5269, priority = 0, weight = 0 };
							};
						});
					elseif qname == "_xmpp-server._tcp.half.example.com"
					   and (qtype or "SRV") == "SRV"
					   and (qclass or "IN") == "IN" then
						cb({
							{
								srv = { target = "xmpp0-a.example.com", port = 5269, priority = 0, weight = 0 };
							};
							{
								srv = { target = "xmpp0-b.example.com", port = 5270, priority = 0, weight = 1 };
							};
						});
					elseif qtype == "A" then
						local l = qname:match("%-(%a)%.example.com$") or "1";
						local d = ("%d"):format(l:byte())
						cb({
							{
								a = "127.0.0."..d;
							};
						});
					elseif qtype == "AAAA" then
						local l = qname:match("%-(%a)%.example.com$") or "1";
						local d = ("%04d"):format(l:byte())
						cb({
							{
								aaaa = "fdeb:9619:649e:c7d9::"..d;
							};
						});
					else
						cb(nil);
					end
				end;
			};
		end;
	};
	package.loaded["net.adns"] = mock(adns);
	local resolver = require "net.resolvers.service";
	math.randomseed(os.time());
	it("works for 99% of deployments", function ()
		-- Most deployments only have a single SRV record, let's make
		-- sure that works okay

		local expected_targets = set.new({
			-- xmpp0-a
			"tcp4  127.0.0.97  5269";
			"tcp6  fdeb:9619:649e:c7d9::0097  5269";
		});
		local received_targets = set.new({});

		local r = resolver.new("single.example.com", "xmpp-server");
		local done = false;
		local function handle_target(...)
			if ... == nil then
				done = true;
				-- No more targets
				return;
			end
			received_targets:add(table.concat({ ... }, "  ", 1, 3));
		end
		r:next(handle_target);
		while not done do
			r:next(handle_target);
		end

		-- We should have received all expected targets, and no unexpected
		-- ones:
		assert.truthy(set.xor(received_targets, expected_targets):empty());
	end);

	it("supports A/AAAA fallback", function ()
		-- Many deployments don't have any SRV records, so we should
		-- fall back to A/AAAA records instead when that is the case

		local expected_targets = set.new({
			-- xmpp0-a
			"tcp4  127.0.0.97  5269";
			"tcp6  fdeb:9619:649e:c7d9::0097  5269";
		});
		local received_targets = set.new({});

		local r = resolver.new("xmpp0-a.example.com", "xmpp-server", "tcp", { default_port = 5269 });
		local done = false;
		local function handle_target(...)
			if ... == nil then
				done = true;
				-- No more targets
				return;
			end
			received_targets:add(table.concat({ ... }, "  ", 1, 3));
		end
		r:next(handle_target);
		while not done do
			r:next(handle_target);
		end

		-- We should have received all expected targets, and no unexpected
		-- ones:
		assert.truthy(set.xor(received_targets, expected_targets):empty());
	end);


	it("works", function ()
		local expected_targets = set.new({
			-- xmpp0-a
			"tcp4  127.0.0.97  5228";
			"tcp6  fdeb:9619:649e:c7d9::0097  5228";
			"tcp4  127.0.0.97  5273";
			"tcp6  fdeb:9619:649e:c7d9::0097  5273";

			-- xmpp0-b
			"tcp4  127.0.0.98  5274";
			"tcp6  fdeb:9619:649e:c7d9::0098  5274";
			"tcp4  127.0.0.98  5216";
			"tcp6  fdeb:9619:649e:c7d9::0098  5216";

			-- xmpp0-c
			"tcp4  127.0.0.99  5200";
			"tcp6  fdeb:9619:649e:c7d9::0099  5200";

			-- xmpp0-d
			"tcp4  127.0.0.100  5256";
			"tcp6  fdeb:9619:649e:c7d9::0100  5256";

			-- xmpp2
			"tcp4  127.0.0.49  5275";
			"tcp6  fdeb:9619:649e:c7d9::0049  5275";

		});
		local received_targets = set.new({});

		local r = resolver.new("example.com", "xmpp-server");
		local done = false;
		local function handle_target(...)
			if ... == nil then
				done = true;
				-- No more targets
				return;
			end
			received_targets:add(table.concat({ ... }, "  ", 1, 3));
		end
		r:next(handle_target);
		while not done do
			r:next(handle_target);
		end

		-- We should have received all expected targets, and no unexpected
		-- ones:
		assert.truthy(set.xor(received_targets, expected_targets):empty());
	end);

	it("balances across weights correctly #slow", function ()
		-- This mimics many repeated connections to 'example.com' (mock
		-- records defined above), and records the port number of the
		-- first target. Therefore it (should) only return priority
		-- 0 records, and the input data is constructed such that the
		-- last two digits of the port number represent the percentage
		-- of times that record should (on average) be picked first.

		-- To prevent random test failures, we test across a handful
		-- of fixed (randomly selected) seeds.
		for _, seed in ipairs({ 8401877, 3943829, 7830992 }) do
			math.randomseed(seed);

			local results = {};
			local function run()
				local run_results = {};
				local r = resolver.new("example.com", "xmpp-server");
				local function record_target(...)
					if ... == nil then
						-- No more targets
						return;
					end
					run_results = { ... };
				end
				r:next(record_target);
				return run_results[3];
			end

			for _ = 1, 1000 do
				local port = run();
				results[port] = (results[port] or 0) + 1;
			end

			local ports = {};
			for port in pairs(results) do
				table.insert(ports, port);
			end
			table.sort(ports);
			for _, port in ipairs(ports) do
				--print("PORT", port, tostring((results[port]/1000) * 100).."% hits (expected "..tostring(port-5200).."%)");
				local hit_pct = (results[port]/1000) * 100;
				local expected_pct = port - 5200;
				--print(hit_pct, expected_pct, math.abs(hit_pct - expected_pct));
				assert.is_true(math.abs(hit_pct - expected_pct) < 5);
			end
			--print("---");
		end
	end);
end);