package dhcpv6

import (
	"encoding/binary"
	"net/netip"
	"testing"

	"github.com/soypat/lneto/dns"
)

// writeOpt6 encodes a single DHCPv6 option into dst using the 4-byte TLV header
// (2-byte code + 2-byte length) and returns the total bytes written.
// Used by tests to build server frames without depending on the stub EncodeOption.
func writeOpt6(dst []byte, code OptCode, data ...byte) int {
	binary.BigEndian.PutUint16(dst[0:2], uint16(code))
	binary.BigEndian.PutUint16(dst[2:4], uint16(len(data)))
	copy(dst[4:], data)
	return 4 + len(data)
}

// buildServerFrame constructs a minimal DHCPv6 Advertise or Reply frame
// containing OptServerID, and an OptIANA with an embedded OptIAAddr.
func buildServerFrame(msgType MsgType, xid uint32, serverDUID []byte, iaid [4]byte, addr [16]byte) []byte {
	// IAAddr payload: addr(16) + preferred(4) + valid(4)
	iaAddrPayload := make([]byte, 24)
	copy(iaAddrPayload[:16], addr[:])
	binary.BigEndian.PutUint32(iaAddrPayload[16:20], 3600)
	binary.BigEndian.PutUint32(iaAddrPayload[20:24], 7200)

	// Encode IAAddr as an option.
	iaAddrOpt := make([]byte, 4+len(iaAddrPayload))
	writeOpt6(iaAddrOpt, OptIAAddr, iaAddrPayload...)

	// IA_NA payload: IAID(4) + T1(4) + T2(4) + IAAddr option.
	iaNAPayload := make([]byte, 12+len(iaAddrOpt))
	copy(iaNAPayload[:4], iaid[:])
	binary.BigEndian.PutUint32(iaNAPayload[4:8], 1800)
	binary.BigEndian.PutUint32(iaNAPayload[8:12], 3600)
	copy(iaNAPayload[12:], iaAddrOpt)

	buf := make([]byte, 1024)
	buf[0] = byte(msgType)
	buf[1] = byte(xid >> 16)
	buf[2] = byte(xid >> 8)
	buf[3] = byte(xid)
	n := OptionsOffset
	n += writeOpt6(buf[n:], OptServerID, serverDUID...)
	n += writeOpt6(buf[n:], OptIANA, iaNAPayload...)
	return buf[:n]
}

// TestFrameForEachOption verifies that ForEachOption correctly delivers
// the option code and data to the callback for a hand-built frame.
func TestFrameForEachOption(t *testing.T) {
	buf := make([]byte, OptionsOffset+8)
	buf[0] = byte(MsgSolicit)
	buf[3] = 42 // XID low byte

	n := writeOpt6(buf[OptionsOffset:], OptClientID, 'A', 'B')
	frm, err := NewFrame(buf[:OptionsOffset+n])
	if err != nil {
		t.Fatal(err)
	}

	var gotCode OptCode
	var gotData []byte
	err = frm.ForEachOption(func(_ int, code OptCode, data []byte) error {
		gotCode = code
		gotData = append(gotData[:0], data...)
		return nil
	})
	if err != nil {
		t.Fatal("ForEachOption:", err)
	}
	if gotCode != OptClientID {
		t.Errorf("option code: want %d (OptClientID), got %d", OptClientID, gotCode)
	}
	if string(gotData) != "AB" {
		t.Errorf("option data: want %q, got %q", "AB", gotData)
	}
}

// TestFrameValidateSize verifies that ValidateSize returns an error when an option's
// declared length extends past the end of the buffer.
func TestFrameValidateSize(t *testing.T) {
	buf := make([]byte, OptionsOffset+6)
	buf[0] = byte(MsgSolicit)
	// Option code = OptClientID, claimed length = 100, actual data = 2 bytes.
	binary.BigEndian.PutUint16(buf[OptionsOffset:], uint16(OptClientID))
	binary.BigEndian.PutUint16(buf[OptionsOffset+2:], 100)
	buf[OptionsOffset+4] = 'A'
	buf[OptionsOffset+5] = 'B'

	frm, err := NewFrame(buf)
	if err != nil {
		t.Fatal(err)
	}
	if err := frm.ValidateSize(); err == nil {
		t.Error("ValidateSize: want error for truncated option, got nil")
	}
}

// TestClientSolicitRequest exercises the full four-step DHCPv6 exchange:
// Solicit → (fabricated) Advertise → Request → (fabricated) Reply → Bound.
func TestClientSolicitRequest(t *testing.T) {
	const xid = 0x112233
	clientMAC := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	assignedAddr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
	iaid := [4]byte{clientMAC[0], clientMAC[1], clientMAC[2], clientMAC[3]}

	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: clientMAC}); err != nil {
		t.Fatal("BeginRequest:", err)
	}
	if cl.State() != StateInit {
		t.Fatalf("initial state: want StateInit, got %v", cl.State())
	}

	buf := make([]byte, 1024)

	// CLIENT: send Solicit.
	n, err := cl.Encapsulate(buf, -1, 0)
	if err != nil {
		t.Fatal("Encapsulate (Solicit):", err)
	}
	if n == 0 {
		t.Fatal("Encapsulate (Solicit): wrote 0 bytes")
	}
	if !frameHasOption(t, buf[:n], OptReconfAccept) {
		t.Fatal("Solicit missing Reconfigure Accept option")
	}
	if cl.State() != StateSoliciting {
		t.Fatalf("after Solicit: want StateSoliciting, got %v", cl.State())
	}

	// SERVER: fabricated Advertise.
	advFrame := buildServerFrame(MsgAdvertise, xid, serverDUID, iaid, assignedAddr)
	if err := cl.Demux(advFrame, 0); err != nil {
		t.Fatal("Demux (Advertise):", err)
	}
	if cl.State() != StateRequesting {
		t.Fatalf("after Advertise: want StateRequesting, got %v", cl.State())
	}

	// CLIENT: send Request.
	n, err = cl.Encapsulate(buf, -1, 0)
	if err != nil {
		t.Fatal("Encapsulate (Request):", err)
	}
	if n == 0 {
		t.Fatal("Encapsulate (Request): wrote 0 bytes")
	}

	// SERVER: fabricated Reply.
	replyFrame := buildServerFrame(MsgReply, xid, serverDUID, iaid, assignedAddr)
	if err := cl.Demux(replyFrame, 0); err != nil {
		t.Fatal("Demux (Reply):", err)
	}
	if cl.State() != StateBound {
		t.Fatalf("after Reply: want StateBound, got %v", cl.State())
	}

	addr, valid := cl.AssignedAddr()
	if !valid {
		t.Fatal("AssignedAddr: not valid after bound")
	}
	if addr != assignedAddr {
		t.Errorf("AssignedAddr: got %v, want %v", addr, assignedAddr)
	}
}

func TestClientReconfigureRenew(t *testing.T) {
	const xid = 0x112233
	clientMAC := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	assignedAddr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
	iaid := [4]byte{clientMAC[0], clientMAC[1], clientMAC[2], clientMAC[3]}
	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: clientMAC}); err != nil {
		t.Fatal(err)
	}
	var buf [1024]byte
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal(err)
	}
	if err := cl.Demux(buildServerFrame(MsgAdvertise, xid, serverDUID, iaid, assignedAddr), 0); err != nil {
		t.Fatal(err)
	}
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal(err)
	}
	if err := cl.Demux(buildServerFrame(MsgReply, xid, serverDUID, iaid, assignedAddr), 0); err != nil {
		t.Fatal(err)
	}
	if cl.State() != StateBound {
		t.Fatalf("state before reconfigure = %v, want %v", cl.State(), StateBound)
	}
	reconf := buildReconfigureFrame(0x445566, serverDUID, cl.duid, MsgRenew)
	if err := cl.Demux(reconf, 0); err != nil {
		t.Fatal(err)
	}
	if cl.State() != StateRenewing {
		t.Fatalf("state after reconfigure = %v, want %v", cl.State(), StateRenewing)
	}
	msg, ok := cl.LastReconfigure()
	if !ok || msg != MsgRenew {
		t.Fatalf("LastReconfigure = %v, %v, want %v, true", msg, ok, MsgRenew)
	}
	n, err := cl.Encapsulate(buf[:], -1, 0)
	if err != nil {
		t.Fatal(err)
	}
	frm, err := NewFrame(buf[:n])
	if err != nil {
		t.Fatal(err)
	}
	if frm.MsgType() != MsgRenew {
		t.Fatalf("Encapsulate after Reconfigure type = %v, want %v", frm.MsgType(), MsgRenew)
	}
}

func buildReconfigureFrame(xid uint32, serverDUID, clientDUID []byte, msg MsgType) []byte {
	buf := make([]byte, 128)
	buf[0] = byte(MsgReconfigure)
	buf[1] = byte(xid >> 16)
	buf[2] = byte(xid >> 8)
	buf[3] = byte(xid)
	n := OptionsOffset
	n += writeOpt6(buf[n:], OptServerID, serverDUID...)
	n += writeOpt6(buf[n:], OptClientID, clientDUID...)
	n += writeOpt6(buf[n:], OptReconfMsg, byte(msg))
	return buf[:n]
}

func frameHasOption(t testing.TB, b []byte, want OptCode) bool {
	t.Helper()
	frm, err := NewFrame(b)
	if err != nil {
		t.Fatal(err)
	}
	found := false
	if err := frm.ForEachOption(func(_ int, code OptCode, _ []byte) error {
		if code == want {
			found = true
		}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	return found
}

func TestClientParsesNTPServerOption(t *testing.T) {
	const xid = 0x102030
	clientMAC := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	assignedAddr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
	ntpAddr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x7b}
	ntpMulticast := [16]byte{0xff, 0x05, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0x01}
	iaid := [4]byte{clientMAC[0], clientMAC[1], clientMAC[2], clientMAC[3]}

	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: clientMAC}); err != nil {
		t.Fatal("BeginRequest:", err)
	}
	var buf [1024]byte
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Solicit):", err)
	}
	advFrame := buildServerFrame(MsgAdvertise, xid, serverDUID, iaid, assignedAddr)
	if err := cl.Demux(advFrame, 0); err != nil {
		t.Fatal("Demux (Advertise):", err)
	}
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Request):", err)
	}

	replyFrame := appendNTPServerOption(t, buildServerFrame(MsgReply, xid, serverDUID, iaid, assignedAddr), ntpAddr, ntpMulticast, "time.example.com")
	if err := cl.Demux(replyFrame, 0); err != nil {
		t.Fatal("Demux (Reply):", err)
	}
	var ntps []netip.Addr
	ntps = cl.AppendNTPServers(ntps)
	if len(ntps) != 1 || ntps[0] != netip.AddrFrom16(ntpAddr) {
		t.Fatalf("AppendNTPServers = %v, want %v", ntps, netip.AddrFrom16(ntpAddr))
	}
	if n := cl.NumNTPServers(); n != 1 {
		t.Fatalf("NumNTPServers = %d, want 1", n)
	}
	var multicasts []netip.Addr
	multicasts = cl.AppendNTPMulticastServers(multicasts)
	if len(multicasts) != 1 || multicasts[0] != netip.AddrFrom16(ntpMulticast) {
		t.Fatalf("AppendNTPMulticastServers = %v, want %v", multicasts, netip.AddrFrom16(ntpMulticast))
	}
	if n := cl.NumNTPMulticastServers(); n != 1 {
		t.Fatalf("NumNTPMulticastServers = %d, want 1", n)
	}
	var names []dns.Name
	names = cl.AppendNTPServerNames(names)
	if len(names) != 1 || !names[0].EqualString("time.example.com") {
		t.Fatalf("AppendNTPServerNames = %v, want time.example.com", names)
	}
	if n := cl.NumNTPServerNames(); n != 1 {
		t.Fatalf("NumNTPServerNames = %d, want 1", n)
	}
}

func TestClientParsesDomainSearchOption(t *testing.T) {
	const xid = 0x203040
	clientMAC := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	assignedAddr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
	iaid := [4]byte{clientMAC[0], clientMAC[1], clientMAC[2], clientMAC[3]}

	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: clientMAC}); err != nil {
		t.Fatal("BeginRequest:", err)
	}
	var buf [1024]byte
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Solicit):", err)
	}
	advFrame := buildServerFrame(MsgAdvertise, xid, serverDUID, iaid, assignedAddr)
	if err := cl.Demux(advFrame, 0); err != nil {
		t.Fatal("Demux (Advertise):", err)
	}
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Request):", err)
	}

	replyFrame := appendDomainSearchOption(t, buildServerFrame(MsgReply, xid, serverDUID, iaid, assignedAddr), "example.com", "corp.example.com")
	if err := cl.Demux(replyFrame, 0); err != nil {
		t.Fatal("Demux (Reply):", err)
	}
	var domains []dns.Name
	domains = cl.AppendDomainSearch(domains)
	if len(domains) != 2 {
		t.Fatalf("AppendDomainSearch len = %d, want 2", len(domains))
	}
	if !domains[0].EqualString("example.com") || !domains[1].EqualString("corp.example.com") {
		t.Fatalf("AppendDomainSearch = %q, %q", domains[0].String(), domains[1].String())
	}
	if n := cl.NumDomainSearch(); n != 2 {
		t.Fatalf("NumDomainSearch = %d, want 2", n)
	}
}

func TestClientParsesDelegatedPrefix(t *testing.T) {
	const xid = 0x304050
	clientMAC := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	assignedAddr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
	delegated := netip.MustParsePrefix("2001:db8:abcd::/48")
	iaid := [4]byte{clientMAC[0], clientMAC[1], clientMAC[2], clientMAC[3]}

	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: clientMAC}); err != nil {
		t.Fatal("BeginRequest:", err)
	}
	var buf [1024]byte
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Solicit):", err)
	}
	advFrame := buildServerFrame(MsgAdvertise, xid, serverDUID, iaid, assignedAddr)
	if err := cl.Demux(advFrame, 0); err != nil {
		t.Fatal("Demux (Advertise):", err)
	}
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Request):", err)
	}

	replyFrame := appendIAPDOption(buildServerFrame(MsgReply, xid, serverDUID, iaid, assignedAddr), iaid, delegated)
	if err := cl.Demux(replyFrame, 0); err != nil {
		t.Fatal("Demux (Reply):", err)
	}
	var prefixes []DelegatedPrefix
	prefixes = cl.AppendDelegatedPrefixes(prefixes)
	if len(prefixes) != 1 {
		t.Fatalf("AppendDelegatedPrefixes len = %d, want 1", len(prefixes))
	}
	if prefixes[0].Prefix != delegated || prefixes[0].PreferredLifetime != 1800 || prefixes[0].ValidLifetime != 3600 {
		t.Fatalf("delegated prefix = %+v, want %s preferred=1800 valid=3600", prefixes[0], delegated)
	}
	if n := cl.NumDelegatedPrefixes(); n != 1 {
		t.Fatalf("NumDelegatedPrefixes = %d, want 1", n)
	}
	if cl.PrefixDelegationRenewalSeconds() != 900 || cl.PrefixDelegationRebindingSeconds() != 1800 {
		t.Fatalf("IA_PD timers = renewal:%d rebind:%d, want 900/1800", cl.PrefixDelegationRenewalSeconds(), cl.PrefixDelegationRebindingSeconds())
	}
}

func appendNTPServerOption(t testing.TB, frame []byte, addr, multicast [16]byte, fqdn string) []byte {
	t.Helper()
	var ntpPayload []byte
	ntpPayload = appendNTPServerAddrSuboption(ntpPayload, 1, addr)
	ntpPayload = appendNTPServerAddrSuboption(ntpPayload, 2, multicast)
	name, err := dns.NewName(fqdn)
	if err != nil {
		t.Fatal(err)
	}
	nameStart := len(ntpPayload) + 4
	ntpPayload = append(ntpPayload, 0, 3, 0, 0)
	ntpPayload, err = name.AppendTo(ntpPayload)
	if err != nil {
		t.Fatal(err)
	}
	binary.BigEndian.PutUint16(ntpPayload[nameStart-2:nameStart], uint16(len(ntpPayload)-nameStart))
	buf := append(frame[:len(frame):len(frame)], make([]byte, 4+len(ntpPayload))...)
	writeOpt6(buf[len(frame):], OptNTPServer, ntpPayload...)
	return buf
}

func appendNTPServerAddrSuboption(dst []byte, code uint16, addr [16]byte) []byte {
	start := len(dst)
	dst = append(dst, 0, 0, 0, 16)
	binary.BigEndian.PutUint16(dst[start:start+2], code)
	return append(dst, addr[:]...)
}

func appendDomainSearchOption(t testing.TB, frame []byte, domains ...string) []byte {
	t.Helper()
	var payload []byte
	for _, domain := range domains {
		name, err := dns.NewName(domain)
		if err != nil {
			t.Fatal(err)
		}
		payload, err = name.AppendTo(payload)
		if err != nil {
			t.Fatal(err)
		}
	}
	buf := append(frame[:len(frame):len(frame)], make([]byte, 4+len(payload))...)
	writeOpt6(buf[len(frame):], OptDomainList, payload...)
	return buf
}

func appendIAPDOption(frame []byte, iaid [4]byte, prefix netip.Prefix) []byte {
	var iaPrefix [29]byte
	binary.BigEndian.PutUint16(iaPrefix[0:2], uint16(OptIAPrefix))
	binary.BigEndian.PutUint16(iaPrefix[2:4], 25)
	binary.BigEndian.PutUint32(iaPrefix[4:8], 1800)
	binary.BigEndian.PutUint32(iaPrefix[8:12], 3600)
	iaPrefix[12] = byte(prefix.Bits())
	prefixAddr := prefix.Addr().As16()
	copy(iaPrefix[13:], prefixAddr[:])
	var iapd [41]byte
	copy(iapd[:4], iaid[:])
	binary.BigEndian.PutUint32(iapd[4:8], 900)
	binary.BigEndian.PutUint32(iapd[8:12], 1800)
	copy(iapd[12:], iaPrefix[:])
	buf := append(frame[:len(frame):len(frame)], make([]byte, 4+len(iapd))...)
	writeOpt6(buf[len(frame):], OptIAPD, iapd[:]...)
	return buf
}

// TestClientEncapsulateSolicit verifies that the first Encapsulate call
// writes a Solicit message with at least OptClientID and OptIANA.
func TestClientEncapsulateSolicit(t *testing.T) {
	var cl Client
	if err := cl.BeginRequest(0xABCDEF, RequestConfig{
		ClientHardwareAddr: [6]byte{1, 2, 3, 4, 5, 6},
	}); err != nil {
		t.Fatal(err)
	}

	buf := make([]byte, 512)
	n, err := cl.Encapsulate(buf, -1, 0)
	if err != nil {
		t.Fatal(err)
	}
	if n == 0 {
		t.Fatal("Encapsulate: wrote 0 bytes")
	}

	frm, err := NewFrame(buf[:n])
	if err != nil {
		t.Fatal(err)
	}
	if frm.MsgType() != MsgSolicit {
		t.Errorf("MsgType: want MsgSolicit, got %v", frm.MsgType())
	}
	if frm.TransactionID() != 0xABCDEF {
		t.Errorf("TransactionID: want 0xABCDEF, got 0x%X", frm.TransactionID())
	}

	var hasClientID, hasIANA bool
	err = frm.ForEachOption(func(_ int, code OptCode, _ []byte) error {
		switch code {
		case OptClientID:
			hasClientID = true
		case OptIANA:
			hasIANA = true
		}
		return nil
	})
	if err != nil {
		t.Fatal("ForEachOption:", err)
	}
	if !hasClientID {
		t.Error("Solicit: missing OptClientID")
	}
	if !hasIANA {
		t.Error("Solicit: missing OptIANA")
	}
}

// TestClientDoubleTapEncapsulate verifies that calling Encapsulate twice in the
// same state returns 0 bytes on the second call (idempotent, no duplicate messages).
func TestClientDoubleTapEncapsulate(t *testing.T) {
	var cl Client
	if err := cl.BeginRequest(1, RequestConfig{
		ClientHardwareAddr: [6]byte{1, 2, 3, 4, 5, 6},
	}); err != nil {
		t.Fatal(err)
	}

	buf := make([]byte, 512)
	n, err := cl.Encapsulate(buf, -1, 0)
	if err != nil {
		t.Fatal("first Encapsulate:", err)
	}
	if n == 0 {
		t.Fatal("first Encapsulate: wrote 0 bytes")
	}

	n2, err := cl.Encapsulate(buf, -1, 0)
	if err != nil {
		t.Fatal("second Encapsulate:", err)
	}
	if n2 != 0 {
		t.Errorf("second Encapsulate: want 0 bytes (idempotent), got %d", n2)
	}
}

// driveToRequesting advances a fresh client through Solicit→Advertise→Request so
// the next Demux of a Reply runs the option-parsing path.
func driveToRequesting(t testing.TB, cl *Client, xid uint32, serverDUID []byte, iaid [4]byte, addr [16]byte) {
	t.Helper()
	var buf [1024]byte
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Solicit):", err)
	}
	if err := cl.Demux(buildServerFrame(MsgAdvertise, xid, serverDUID, iaid, addr), 0); err != nil {
		t.Fatal("Demux (Advertise):", err)
	}
	if _, err := cl.Encapsulate(buf[:], -1, 0); err != nil {
		t.Fatal("Encapsulate (Request):", err)
	}
}

// floodDNSServersOption appends an OptDNSServers carrying n distinct addresses.
func floodDNSServersOption(frame []byte, n int) []byte {
	payload := make([]byte, 0, n*16)
	for i := range n {
		var a [16]byte
		a[0], a[1], a[15] = 0x20, 0x01, byte(i+1)
		payload = append(payload, a[:]...)
	}
	buf := append(frame[:len(frame):len(frame)], make([]byte, 4+len(payload))...)
	writeOpt6(buf[len(frame):], OptDNSServers, payload...)
	return buf
}

// floodIAPDOption appends an OptIAPD carrying n OptIAPrefix sub-options.
func floodIAPDOption(frame []byte, iaid [4]byte, n int) []byte {
	payload := make([]byte, 12)
	copy(payload[:4], iaid[:])
	binary.BigEndian.PutUint32(payload[4:8], 900)
	binary.BigEndian.PutUint32(payload[8:12], 1800)
	for i := range n {
		var iaPrefix [29]byte
		binary.BigEndian.PutUint16(iaPrefix[0:2], uint16(OptIAPrefix))
		binary.BigEndian.PutUint16(iaPrefix[2:4], 25)
		binary.BigEndian.PutUint32(iaPrefix[4:8], 1800)
		binary.BigEndian.PutUint32(iaPrefix[8:12], 3600)
		iaPrefix[12] = 48
		iaPrefix[13], iaPrefix[14], iaPrefix[15] = 0x20, 0x01, byte(i+1)
		payload = append(payload, iaPrefix[:]...)
	}
	buf := append(frame[:len(frame):len(frame)], make([]byte, 4+len(payload))...)
	writeOpt6(buf[len(frame):], OptIAPD, payload...)
	return buf
}

// floodNTPOption appends an OptNTPServer carrying nAddr unicast (suboption 1),
// nMulticast (suboption 2) and nNames FQDN (suboption 3) entries.
func floodNTPOption(t testing.TB, frame []byte, nAddr, nMulticast, nNames int) []byte {
	t.Helper()
	var payload []byte
	for i := range nAddr {
		var a [16]byte
		a[0], a[15] = 0x20, byte(i+1)
		payload = appendNTPServerAddrSuboption(payload, 1, a)
	}
	for i := range nMulticast {
		var a [16]byte
		a[0], a[1], a[15] = 0xff, 0x05, byte(i+1)
		payload = appendNTPServerAddrSuboption(payload, 2, a)
	}
	name, err := dns.NewName("ntp.example.com")
	if err != nil {
		t.Fatal(err)
	}
	for range nNames {
		start := len(payload) + 4
		payload = append(payload, 0, 3, 0, 0)
		payload, err = name.AppendTo(payload)
		if err != nil {
			t.Fatal(err)
		}
		binary.BigEndian.PutUint16(payload[start-2:start], uint16(len(payload)-start))
	}
	buf := append(frame[:len(frame):len(frame)], make([]byte, 4+len(payload))...)
	writeOpt6(buf[len(frame):], OptNTPServer, payload...)
	return buf
}

func repeatStrings(s string, n int) []string {
	out := make([]string, n)
	for i := range out {
		out[i] = s
	}
	return out
}

// TestClientLimitsCapServerData verifies that the client stores no more than
// the configured number of each repeated option, so a server flooding a Reply
// cannot drive unbounded allocation (the OOM concern raised in the PR review).
func TestClientLimitsCapServerData(t *testing.T) {
	const xid = 0x515151
	mac := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	iaid := [4]byte{mac[0], mac[1], mac[2], mac[3]}
	addr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
	limits := Limits{
		MaxDNSServers: 2, MaxDomainSearch: 2, MaxNTPServers: 2,
		MaxNTPMulticastServers: 1, MaxNTPServerNames: 2, MaxDelegatedPrefixes: 2,
	}

	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: mac, Limits: limits}); err != nil {
		t.Fatal("BeginRequest:", err)
	}
	driveToRequesting(t, &cl, xid, serverDUID, iaid, addr)

	// A Reply far exceeding every configured cap.
	reply := buildServerFrame(MsgReply, xid, serverDUID, iaid, addr)
	reply = floodDNSServersOption(reply, 8)
	reply = floodIAPDOption(reply, iaid, 8)
	reply = appendDomainSearchOption(t, reply, repeatStrings("example.com", 8)...)
	reply = floodNTPOption(t, reply, 8, 8, 8)
	if err := cl.Demux(reply, 0); err != nil {
		t.Fatal("Demux (Reply):", err)
	}

	for _, tc := range []struct {
		name string
		got  int
		want int
	}{
		{"DNS servers", cl.NumDNSServers(), limits.MaxDNSServers},
		{"domain search", cl.NumDomainSearch(), limits.MaxDomainSearch},
		{"NTP servers", cl.NumNTPServers(), limits.MaxNTPServers},
		{"NTP multicast", cl.NumNTPMulticastServers(), limits.MaxNTPMulticastServers},
		{"NTP names", cl.NumNTPServerNames(), limits.MaxNTPServerNames},
		{"delegated prefixes", cl.NumDelegatedPrefixes(), limits.MaxDelegatedPrefixes},
	} {
		if tc.got != tc.want {
			t.Errorf("%s stored = %d, want capped at %d", tc.name, tc.got, tc.want)
		}
	}
}

// TestClientParseNoAllocs verifies that once BeginRequest has sized the option
// backing arrays, parsing a server message performs no further allocation: the
// bounded slices are reused, keeping the network-facing path off the heap.
func TestClientParseNoAllocs(t *testing.T) {
	const xid = 0x123456
	mac := [6]byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}
	serverDUID := []byte{0, 3, 0, 1, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66}
	iaid := [4]byte{mac[0], mac[1], mac[2], mac[3]}
	addr := [16]byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}

	var cl Client
	if err := cl.BeginRequest(xid, RequestConfig{ClientHardwareAddr: mac}); err != nil {
		t.Fatal("BeginRequest:", err)
	}
	reply := buildServerFrame(MsgReply, xid, serverDUID, iaid, addr)
	reply = floodDNSServersOption(reply, cl.limits.MaxDNSServers)
	reply = floodIAPDOption(reply, iaid, cl.limits.MaxDelegatedPrefixes)
	reply = appendDomainSearchOption(t, reply, repeatStrings("example.com", cl.limits.MaxDomainSearch)...)
	reply = floodNTPOption(t, reply, cl.limits.MaxNTPServers, cl.limits.MaxNTPMulticastServers, cl.limits.MaxNTPServerNames)
	frm, err := NewFrame(reply)
	if err != nil {
		t.Fatal("NewFrame:", err)
	}

	clearStores := func() {
		cl.serverDUID = cl.serverDUID[:0]
		cl.dns = cl.dns[:0]
		cl.domainSearch = cl.domainSearch[:0]
		cl.ntps = cl.ntps[:0]
		cl.ntpMulticast = cl.ntpMulticast[:0]
		cl.ntpNames = cl.ntpNames[:0]
		cl.delegatedPrefixes = cl.delegatedPrefixes[:0]
	}
	clearStores()
	if err := cl.setOptions(frm); err != nil { // warm up the dns.Name backing arrays.
		t.Fatal("setOptions:", err)
	}
	allocs := testing.AllocsPerRun(50, func() {
		clearStores()
		_ = cl.setOptions(frm)
	})
	if allocs != 0 {
		t.Errorf("setOptions must not allocate after BeginRequest sizing, got %v allocs/op", allocs)
	}
}