-- 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 wire = require("dns.wire")
local types = require("dns.types")

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

---------------------------------------------------------------------------
-- Helper: build raw bytes from a table of integer byte values
---------------------------------------------------------------------------
local function raw(...)
	local bytes = { ... }
	local chars = {}
	for i = 1, #bytes do
		chars[i] = string.char(bytes[i])
	end
	return table.concat(chars)
end

-- =========================================================================
-- Constants
-- =========================================================================

testify:that("TYPE forward lookup works", function()
	testimony.assert_equal(1, types.TYPE.A)
	testimony.assert_equal(2, types.TYPE.NS)
	testimony.assert_equal(5, types.TYPE.CNAME)
	testimony.assert_equal(6, types.TYPE.SOA)
	testimony.assert_equal(12, types.TYPE.PTR)
	testimony.assert_equal(15, types.TYPE.MX)
	testimony.assert_equal(16, types.TYPE.TXT)
	testimony.assert_equal(28, types.TYPE.AAAA)
	testimony.assert_equal(33, types.TYPE.SRV)
	testimony.assert_equal(41, types.TYPE.OPT)
	testimony.assert_equal(255, types.TYPE.ANY)
end)

testify:that("TYPE reverse lookup works", function()
	testimony.assert_equal("A", types.TYPE_NAME[1])
	testimony.assert_equal("NS", types.TYPE_NAME[2])
	testimony.assert_equal("CNAME", types.TYPE_NAME[5])
	testimony.assert_equal("SOA", types.TYPE_NAME[6])
	testimony.assert_equal("PTR", types.TYPE_NAME[12])
	testimony.assert_equal("MX", types.TYPE_NAME[15])
	testimony.assert_equal("TXT", types.TYPE_NAME[16])
	testimony.assert_equal("AAAA", types.TYPE_NAME[28])
	testimony.assert_equal("SRV", types.TYPE_NAME[33])
	testimony.assert_equal("OPT", types.TYPE_NAME[41])
	testimony.assert_equal("ANY", types.TYPE_NAME[255])
end)

testify:that("CLASS forward and reverse lookup works", function()
	testimony.assert_equal(1, types.CLASS.IN)
	testimony.assert_equal(3, types.CLASS.CH)
	testimony.assert_equal(4, types.CLASS.HS)
	testimony.assert_equal(255, types.CLASS.ANY)
	testimony.assert_equal("IN", types.CLASS_NAME[1])
	testimony.assert_equal("CH", types.CLASS_NAME[3])
end)

testify:that("RCODE forward and reverse lookup works", function()
	testimony.assert_equal(0, types.RCODE.NOERROR)
	testimony.assert_equal(3, types.RCODE.NXDOMAIN)
	testimony.assert_equal(5, types.RCODE.REFUSED)
	testimony.assert_equal("NOERROR", types.RCODE_NAME[0])
	testimony.assert_equal("NXDOMAIN", types.RCODE_NAME[3])
end)

testify:that("OPCODE forward and reverse lookup works", function()
	testimony.assert_equal(0, types.OPCODE.QUERY)
	testimony.assert_equal(1, types.OPCODE.IQUERY)
	testimony.assert_equal(2, types.OPCODE.STATUS)
	testimony.assert_equal("QUERY", types.OPCODE_NAME[0])
end)

-- =========================================================================
-- Codec registry
-- =========================================================================

testify:that("get_codec returns nil for unknown type", function()
	testimony.assert_nil(types.get_codec(65534))
end)

testify:that("get_codec returns codec for known type", function()
	testimony.assert_not_nil(types.get_codec(types.TYPE.A))
end)

-- =========================================================================
-- A codec
-- =========================================================================

testify:that("A codec decodes 4 bytes to dotted quad", function()
	local r = wire.reader(raw(0x5d, 0xb8, 0xd8, 0x22))
	local codec = types.get_codec(types.TYPE.A)
	local rdata = codec.decode(r)
	testimony.assert_equal("93.184.216.34", rdata.address)
end)

testify:that("A codec encode roundtrips", function()
	local codec = types.get_codec(types.TYPE.A)
	local w = wire.writer()
	codec.encode(w, { address = "10.0.0.1" })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal("10.0.0.1", rdata.address)
end)

testify:that("A codec decode fails on truncated input", function()
	local r = wire.reader(raw(0x0A, 0x00))
	local codec = types.get_codec(types.TYPE.A)
	local rdata, err = codec.decode(r)
	testimony.assert_nil(rdata)
	testimony.assert_not_nil(err)
end)

-- =========================================================================
-- AAAA codec
-- =========================================================================

testify:that("AAAA codec decodes 2001:db8::1", function()
	local r =
		wire.reader(raw(0x20, 0x01, 0x0d, 0xb8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01))
	local codec = types.get_codec(types.TYPE.AAAA)
	local rdata = codec.decode(r)
	testimony.assert_equal("2001:db8::1", rdata.address)
end)

testify:that("AAAA codec decodes all zeros as ::", function()
	local r = wire.reader(string.rep(raw(0x00), 16))
	local codec = types.get_codec(types.TYPE.AAAA)
	local rdata = codec.decode(r)
	testimony.assert_equal("::", rdata.address)
end)

testify:that("AAAA codec decodes ::1 (loopback)", function()
	local r =
		wire.reader(raw(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01))
	local codec = types.get_codec(types.TYPE.AAAA)
	local rdata = codec.decode(r)
	testimony.assert_equal("::1", rdata.address)
end)

testify:that("AAAA codec does not use :: for single zero group", function()
	-- fe80:0:0:0:0:0:0:1 → has 6 zero groups → should use ::
	-- fe80:1:0:2:0:0:0:3 → has groups: fe80, 1, 0, 2, 0, 0, 0, 3
	-- longest run is groups 5-7 (3 zeros) → fe80:1:0:2::3
	-- But for single zero: fe80:0:1:2:3:4:5:6 → fe80:0:1:2:3:4:5:6 (no ::)
	local r =
		wire.reader(raw(0xfe, 0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x04, 0x00, 0x05, 0x00, 0x06))
	local codec = types.get_codec(types.TYPE.AAAA)
	local rdata = codec.decode(r)
	-- Single zero group at position 2 → no :: per RFC5952
	testimony.assert_equal("fe80:0:1:2:3:4:5:6", rdata.address)
end)

testify:that("AAAA codec encode roundtrips", function()
	local codec = types.get_codec(types.TYPE.AAAA)
	local w = wire.writer()
	codec.encode(w, { address = "2001:db8::1" })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal("2001:db8::1", rdata.address)
end)

testify:that("AAAA codec roundtrips all zeros", function()
	local codec = types.get_codec(types.TYPE.AAAA)
	local w = wire.writer()
	codec.encode(w, { address = "::" })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal("::", rdata.address)
end)

testify:that("AAAA codec decode fails on truncated input", function()
	local r = wire.reader(raw(0x20, 0x01, 0x0d, 0xb8))
	local codec = types.get_codec(types.TYPE.AAAA)
	local rdata, err = codec.decode(r)
	testimony.assert_nil(rdata)
	testimony.assert_not_nil(err)
end)

-- =========================================================================
-- CNAME codec
-- =========================================================================

testify:that("CNAME codec decodes name", function()
	local data = raw(0x06) .. "target" .. raw(0x03) .. "com" .. raw(0x00)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.CNAME)
	local rdata = codec.decode(r)
	testimony.assert_equal("target.com.", rdata.cname)
end)

testify:that("CNAME codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.CNAME)
	local w = wire.writer()
	codec.encode(w, { cname = "alias.example.com." })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal("alias.example.com.", rdata.cname)
end)

testify:that("CNAME codec resolves compression pointer via sub_reader", function()
	-- Simulate message with name at offset 12, sub_reader RDATA with pointer
	local header = string.rep(raw(0x00), 12)
	local name1 = raw(0x07) .. "example" .. raw(0x03) .. "com" .. raw(0x00)
	local rdata_bytes = raw(0xC0, 0x0C) -- pointer to offset 12
	local data = header .. name1 .. rdata_bytes

	local r = wire.reader(data)
	r:seek(26) -- position at rdata_bytes
	local sub = r:sub_reader(2)
	local codec = types.get_codec(types.TYPE.CNAME)
	local rdata = codec.decode(sub)
	testimony.assert_equal("example.com.", rdata.cname)
end)

-- =========================================================================
-- NS codec
-- =========================================================================

testify:that("NS codec decodes name", function()
	local data = raw(0x03) .. "ns1" .. raw(0x07) .. "example" .. raw(0x03) .. "com" .. raw(0x00)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.NS)
	local rdata = codec.decode(r)
	testimony.assert_equal("ns1.example.com.", rdata.nsdname)
end)

testify:that("NS codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.NS)
	local w = wire.writer()
	codec.encode(w, { nsdname = "ns2.example.com." })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal("ns2.example.com.", rdata.nsdname)
end)

-- =========================================================================
-- PTR codec
-- =========================================================================

testify:that("PTR codec decodes name", function()
	local data = raw(0x07) .. "example" .. raw(0x03) .. "com" .. raw(0x00)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.PTR)
	local rdata = codec.decode(r)
	testimony.assert_equal("example.com.", rdata.ptrdname)
end)

testify:that("PTR codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.PTR)
	local w = wire.writer()
	codec.encode(w, { ptrdname = "one.one.one.one." })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal("one.one.one.one.", rdata.ptrdname)
end)

-- =========================================================================
-- MX codec
-- =========================================================================

testify:that("MX codec decodes preference and exchange", function()
	local data = raw(0x00, 0x0A) -- preference=10
		.. raw(0x04)
		.. "mail"
		.. raw(0x07)
		.. "example"
		.. raw(0x03)
		.. "com"
		.. raw(0x00)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.MX)
	local rdata = codec.decode(r)
	testimony.assert_equal(10, rdata.preference)
	testimony.assert_equal("mail.example.com.", rdata.exchange)
end)

testify:that("MX codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.MX)
	local w = wire.writer()
	codec.encode(w, { preference = 20, exchange = "mx.example.com." })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal(20, rdata.preference)
	testimony.assert_equal("mx.example.com.", rdata.exchange)
end)

testify:that("MX codec with compression pointer via sub_reader", function()
	local header = string.rep(raw(0x00), 12)
	local name1 = raw(0x07) .. "example" .. raw(0x03) .. "com" .. raw(0x00)
	-- RDATA: preference=10 + "mail" + pointer to offset 12
	local rdata_bytes = raw(0x00, 0x0A, 0x04) .. "mail" .. raw(0xC0, 0x0C)
	local data = header .. name1 .. rdata_bytes

	local r = wire.reader(data)
	r:seek(26) -- position at rdata_bytes
	local sub = r:sub_reader(9) -- 2 + 1 + 4 + 2 = 9
	local codec = types.get_codec(types.TYPE.MX)
	local rdata = codec.decode(sub)
	testimony.assert_equal(10, rdata.preference)
	testimony.assert_equal("mail.example.com.", rdata.exchange)
end)

-- =========================================================================
-- SOA codec
-- =========================================================================

testify:that("SOA codec decodes all 7 fields", function()
	-- Pad with 12-byte fake header so compression pointers are valid
	local w = wire.writer()
	for _ = 1, 12 do
		w:u8(0)
	end
	w:name("ns1.example.com.")
	w:name("admin.example.com.")
	w:u32(2024010101)
	w:u32(7200)
	w:u32(1800)
	w:u32(604800)
	w:u32(300)
	local r = wire.reader(w:output())
	r:skip(12) -- skip fake header
	local codec = types.get_codec(types.TYPE.SOA)
	local rdata = codec.decode(r)
	testimony.assert_equal("ns1.example.com.", rdata.mname)
	testimony.assert_equal("admin.example.com.", rdata.rname)
	testimony.assert_equal(2024010101, rdata.serial)
	testimony.assert_equal(7200, rdata.refresh)
	testimony.assert_equal(1800, rdata.retry)
	testimony.assert_equal(604800, rdata.expire)
	testimony.assert_equal(300, rdata.minimum)
end)

testify:that("SOA codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.SOA)
	local original = {
		mname = "ns.example.com.",
		rname = "hostmaster.example.com.",
		serial = 1000000,
		refresh = 3600,
		retry = 900,
		expire = 86400,
		minimum = 60,
	}
	-- Pad with 12-byte fake header so compression pointers are valid
	local w = wire.writer()
	for _ = 1, 12 do
		w:u8(0)
	end
	codec.encode(w, original)
	local r = wire.reader(w:output())
	r:skip(12) -- skip fake header
	local rdata = codec.decode(r)
	testimony.assert_equal(original.mname, rdata.mname)
	testimony.assert_equal(original.rname, rdata.rname)
	testimony.assert_equal(original.serial, rdata.serial)
	testimony.assert_equal(original.refresh, rdata.refresh)
	testimony.assert_equal(original.retry, rdata.retry)
	testimony.assert_equal(original.expire, rdata.expire)
	testimony.assert_equal(original.minimum, rdata.minimum)
end)

testify:that("SOA codec with compression via sub_reader", function()
	local header = string.rep(raw(0x00), 12)
	local name1 = raw(0x07) .. "example" .. raw(0x03) .. "com" .. raw(0x00)
	-- SOA RDATA: ns.example.com. + admin.example.com. + 5 u32s
	-- ns: 02 "ns" + C00C, admin: 05 "admin" + C00C, then 5×u32
	local soa_rdata = raw(0x02)
		.. "ns"
		.. raw(0xC0, 0x0C)
		.. raw(0x05)
		.. "admin"
		.. raw(0xC0, 0x0C)
		.. raw(0x00, 0x00, 0x00, 0x01) -- serial
		.. raw(0x00, 0x00, 0x1C, 0x20) -- refresh=7200
		.. raw(0x00, 0x00, 0x07, 0x08) -- retry=1800
		.. raw(0x00, 0x09, 0x3A, 0x80) -- expire=604800
		.. raw(0x00, 0x00, 0x01, 0x2C) -- minimum=300
	local data = header .. name1 .. soa_rdata

	local r = wire.reader(data)
	r:seek(26) -- start of soa_rdata
	local rdlen = 5 + 8 + 20 -- ns(5) + admin(8) + 5*u32(20) = 33
	local sub = r:sub_reader(rdlen)
	local codec = types.get_codec(types.TYPE.SOA)
	local rdata = codec.decode(sub)
	testimony.assert_equal("ns.example.com.", rdata.mname)
	testimony.assert_equal("admin.example.com.", rdata.rname)
	testimony.assert_equal(1, rdata.serial)
	testimony.assert_equal(7200, rdata.refresh)
end)

-- =========================================================================
-- TXT codec
-- =========================================================================

testify:that("TXT codec decodes single string", function()
	local data = raw(0x05) .. "hello"
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.TXT)
	local rdata = codec.decode(r)
	testimony.assert_equal(1, #rdata.strings)
	testimony.assert_equal("hello", rdata.strings[1])
	testimony.assert_equal("hello", rdata.text)
end)

testify:that("TXT codec decodes multiple strings", function()
	local data = raw(0x05) .. "hello" .. raw(0x05) .. "world"
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.TXT)
	local rdata = codec.decode(r)
	testimony.assert_equal(2, #rdata.strings)
	testimony.assert_equal("hello", rdata.strings[1])
	testimony.assert_equal("world", rdata.strings[2])
	testimony.assert_equal("helloworld", rdata.text)
end)

testify:that("TXT codec decodes empty string", function()
	local data = raw(0x00)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.TXT)
	local rdata = codec.decode(r)
	testimony.assert_equal(1, #rdata.strings)
	testimony.assert_equal("", rdata.strings[1])
end)

testify:that("TXT codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.TXT)
	local w = wire.writer()
	codec.encode(w, { strings = { "v=spf1 include:example.com", "~all" } })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal(2, #rdata.strings)
	testimony.assert_equal("v=spf1 include:example.com", rdata.strings[1])
	testimony.assert_equal("~all", rdata.strings[2])
end)

-- =========================================================================
-- SRV codec
-- =========================================================================

testify:that("SRV codec decodes all fields", function()
	local w = wire.writer()
	w:u16(10) -- priority
	w:u16(60) -- weight
	w:u16(5060) -- port
	w:name("sipserver.example.com.")
	local r = wire.reader(w:output())
	local codec = types.get_codec(types.TYPE.SRV)
	local rdata = codec.decode(r)
	testimony.assert_equal(10, rdata.priority)
	testimony.assert_equal(60, rdata.weight)
	testimony.assert_equal(5060, rdata.port)
	testimony.assert_equal("sipserver.example.com.", rdata.target)
end)

testify:that("SRV codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.SRV)
	local original = { priority = 5, weight = 100, port = 443, target = "web.example.com." }
	local w = wire.writer()
	codec.encode(w, original)
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_equal(original.priority, rdata.priority)
	testimony.assert_equal(original.weight, rdata.weight)
	testimony.assert_equal(original.port, rdata.port)
	testimony.assert_equal(original.target, rdata.target)
end)

testify:that("SRV codec with compression via sub_reader", function()
	local header = string.rep(raw(0x00), 12)
	local name1 = raw(0x07) .. "example" .. raw(0x03) .. "com" .. raw(0x00)
	-- SRV RDATA: priority=10, weight=60, port=5060, target="sip" + C00C
	local srv_rdata = raw(0x00, 0x0A, 0x00, 0x3C, 0x13, 0xC4) .. raw(0x03) .. "sip" .. raw(0xC0, 0x0C)
	local data = header .. name1 .. srv_rdata

	local r = wire.reader(data)
	r:seek(26) -- start of srv_rdata
	local sub = r:sub_reader(12) -- 6 + 1+3+2 = 12
	local codec = types.get_codec(types.TYPE.SRV)
	local rdata = codec.decode(sub)
	testimony.assert_equal(10, rdata.priority)
	testimony.assert_equal(60, rdata.weight)
	testimony.assert_equal(5060, rdata.port)
	testimony.assert_equal("sip.example.com.", rdata.target)
end)

-- =========================================================================
-- OPT codec
-- =========================================================================

testify:that("OPT codec decodes empty options", function()
	local r = wire.reader("")
	local codec = types.get_codec(types.TYPE.OPT)
	local rdata = codec.decode(r)
	testimony.assert_not_nil(rdata.options)
end)

testify:that("OPT codec decodes one option", function()
	-- Option code=10, length=4, data=0xDEADBEEF
	local data = raw(0x00, 0x0A, 0x00, 0x04, 0xDE, 0xAD, 0xBE, 0xEF)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.OPT)
	local rdata = codec.decode(r)
	testimony.assert_not_nil(rdata.options[10])
	testimony.assert_equal(4, #rdata.options[10])
end)

testify:that("OPT codec decodes multiple options", function()
	-- Option code=3, length=2, data=0x0102
	-- Option code=8, length=0, data=""
	local data = raw(0x00, 0x03, 0x00, 0x02, 0x01, 0x02, 0x00, 0x08, 0x00, 0x00)
	local r = wire.reader(data)
	local codec = types.get_codec(types.TYPE.OPT)
	local rdata = codec.decode(r)
	testimony.assert_not_nil(rdata.options[3])
	testimony.assert_equal(2, #rdata.options[3])
	testimony.assert_not_nil(rdata.options[8])
	testimony.assert_equal(0, #rdata.options[8])
end)

testify:that("OPT codec roundtrips", function()
	local codec = types.get_codec(types.TYPE.OPT)
	local w = wire.writer()
	codec.encode(w, { options = { [10] = raw(0xDE, 0xAD) } })
	local r = wire.reader(w:output())
	local rdata = codec.decode(r)
	testimony.assert_not_nil(rdata.options[10])
	testimony.assert_equal(2, #rdata.options[10])
end)

-- =========================================================================
-- IPv6 formatting edge cases
-- =========================================================================

testify:that("format_ipv6 leftmost longest run wins", function()
	-- Groups: 1, 0, 0, 2, 0, 0, 0, 3
	-- Runs: pos 2-3 (len 2), pos 5-7 (len 3)
	-- Longest is 5-7 → 1:0:0:2::3
	-- Wait, group indices are 1-based: groups[1]=1, [2]=0, [3]=0, [4]=2, [5]=0, [6]=0, [7]=0, [8]=3
	-- Run at 2-3 (len 2), run at 5-7 (len 3) → longest is 5-7
	local result = types.format_ipv6({ 1, 0, 0, 2, 0, 0, 0, 3 })
	testimony.assert_equal("1:0:0:2::3", result)
end)

testify:that("format_ipv6 all ones has no ::", function()
	local result = types.format_ipv6({ 1, 2, 3, 4, 5, 6, 7, 8 })
	testimony.assert_equal("1:2:3:4:5:6:7:8", result)
end)

testify:that("parse_ipv6 handles full address", function()
	local groups = types.parse_ipv6("2001:db8:0:0:0:0:0:1")
	testimony.assert_equal(8, #groups)
	testimony.assert_equal(0x2001, groups[1])
	testimony.assert_equal(0x0db8, groups[2])
	testimony.assert_equal(1, groups[8])
end)

testify:that("parse_ipv6 handles :: at start", function()
	local groups = types.parse_ipv6("::1")
	testimony.assert_equal(8, #groups)
	testimony.assert_equal(0, groups[1])
	testimony.assert_equal(1, groups[8])
end)

testify:that("parse_ipv6 handles :: at end", function()
	local groups = types.parse_ipv6("fe80::")
	testimony.assert_equal(8, #groups)
	testimony.assert_equal(0xfe80, groups[1])
	testimony.assert_equal(0, groups[8])
end)

testify:that("parse_ipv6 rejects too many groups", function()
	local groups, err = types.parse_ipv6("1:2:3:4:5:6:7:8:9")
	testimony.assert_nil(groups)
	testimony.assert_not_nil(err)
end)

testify:conclude()