-- SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
-- SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
-- Licensed under OWL v1.0+. See LICENSE.

local testimony = require("testimony")
local iter = require("dns.iter")
local types = require("dns.types")

local TYPE = types.TYPE
local RCODE = types.RCODE

local testify = testimony.new("== dns.iter ==")

---------------------------------------------------------------------------
-- Helpers: build hand-crafted response tables for classify/extract tests
---------------------------------------------------------------------------

local make_header = function(overrides)
	local hdr = {
		id = 0xABCD,
		qr = 1,
		opcode = 0,
		aa = false,
		tc = false,
		rd = false,
		ra = false,
		rcode = 0,
	}
	if overrides then
		for k, v in pairs(overrides) do
			hdr[k] = v
		end
	end
	return hdr
end

local make_question = function(name, qtype)
	return { { name = name or "example.com.", qtype = qtype or TYPE.A, qclass = 1 } }
end

-- =========================================================================
-- ROOT_SERVERS
-- =========================================================================

testify:that("ROOT_SERVERS has 13 entries", function()
	testimony.assert_equal(13, #iter.ROOT_SERVERS)
end)

testify:that("ROOT_SERVERS entries have required fields", function()
	for _, rs in ipairs(iter.ROOT_SERVERS) do
		testimony.assert_not_nil(rs.name, "missing name field")
		testimony.assert_not_nil(rs.host, "missing host field")
		testimony.assert_not_nil(rs.v6, "missing v6 field")
		-- IPv4 format check
		testimony.assert_match("^%d+%.%d+%.%d+%.%d+$", rs.host, "bad IPv4: " .. rs.host)
	end
end)

-- =========================================================================
-- classify: answer
-- =========================================================================

testify:that("classify: answer with matching records", function()
	local resp = {
		header = make_header({ aa = true }),
		question = make_question(),
		answer = {
			{ name = "example.com.", type = TYPE.A, class = 1, ttl = 300, rdata = { address = "93.184.216.34" } },
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("answer", kind)
	testimony.assert_equal(1, #info.records)
	testimony.assert_equal(true, info.authoritative)
end)

testify:that("classify: non-authoritative answer", function()
	local resp = {
		header = make_header({ aa = false }),
		question = make_question(),
		answer = {
			{ name = "example.com.", type = TYPE.A, class = 1, ttl = 300, rdata = { address = "93.184.216.34" } },
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("answer", kind)
	testimony.assert_equal(false, info.authoritative)
end)

testify:that("classify: multiple answer records", function()
	local resp = {
		header = make_header({ aa = true }),
		question = make_question(),
		answer = {
			{ name = "example.com.", type = TYPE.A, class = 1, ttl = 300, rdata = { address = "93.184.216.34" } },
			{ name = "example.com.", type = TYPE.A, class = 1, ttl = 300, rdata = { address = "93.184.216.35" } },
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("answer", kind)
	testimony.assert_equal(2, #info.records)
end)

-- =========================================================================
-- classify: cname
-- =========================================================================

testify:that("classify: CNAME without matching qtype", function()
	local resp = {
		header = make_header(),
		question = make_question(),
		answer = {
			{ name = "example.com.", type = TYPE.CNAME, class = 1, ttl = 300, rdata = { cname = "real.example.com." } },
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("cname", kind)
	testimony.assert_equal("real.example.com.", info.target)
end)

testify:that("classify: CNAME + A for target returns cname not answer", function()
	local resp = {
		header = make_header({ aa = true }),
		question = make_question(),
		answer = {
			{ name = "example.com.", type = TYPE.CNAME, class = 1, ttl = 300, rdata = { cname = "real.example.com." } },
			{ name = "real.example.com.", type = TYPE.A, class = 1, ttl = 300, rdata = { address = "1.2.3.4" } },
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("cname", kind)
	testimony.assert_equal("real.example.com.", info.target)
end)

-- =========================================================================
-- classify: referral
-- =========================================================================

testify:that("classify: referral with NS in authority", function()
	local resp = {
		header = make_header(),
		question = make_question(),
		answer = {},
		authority = {
			{ name = "com.", type = TYPE.NS, class = 1, ttl = 172800, rdata = { nsdname = "a.gtld-servers.net." } },
			{ name = "com.", type = TYPE.NS, class = 1, ttl = 172800, rdata = { nsdname = "b.gtld-servers.net." } },
		},
		additional = {
			{
				name = "a.gtld-servers.net.",
				type = TYPE.A,
				class = 1,
				ttl = 172800,
				rdata = { address = "192.5.6.30" },
			},
			{
				name = "b.gtld-servers.net.",
				type = TYPE.A,
				class = 1,
				ttl = 172800,
				rdata = { address = "192.33.14.30" },
			},
		},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("referral", kind)
	testimony.assert_equal("com.", info.zone)
	testimony.assert_equal(2, #info.ns)
	testimony.assert_equal("a.gtld-servers.net.", info.ns[1].name)
	testimony.assert_equal("192.5.6.30", info.glue["a.gtld-servers.net."].ipv4)
	testimony.assert_equal("192.33.14.30", info.glue["b.gtld-servers.net."].ipv4)
end)

testify:that("classify: referral without glue", function()
	local resp = {
		header = make_header(),
		question = make_question(),
		answer = {},
		authority = {
			{
				name = "example.com.",
				type = TYPE.NS,
				class = 1,
				ttl = 86400,
				rdata = { nsdname = "ns1.otherdomain.com." },
			},
		},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("referral", kind)
	testimony.assert_equal("example.com.", info.zone)
	testimony.assert_equal(1, #info.ns)
	-- No glue for this NS
	testimony.assert_nil(info.glue["ns1.otherdomain.com."])
end)

-- =========================================================================
-- classify: nxdomain
-- =========================================================================

testify:that("classify: NXDOMAIN", function()
	local resp = {
		header = make_header({ rcode = RCODE.NXDOMAIN }),
		question = make_question("nonexistent.example.com."),
		answer = {},
		authority = {
			{
				name = "example.com.",
				type = TYPE.SOA,
				class = 1,
				ttl = 900,
				rdata = {
					mname = "ns1.example.com.",
					rname = "admin.example.com.",
					serial = 2024010101,
					refresh = 3600,
					retry = 900,
					expire = 604800,
					minimum = 86400,
				},
			},
		},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "nonexistent.example.com.")
	testimony.assert_equal("nxdomain", kind)
	testimony.assert_not_nil(info.soa)
	testimony.assert_equal("ns1.example.com.", info.soa.rdata.mname)
end)

testify:that("classify: NXDOMAIN without SOA", function()
	local resp = {
		header = make_header({ rcode = RCODE.NXDOMAIN }),
		question = make_question(),
		answer = {},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("nxdomain", kind)
	testimony.assert_nil(info.soa)
end)

-- =========================================================================
-- classify: nodata
-- =========================================================================

testify:that("classify: NODATA (NOERROR, empty answer, no NS)", function()
	local resp = {
		header = make_header(),
		question = make_question(),
		answer = {},
		authority = {
			{
				name = "example.com.",
				type = TYPE.SOA,
				class = 1,
				ttl = 900,
				rdata = {
					mname = "ns1.example.com.",
					rname = "admin.example.com.",
					serial = 2024010101,
					refresh = 3600,
					retry = 900,
					expire = 604800,
					minimum = 86400,
				},
			},
		},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.AAAA, "example.com.")
	testimony.assert_equal("nodata", kind)
	testimony.assert_not_nil(info.soa)
end)

testify:that("classify: NODATA without SOA", function()
	local resp = {
		header = make_header(),
		question = make_question(),
		answer = {},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("nodata", kind)
	testimony.assert_nil(info.soa)
end)

-- =========================================================================
-- classify: error
-- =========================================================================

testify:that("classify: SERVFAIL", function()
	local resp = {
		header = make_header({ rcode = RCODE.SERVFAIL }),
		question = make_question(),
		answer = {},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("error", kind)
	testimony.assert_equal(RCODE.SERVFAIL, info.rcode)
	testimony.assert_equal("SERVFAIL", info.rcode_name)
end)

testify:that("classify: REFUSED", function()
	local resp = {
		header = make_header({ rcode = RCODE.REFUSED }),
		question = make_question(),
		answer = {},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "example.com.")
	testimony.assert_equal("error", kind)
	testimony.assert_equal(RCODE.REFUSED, info.rcode)
	testimony.assert_equal("REFUSED", info.rcode_name)
end)

testify:that("classify: ignores CNAME with non-matching owner name", function()
	local resp = {
		header = make_header(),
		question = make_question("my.vultr.com.cdn.cloudflare.net."),
		answer = {
			{
				name = "cdn.cloudflare.net.",
				type = TYPE.CNAME,
				class = 1,
				ttl = 300,
				rdata = { cname = "cloudflare.net." },
			},
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "my.vultr.com.cdn.cloudflare.net.")
	-- Should NOT be "cname" — the CNAME's owner doesn't match qname
	testimony.assert_equal("nodata", kind)
end)

testify:that("classify: ignores A records with non-matching owner name", function()
	local resp = {
		header = make_header(),
		question = make_question("my.vultr.com.cdn.cloudflare.net."),
		answer = {
			{
				name = "cloudflare.net.",
				type = TYPE.A,
				class = 1,
				ttl = 300,
				rdata = { address = "167.179.118.50" },
			},
		},
		authority = {},
		additional = {},
	}
	local kind, info = iter.classify(resp, TYPE.A, "my.vultr.com.cdn.cloudflare.net.")
	testimony.assert_equal("nodata", kind)
end)

-- =========================================================================
-- extract_referral
-- =========================================================================

testify:that("extract_referral: full glue (A records)", function()
	local resp = {
		authority = {
			{ name = "com.", type = TYPE.NS, class = 1, ttl = 172800, rdata = { nsdname = "a.gtld-servers.net." } },
			{ name = "com.", type = TYPE.NS, class = 1, ttl = 172800, rdata = { nsdname = "b.gtld-servers.net." } },
		},
		additional = {
			{
				name = "a.gtld-servers.net.",
				type = TYPE.A,
				class = 1,
				ttl = 172800,
				rdata = { address = "192.5.6.30" },
			},
			{
				name = "b.gtld-servers.net.",
				type = TYPE.A,
				class = 1,
				ttl = 172800,
				rdata = { address = "192.33.14.30" },
			},
		},
	}
	local zone, ns, glue = iter.extract_referral(resp)
	testimony.assert_equal("com.", zone)
	testimony.assert_equal(2, #ns)
	testimony.assert_equal("192.5.6.30", glue["a.gtld-servers.net."].ipv4)
	testimony.assert_equal("192.33.14.30", glue["b.gtld-servers.net."].ipv4)
end)

testify:that("extract_referral: partial glue", function()
	local resp = {
		authority = {
			{ name = "example.com.", type = TYPE.NS, class = 1, ttl = 86400, rdata = { nsdname = "ns1.example.com." } },
			{
				name = "example.com.",
				type = TYPE.NS,
				class = 1,
				ttl = 86400,
				rdata = { nsdname = "ns2.otherdomain.net." },
			},
		},
		additional = {
			{ name = "ns1.example.com.", type = TYPE.A, class = 1, ttl = 86400, rdata = { address = "10.0.0.1" } },
			-- no glue for ns2.otherdomain.net. (glueless)
		},
	}
	local zone, ns, glue = iter.extract_referral(resp)
	testimony.assert_equal("example.com.", zone)
	testimony.assert_equal(2, #ns)
	testimony.assert_equal("10.0.0.1", glue["ns1.example.com."].ipv4)
	testimony.assert_nil(glue["ns2.otherdomain.net."])
end)

testify:that("extract_referral: no glue at all", function()
	local resp = {
		authority = {
			{
				name = "example.com.",
				type = TYPE.NS,
				class = 1,
				ttl = 86400,
				rdata = { nsdname = "ns1.otherdomain.com." },
			},
		},
		additional = {},
	}
	local zone, ns, glue = iter.extract_referral(resp)
	testimony.assert_equal("example.com.", zone)
	testimony.assert_equal(1, #ns)
	testimony.assert_nil(glue["ns1.otherdomain.com."])
end)

testify:that("extract_referral: both A and AAAA glue", function()
	local resp = {
		authority = {
			{ name = "example.com.", type = TYPE.NS, class = 1, ttl = 86400, rdata = { nsdname = "ns1.example.com." } },
		},
		additional = {
			{ name = "ns1.example.com.", type = TYPE.A, class = 1, ttl = 86400, rdata = { address = "10.0.0.1" } },
			{
				name = "ns1.example.com.",
				type = TYPE.AAAA,
				class = 1,
				ttl = 86400,
				rdata = { address = "2001:db8::1" },
			},
		},
	}
	local _, _, glue = iter.extract_referral(resp)
	testimony.assert_equal("10.0.0.1", glue["ns1.example.com."].ipv4)
	testimony.assert_equal("2001:db8::1", glue["ns1.example.com."].ipv6)
end)

testify:that("extract_referral: no authority returns nil zone", function()
	local resp = {
		authority = {},
		additional = {},
	}
	local zone, ns, _ = iter.extract_referral(resp)
	testimony.assert_nil(zone)
	testimony.assert_equal(0, #ns)
end)

testify:that("extract_referral: mixed authority records (NS + SOA)", function()
	local resp = {
		authority = {
			{ name = "com.", type = TYPE.NS, class = 1, ttl = 172800, rdata = { nsdname = "a.gtld-servers.net." } },
			{
				name = "com.",
				type = TYPE.SOA,
				class = 1,
				ttl = 900,
				rdata = {
					mname = "a.gtld-servers.net.",
					rname = "nstld.verisign-grs.com.",
					serial = 2024010101,
					refresh = 1800,
					retry = 900,
					expire = 604800,
					minimum = 86400,
				},
			},
		},
		additional = {
			{
				name = "a.gtld-servers.net.",
				type = TYPE.A,
				class = 1,
				ttl = 172800,
				rdata = { address = "192.5.6.30" },
			},
		},
	}
	local zone, ns, glue = iter.extract_referral(resp)
	testimony.assert_equal("com.", zone)
	testimony.assert_equal(1, #ns) -- Only the NS record, not SOA
	testimony.assert_equal("192.5.6.30", glue["a.gtld-servers.net."].ipv4)
end)

-- =========================================================================
-- query_server: input validation
-- =========================================================================

testify:that("query_server: rejects invalid qtype", function()
	local resp, err = iter.query_server("198.41.0.4", "example.com.", "BOGUS")
	testimony.assert_nil(resp)
	testimony.assert_match("unknown query type", err)
end)

testify:that("query_server: rejects invalid domain name", function()
	local resp, err = iter.query_server("198.41.0.4", "bad!name.com", "A")
	testimony.assert_nil(resp)
	testimony.assert_match("invalid domain name", err)
end)

-- =========================================================================
-- trace: input validation
-- =========================================================================

testify:that("trace: rejects invalid qtype", function()
	local steps, err = iter.trace("example.com", "BOGUS")
	testimony.assert_nil(steps)
	testimony.assert_match("unknown query type", err)
end)

testify:that("trace: rejects invalid domain name", function()
	local steps, err = iter.trace("bad!name.com", "A")
	testimony.assert_nil(steps)
	testimony.assert_match("invalid domain name", err)
end)

testify:that("trace: rejects empty root servers", function()
	local steps, err = iter.trace("example.com", "A", { root_servers = {} })
	testimony.assert_nil(steps)
	testimony.assert_match("no root servers", err)
end)

testify:conclude()