package server

import (
	"net"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
)

// readTick mirrors the unexported read deadline serveUDP uses. The test has to
// know it to be able to sleep past it; a copy here means a change to the
// product constant would quietly weaken this test rather than break it, so the
// value is asserted against the listener's behaviour below.
const readTick = 500 * time.Millisecond

// TestUDPReadDeadlineIsResetEveryIteration is the regression test for the
// worst bug this package has had (AGENTS.md): serveUDP sets a read deadline so
// the loop notices cancellation, and a deadline set ONCE expires for good.
// After that the loop spins — every ReadFrom returns "i/o timeout" instantly —
// and never reads another packet for the life of the process. The listener
// looks healthy: it is bound, it is in the admin API, it logs nothing.
//
// The invariant pinned here: a packet that arrives more than one read-tick
// after the previous one is still answered. Send, sleep past the tick, send
// again, and both must come back.
func TestUDPReadDeadlineIsResetEveryIteration(t *testing.T) {
	h := newHarness(t, nil)
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer pc.Close()
	addr, err := net.ResolveUDPAddr("udp", h.addr("udp"))
	if err != nil {
		t.Fatalf("resolve: %v", err)
	}

	ask := func(id uint16) {
		t.Helper()
		q := question("deadline.test.", dns.TypeA)
		q.Id = id
		wire, err := q.Pack()
		if err != nil {
			t.Fatalf("pack: %v", err)
		}
		if _, err := pc.WriteTo(wire, addr); err != nil {
			t.Fatalf("write: %v", err)
		}
		_ = pc.SetReadDeadline(time.Now().Add(3 * time.Second))
		buf := make([]byte, 4096)
		n, _, err := pc.ReadFrom(buf)
		if err != nil {
			t.Fatalf("no answer to query %d: %v", id, err)
		}
		res := new(dns.Msg)
		if err := res.Unpack(buf[:n]); err != nil {
			t.Fatalf("unpack: %v", err)
		}
		if res.Id != id {
			t.Errorf("reply id = %d, want %d", res.Id, id)
		}
		if a := answerAddr(res); a != "203.0.113.9" {
			t.Errorf("reply A = %q, want 203.0.113.9", a)
		}
	}

	ask(0x3001)
	// Longer than the read deadline, so a deadline that is not reset has
	// certainly expired by the time the second packet lands.
	time.Sleep(readTick + 250*time.Millisecond)
	ask(0x3002)

	waitFor(t, 2*time.Second, "both queries at the engine", func() bool {
		return len(h.fake.seen()) == 2
	})
}

// TestUDPTruncatesAnOversizedAnswer: a multi-kilobyte answer to a client that
// advertised nothing (the 512-byte floor) must come back as a question-only
// TC reply. Sending it whole would be silently dropped by the client and read
// as a dropped query, which is exactly the bug the engine suite found once
// already (STATUS.md: "the uncached path never truncated").
func TestUDPTruncatesAnOversizedAnswer(t *testing.T) {
	h := newHarness(t, nil)
	res := h.askUDP(question("big.test.", dns.TypeTXT))
	if !res.Truncated {
		t.Error("TC not set on a 3 KiB answer to a 512-byte client")
	}
	if len(res.Answer) != 0 {
		t.Errorf("truncated reply carries %d answers, want none", len(res.Answer))
	}
	if len(res.Question) != 1 || res.Question[0].Name != "big.test." {
		t.Errorf("truncated reply question = %+v, want big.test.", res.Question)
	}
	if got := h.srv.load("udp", "truncated"); got != 1 {
		t.Errorf("udp.truncated = %d, want 1", got)
	}
	if err := h.fake.lastPackErr(); err != nil {
		t.Fatalf("the fake's big answer did not pack, so this test is measuring a fixture failure: %v", err)
	}
}

// TestUDPUsesTheAdvertisedEDNSSize: with EDNS0 the client says how much it can
// take, and the server must honour it — the same answer, delivered whole.
func TestUDPUsesTheAdvertisedEDNSSize(t *testing.T) {
	h := newHarness(t, nil)
	m := question("big.test.", dns.TypeTXT)
	m.SetEdns0(4096, false)
	res := h.askUDP(m)
	if res.Truncated {
		t.Error("3 KiB answer truncated for a client that advertised 4096")
	}
	if len(res.Answer) != 1 {
		t.Fatalf("reply has %d answers, want the whole TXT", len(res.Answer))
	}
	if o := res.IsEdns0(); o == nil {
		t.Error("reply has no OPT even though the query had one")
	}
}

// TestUDPCapsTheAdvertisedEDNS0Buffer is the amplification guard: a client may
// claim 64 KiB, and the server rewrites the request to 4096 before the engine
// ever sees it, so no answer is built for a buffer nobody will accept.
func TestUDPCapsTheAdvertisedEDNS0Buffer(t *testing.T) {
	h := newHarness(t, nil)
	m := question("edns.test.", dns.TypeA)
	m.SetEdns0(65535, false)
	h.askUDP(m)

	got := h.fake.only(t)
	if got.ednsSize == 0 {
		t.Fatal("the engine saw no EDNS0 at all")
	}
	if got.ednsSize > 4096 {
		t.Errorf("engine saw an advertised buffer of %d, want it capped at 4096", got.ednsSize)
	}
	if got.ednsSize != 4096 {
		t.Errorf("advertised buffer = %d, want the 4096 cap", got.ednsSize)
	}
}

// TestFitUDP covers the packing helper directly, where the buffer arithmetic is
// visible without a socket in the way. The three cases are the three answers
// the function can give: pass through, truncate to the question, or refuse to
// answer at all.
func TestFitUDP(t *testing.T) {
	t.Run("small answer passes through untouched", func(t *testing.T) {
		s := &Server{}
		q := question("small.test.", dns.TypeA)
		res := answerA(newReply(q))
		out := s.fitUDP(res, q, 512)
		want, err := res.Pack()
		if err != nil {
			t.Fatalf("pack: %v", err)
		}
		if string(out) != string(want) {
			t.Errorf("fitUDP rewrote an answer that already fit (%d vs %d bytes)", len(out), len(want))
		}
		if got := s.load("udp", "truncated"); got != 0 {
			t.Errorf("truncated = %d for an answer that fit", got)
		}
	})

	t.Run("oversized answer becomes a question-only TC reply", func(t *testing.T) {
		s := &Server{}
		q := question("big.test.", dns.TypeTXT)
		res := bigAnswer(newReply(q))
		full, err := res.Pack()
		if err != nil {
			t.Fatalf("the big answer must pack for this test to mean anything: %v", err)
		}
		if len(full) <= 512 {
			t.Fatalf("fixture is only %d bytes, it cannot test truncation", len(full))
		}
		out := s.fitUDP(res, q, 512)
		if len(out) > 512 {
			t.Errorf("truncated reply is %d bytes, over the 512 floor", len(out))
		}
		got := new(dns.Msg)
		if err := got.Unpack(out); err != nil {
			t.Fatalf("unpack truncated reply: %v", err)
		}
		if !got.Truncated {
			t.Error("TC not set")
		}
		if len(got.Answer) != 0 {
			t.Errorf("truncated reply carries %d answers", len(got.Answer))
		}
		if len(got.Question) != 1 || got.Question[0].Name != q.Question[0].Name {
			t.Errorf("question = %+v, want %v", got.Question, q.Question)
		}
		if got.Rcode != res.Rcode {
			t.Errorf("rcode = %d, want the answer's %d", got.Rcode, res.Rcode)
		}
		if s.load("udp", "truncated") != 1 {
			t.Errorf("truncated counter = %d, want 1", s.load("udp", "truncated"))
		}
	})

	t.Run("an advertised EDNS0 buffer is the limit", func(t *testing.T) {
		s := &Server{}
		q := question("big.test.", dns.TypeTXT)
		q.SetEdns0(4096, false)
		res := bigAnswer(newReply(q))
		out := s.fitUDP(res, q, 512)
		got := new(dns.Msg)
		if err := got.Unpack(out); err != nil {
			t.Fatalf("unpack: %v", err)
		}
		if got.Truncated {
			t.Error("3 KiB answer truncated against a 4096-byte buffer")
		}
		if len(got.Answer) != 1 {
			t.Errorf("answers = %d, want 1", len(got.Answer))
		}
		if o := got.IsEdns0(); o == nil || o.UDPSize() != 4096 {
			t.Errorf("truncated reply lost the OPT size: %+v", o)
		}
	})

	t.Run("an unadvertised buffer below 512 cannot raise the floor", func(t *testing.T) {
		s := &Server{}
		q := question("big.test.", dns.TypeTXT)
		// 300 is below the 512 floor RFC 1035 guarantees; using it as the
		// limit would produce a reply no client is obliged to read.
		q.SetEdns0(300, false)
		res := bigAnswer(newReply(q))
		out := s.fitUDP(res, q, 512)
		if len(out) > 512 {
			t.Errorf("reply is %d bytes, over the 512 floor", len(out))
		}
		got := new(dns.Msg)
		if err := got.Unpack(out); err != nil {
			t.Fatalf("unpack: %v", err)
		}
		if !got.Truncated {
			t.Error("TC not set")
		}
	})

	t.Run("an answer that will not pack is dropped, not truncated", func(t *testing.T) {
		s := &Server{}
		q := question("unpackable.test.", dns.TypeTXT)
		res := new(dns.Msg)
		res.SetReply(q)
		// A character-string over 255 bytes: miekg/dns v1.1.73 reports its
		// length as zero, so Pack fails. Dropping is the only safe answer —
		// a half-written datagram is worse than none.
		res.Answer = []dns.RR{&dns.TXT{
			Hdr: dns.RR_Header{Name: q.Question[0].Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 60},
			Txt: []string{strings.Repeat("x", 300)},
		}}
		if _, err := res.Pack(); err == nil {
			t.Fatal("fixture packs; this case is no longer testing the pack failure")
		}
		if out := s.fitUDP(res, q, 512); out != nil {
			t.Errorf("fitUDP returned %d bytes for an unpackable answer, want nil", len(out))
		}
	})
}

// TestUDPDropSendsNothing: a dropped query is answered with silence, which on
// UDP means no datagram at all. The counter is the only evidence a client-side
// timeout was a policy decision.
func TestUDPDropSendsNothing(t *testing.T) {
	h := newHarness(t, nil)
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer pc.Close()
	addr, err := net.ResolveUDPAddr("udp", h.addr("udp"))
	if err != nil {
		t.Fatalf("resolve: %v", err)
	}
	wire, err := question("drop.test.", dns.TypeA).Pack()
	if err != nil {
		t.Fatalf("pack: %v", err)
	}
	if _, err := pc.WriteTo(wire, addr); err != nil {
		t.Fatalf("write: %v", err)
	}
	_ = pc.SetReadDeadline(time.Now().Add(300 * time.Millisecond))
	buf := make([]byte, 4096)
	if n, _, err := pc.ReadFrom(buf); err == nil {
		t.Errorf("a dropped query was answered with %d bytes", n)
	}
	waitFor(t, 2*time.Second, "the drop at the engine", func() bool { return len(h.fake.seen()) == 1 })
	if got := h.srv.load("udp", "dropped"); got != 1 {
		t.Errorf("udp.dropped = %d, want 1", got)
	}
	if got := h.srv.load("udp", "queries"); got != 1 {
		t.Errorf("udp.queries = %d, want 1", got)
	}
}

// TestUDPMalformedPacketIsCountedNotAnswered: an unparsable datagram has no ID
// and no question, so there is nothing legal to reply with. It must be counted
// and dropped — a server that echoes a FORMERR to arbitrary garbage is an
// amplifier.
func TestUDPMalformedPacketIsCountedNotAnswered(t *testing.T) {
	h := newHarness(t, nil)
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer pc.Close()
	addr, err := net.ResolveUDPAddr("udp", h.addr("udp"))
	if err != nil {
		t.Fatalf("resolve: %v", err)
	}
	if _, err := pc.WriteTo([]byte{0x01, 0x02, 0x03}, addr); err != nil {
		t.Fatalf("write: %v", err)
	}
	waitFor(t, 2*time.Second, "the parse failure", func() bool { return h.srv.load("udp", "errors") == 1 })
	if n := len(h.fake.seen()); n != 0 {
		t.Errorf("garbage reached the engine %d times", n)
	}
	_ = pc.SetReadDeadline(time.Now().Add(300 * time.Millisecond))
	buf := make([]byte, 4096)
	if n, _, err := pc.ReadFrom(buf); err == nil {
		t.Errorf("garbage was answered with %d bytes", n)
	}
}

// TestUDPQuestionlessMessageReachesTheEngine: a parseable message with no
// question is answerable with FORMERR, and servePacket deliberately hands it
// to the engine rather than dropping it — the comment in the code says so and
// it is the difference between a client that learns why it was refused and a
// client that times out.
func TestUDPQuestionlessMessageReachesTheEngine(t *testing.T) {
	h := newHarness(t, nil)
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer pc.Close()
	addr, err := net.ResolveUDPAddr("udp", h.addr("udp"))
	if err != nil {
		t.Fatalf("resolve: %v", err)
	}
	// Header only: valid DNS, qdcount 0.
	msg := []byte{0x12, 0x34, 0x01, 0x00, 0, 0, 0, 0, 0, 0, 0, 0}
	if _, err := pc.WriteTo(msg, addr); err != nil {
		t.Fatalf("write: %v", err)
	}
	_ = pc.SetReadDeadline(time.Now().Add(3 * time.Second))
	buf := make([]byte, 4096)
	n, _, err := pc.ReadFrom(buf)
	if err != nil {
		t.Fatalf("no answer to a questionless query: %v", err)
	}
	res := new(dns.Msg)
	if err := res.Unpack(buf[:n]); err != nil {
		t.Fatalf("unpack: %v", err)
	}
	if res.Rcode != dns.RcodeFormatError {
		t.Errorf("rcode = %s, want FORMERR", dns.RcodeToString[res.Rcode])
	}
	if got := h.fake.only(t); got.name != "" {
		t.Errorf("engine saw question %q, want none", got.name)
	}
}