package server

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

// The counters are what the console's status page and the /status endpoint
// report, so they are the only account an operator has of what the server is
// doing. Everything in this file is about the numbers being right rather than
// merely present: an under-counted drop reads as a working resolver, and a
// double-counted query reads as a busy one.
//
// The harness's fake handler selects its outcome from the question name
// (block., cached., fail., drop., big.), so a test drives the exact counter
// move it wants by asking for it by name. No counter is bumped by hand.

var statTransports = []string{"udp", "tcp", "dot", "doh"}

// silenceWait is how long a transport is given to produce the answer the fake
// intends to produce. A drop is answered with nothing at all, so over UDP this
// is the price of asking about one.
const silenceWait = 300 * time.Millisecond

// exchangeOnce sends m over transport and returns the reply, or nil when the
// transport answered with silence or hung up. h.ask cannot be used here: it
// treats a missing answer as a test failure, and "no answer" is a legitimate
// outcome the fake produces on purpose.
func exchangeOnce(t *testing.T, h *harness, transport string, m *dns.Msg) *dns.Msg {
	t.Helper()
	switch transport {
	case "udp":
		return exchangeUDP(t, h, m)
	case "tcp":
		return exchangeStream(t, h.addr("tcp"), nil, m)
	case "dot":
		tc := h.clientTLS()
		tc.NextProtos = []string{"dot"}
		return exchangeStream(t, h.addr("dot"), tc, m)
	case "doh":
		res, body := h.doHPOST(m)
		if res.StatusCode != http.StatusOK {
			t.Fatalf("doh POST = %d, want 200: %s", res.StatusCode, body)
		}
		return unpackDoH(t, body)
	}
	t.Fatalf("unknown transport %q", transport)
	return nil
}

func exchangeUDP(t *testing.T, h *harness, m *dns.Msg) *dns.Msg {
	t.Helper()
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { _ = pc.Close() })
	addr, err := net.ResolveUDPAddr("udp", h.addr("udp"))
	if err != nil {
		t.Fatalf("resolve: %v", err)
	}
	wire, err := m.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(silenceWait))
	buf := make([]byte, 4096)
	n, _, err := pc.ReadFrom(buf)
	if err != nil {
		return nil // silence, or a datagram that never came
	}
	res := new(dns.Msg)
	if err := res.Unpack(buf[:n]); err != nil {
		t.Fatalf("unpack reply: %v", err)
	}
	return res
}

func exchangeStream(t *testing.T, addr string, tc *tls.Config, m *dns.Msg) *dns.Msg {
	t.Helper()
	d := &net.Dialer{Timeout: 2 * time.Second}
	var (
		conn net.Conn
		err  error
	)
	if tc != nil {
		conn, err = tls.DialWithDialer(d, "tcp", addr, tc)
	} else {
		conn, err = d.Dial("tcp", addr)
	}
	if err != nil {
		t.Fatalf("dial %s: %v", addr, err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	_ = conn.SetDeadline(time.Now().Add(3 * time.Second))
	co := &dns.Conn{Conn: conn}
	if err := co.WriteMsg(m); err != nil {
		t.Fatalf("write: %v", err)
	}
	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	res, err := co.ReadMsg()
	if err != nil {
		return nil // a drop closes the connection instead of answering
	}
	return res
}

// counterDelta is the change one query is expected to make to a transport's
// counter set. Every field is a delta, never a total: a test that asserts a
// total is asserting the order the tests happened to run in.
type counterDelta struct {
	queries   uint64
	blocked   uint64
	cached    uint64
	errors    uint64
	dropped   uint64
	truncated uint64
	sent      uint64
}

// applyDelta is the expected counter set after the delta, used to wait for the
// server to settle on it and to report where it differs.
func applyDelta(base TransportStats, d counterDelta) TransportStats {
	return TransportStats{
		Transport: base.Transport,
		Queries:   base.Queries + d.queries,
		Blocked:   base.Blocked + d.blocked,
		Cached:    base.Cached + d.cached,
		Errors:    base.Errors + d.errors,
		Dropped:   base.Dropped + d.dropped,
		Truncated: base.Truncated + d.truncated,
		Sent:      base.Sent + d.sent,
	}
}

// TestCountersFollowTheEngineOutcome is the main counter contract: one query,
// one answer, and exactly the counters that outcome implies — no more, no
// fewer. The full counter set is compared, not just the counter under test,
// because a counter that moves on every path is exactly the bug a
// single-counter assertion would miss.
func TestCountersFollowTheEngineOutcome(t *testing.T) {
	cases := []struct {
		name string
		// prefix is the first label of the question name; the fake's recipe
		// keys the outcome off it.
		prefix string
		qtype  uint16
		// deltas is the expected move per transport.
		deltas map[string]counterDelta
		why    string
	}{
		{
			prefix: "plain", qtype: dns.TypeA,
			deltas: map[string]counterDelta{
				"udp": {queries: 1, sent: 1},
				"tcp": {queries: 1, sent: 1},
				"dot": {queries: 1, sent: 1},
				// The DoH path writes the body through net/http and never
				// counts a send. Pinned exactly (sent stays 0) so the
				// transport tab's "sent" column reads 0 next to a non-zero
				// "queries" for DoH today, and so adding the bump later has
				// to come with a decision about what it means.
				"doh": {queries: 1},
			},
			why: "an ordinary answer: accepted and written back",
		},
		{
			prefix: "block", qtype: dns.TypeA,
			deltas: map[string]counterDelta{
				"udp": {queries: 1, blocked: 1, sent: 1},
				"tcp": {queries: 1, blocked: 1, sent: 1},
				"dot": {queries: 1, blocked: 1, sent: 1},
				"doh": {queries: 1, blocked: 1},
			},
			why: "a blocklist hit is a normal answer to the client, so it counts as sent too",
		},
		{
			prefix: "cached", qtype: dns.TypeA,
			deltas: map[string]counterDelta{
				"udp": {queries: 1, cached: 1, sent: 1},
				"tcp": {queries: 1, cached: 1, sent: 1},
				"dot": {queries: 1, cached: 1, sent: 1},
				"doh": {queries: 1, cached: 1},
			},
			why: "a cache hit is the number an operator uses to prove the cache is working",
		},
		{
			prefix: "fail", qtype: dns.TypeA,
			deltas: map[string]counterDelta{
				"udp": {queries: 1, errors: 1, sent: 1},
				"tcp": {queries: 1, errors: 1, sent: 1},
				"dot": {queries: 1, errors: 1, sent: 1},
				"doh": {queries: 1, errors: 1},
			},
			why: "SERVFAIL is still an answer, so errors and sent both move",
		},
		{
			prefix: "drop", qtype: dns.TypeA,
			deltas: map[string]counterDelta{
				// A drop counts as an error as well as a drop: HandlerStats
				// documents Errors as "SERVFAIL, dropped or malformed". It is
				// the only evidence a client-side timeout was a policy
				// decision, so it must not be under-counted.
				"udp": {queries: 1, errors: 1, dropped: 1},
				"tcp": {queries: 1, errors: 1, dropped: 1},
				"dot": {queries: 1, errors: 1, dropped: 1},
				// This is the regression pin: a DoH drop used to be counted
				// as an error only, so the console's transport tab
				// under-reported the drops an operator was told about in the
				// aggregate. The body is still written (a DoH client waiting
				// on a socket must not read a hang-up as a transport error),
				// which is why sent is not bumped either: the write is
				// net/http's, and the DoH path counts none.
				"doh": {queries: 1, errors: 1, dropped: 1},
			},
			why: "a drop is a counter the client cannot see",
		},
		{
			prefix: "big", qtype: dns.TypeTXT,
			deltas: map[string]counterDelta{
				// Only UDP can truncate: a stream transport sends the whole
				// answer and lets the client reassemble it, so counting a
				// truncation there would be a lie about what happened.
				"udp": {queries: 1, truncated: 1, sent: 1},
				"tcp": {queries: 1, sent: 1},
				"dot": {queries: 1, sent: 1},
				"doh": {queries: 1},
			},
			why: "a 3 KiB answer to a 512-byte client is truncated, and that is a fact worth counting",
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			h := newHarness(t, nil)
			before := h.srv.ByTransport()
			want := map[string]TransportStats{}
			for _, transport := range statTransports {
				d, ok := tc.deltas[transport]
				if !ok {
					t.Fatalf("fixture: no expected delta for %s", transport)
				}
				want[transport] = applyDelta(before[transport], d)
			}
			for _, transport := range statTransports {
				exchangeOnce(t, h, transport, question(tc.prefix+".counters.test.", tc.qtype))
			}
			// The client can be answered before the last counter is bumped
			// (the stream paths bump "sent" after the write returns), so the
			// assertion waits for the settled state rather than for a
			// duration.
			waitForCounters(t, h, want)
		})
	}
}

// waitForCounters blocks until every counter set named in want is exactly
// matched. The comparison is exact in both directions: a counter that moved
// too far never matches, so double counting fails here rather than hiding
// behind a >= . want may name a subset of the transports, which is how a
// focused test states "this transport moved by this much".
func waitForCounters(t *testing.T, h *harness, want map[string]TransportStats) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for {
		got := h.srv.ByTransport()
		if countersMatch(got, want) {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("counters did not settle within 3s:%s", counterDiff(got, want))
		}
		time.Sleep(2 * time.Millisecond)
	}
}

func countersMatch(got, want map[string]TransportStats) bool {
	for transport, w := range want {
		if got[transport] != w {
			return false
		}
	}
	return true
}

func counterDiff(got, want map[string]TransportStats) string {
	var b strings.Builder
	for _, transport := range statTransports {
		if w, ok := want[transport]; ok && got[transport] != w {
			fmt.Fprintf(&b, "\n  %s: got %+v\n     want %+v", transport, got[transport], w)
		}
	}
	return b.String()
}

// TestPerTransportCountersSumToTheAggregate: /status reports both the
// aggregate (HandlerStats) and the per-transport breakdown, and the console
// shows them side by side. If the aggregate ever counted something the
// breakdown does not — or double counted it — the two views of the same server
// would disagree, which is worse than either being wrong alone. The per-
// transport numbers are the ones the aggregate is built from, so summing them
// has to reproduce it exactly.
func TestPerTransportCountersSumToTheAggregate(t *testing.T) {
	h := newHarness(t, nil)
	// One of every outcome over every transport, so the sum is not trivially
	// equal to a single counter.
	names := []string{"plain.test.", "block.test.", "cached.test.", "fail.test.", "drop.test."}
	for _, transport := range statTransports {
		for _, name := range names {
			exchangeOnce(t, h, transport, question(name, dns.TypeA))
		}
	}
	waitFor(t, 3*time.Second, "every query to be counted", func() bool {
		return h.srv.Stats().Queries == uint64(len(statTransports)*len(names))
	})
	agg := h.srv.Stats()
	byT := h.srv.ByTransport()

	var sum HandlerStats
	for _, transport := range statTransports {
		ts := byT[transport]
		sum.Queries += ts.Queries
		sum.Blocked += ts.Blocked
		sum.Cached += ts.Cached
		sum.Errors += ts.Errors
		sum.Dropped += ts.Dropped
		if ts.Queries != uint64(len(names)) {
			t.Errorf("%s.Queries = %d, want %d", transport, ts.Queries, len(names))
		}
		if ts.Blocked != 1 {
			t.Errorf("%s.Blocked = %d, want 1", transport, ts.Blocked)
		}
		if ts.Cached != 1 {
			t.Errorf("%s.Cached = %d, want 1", transport, ts.Cached)
		}
		// One SERVFAIL and one drop: HandlerStats documents Errors as
		// including drops.
		if ts.Errors != 2 {
			t.Errorf("%s.Errors = %d, want 2 (one SERVFAIL, one drop)", transport, ts.Errors)
		}
		if ts.Dropped != 1 {
			t.Errorf("%s.Dropped = %d, want 1", transport, ts.Dropped)
		}
	}
	if agg.Queries != sum.Queries {
		t.Errorf("Stats().Queries = %d, want the sum of the transports %d", agg.Queries, sum.Queries)
	}
	if agg.Blocked != sum.Blocked {
		t.Errorf("Stats().Blocked = %d, want %d", agg.Blocked, sum.Blocked)
	}
	if agg.Cached != sum.Cached {
		t.Errorf("Stats().Cached = %d, want %d", agg.Cached, sum.Cached)
	}
	if agg.Errors != sum.Errors {
		t.Errorf("Stats().Errors = %d, want %d", agg.Errors, sum.Errors)
	}
	if agg.Dropped != sum.Dropped {
		t.Errorf("Stats().Dropped = %d, want %d", agg.Dropped, sum.Dropped)
	}
	if agg.Queries == 0 || agg.Dropped == 0 {
		t.Fatalf("the fixture counted nothing: %+v", agg)
	}
}

// TestDoHDropIsCountedAsDropped is the named pin for the DoH drop counter. It
// is separate from the table above because it is the one that was wrong: the
// DoH handler answered a dropped query with an empty body but bumped only
// "errors", so /status said the query was dropped while the transport tab
// said it was an error. An operator reading the transport tab to work out
// which listener is dropping traffic could not tell DoH from a resolver that
// was failing outright.
func TestDoHDropIsCountedAsDropped(t *testing.T) {
	h := newHarness(t, nil)
	before := h.srv.ByTransport()["doh"]
	res := h.askDoHPOST(question("drop.test.", dns.TypeA))
	// The body is still required: unlike UDP, hanging up on an HTTP exchange
	// reads to the client as a transport error and it retries elsewhere.
	if len(res.Answer) != 0 {
		t.Errorf("dropped query came back with %d answers: %v", len(res.Answer), res.Answer)
	}
	waitForCounters(t, h, map[string]TransportStats{
		"doh": applyDelta(before, counterDelta{queries: 1, errors: 1, dropped: 1}),
	})
	if got := h.srv.Stats().Dropped; got != 1 {
		t.Errorf("Stats().Dropped = %d, want 1 for a single DoH drop", got)
	}
	// And the other transports are untouched: a drop is counted against the
	// transport it arrived on, not smeared over the server.
	byT := h.srv.ByTransport()
	for _, transport := range []string{"udp", "tcp", "dot"} {
		if byT[transport] != (TransportStats{Transport: transport}) {
			t.Errorf("a DoH query moved %s counters: %+v", transport, byT[transport])
		}
	}
}

// TestMalformedDatagramIsAnErrorAndNotAQuery is the aggregate half of
// TestUDPMalformedPacketIsCountedNotAnswered (which pins the per-transport
// counter and that nothing is written back). It matters separately because it
// makes Errors able to exceed Queries: a datagram that cannot be parsed is
// counted as an error, and never as a query, because the server never accepted
// a question out of it. Anything rendering one as a ratio of the other — the
// console, a dashboard, an alert — has to cope with that, and a server under
// a garbage flood is exactly when somebody looks.
func TestMalformedDatagramIsAnErrorAndNotAQuery(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)
	}
	t.Cleanup(func() { _ = 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 to be counted", func() bool {
		return h.srv.Stats().Errors == 1
	})
	st := h.srv.Stats()
	if st.Queries != 0 {
		t.Errorf("Stats().Queries = %d, want 0: an unparsable datagram is not a query", st.Queries)
	}
	if st.Dropped != 0 {
		t.Errorf("Stats().Dropped = %d, want 0: it is counted as an error, not as a drop", st.Dropped)
	}
	if n := len(h.fake.seen()); n != 0 {
		t.Errorf("garbage reached the engine %d times: %+v", n, h.fake.seen())
	}
	if got := h.srv.ByTransport()["udp"].Sent; got != 0 {
		t.Errorf("udp.Sent = %d, want 0: nothing was written back", got)
	}
}

// TestInflightTracksTheHandler pins the one field of HandlerStats that is not a
// total. It is read by the health endpoint and the console's "in flight" gauge
// to tell a busy resolver from a stuck one, so it has to rise while a query
// is inside the pipeline and return to zero when the answer goes out. The
// second half is the part a broken implementation gets wrong: a gauge that
// ratchets up and never comes down turns the health check into a permanent
// alarm.
func TestInflightTracksTheHandler(t *testing.T) {
	h := newHarness(t, nil)
	release := holdOnce(t, h.fake)
	conn, err := net.DialTimeout("tcp", h.addr("tcp"), 2*time.Second)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
	co := &dns.Conn{Conn: conn}
	if err := co.WriteMsg(question("inflight.test.", dns.TypeA)); err != nil {
		t.Fatalf("write: %v", err)
	}
	// The handler signals arrival before it parks on the gate, and the
	// in-flight counter is incremented before that, so it is already non-zero
	// here — no sleeping on "it has probably got there by now".
	h.fake.waitArrived(t)
	if got := h.srv.Stats().Inflight; got < 1 {
		t.Fatalf("Stats().Inflight = %d while a query is inside the handler, want at least 1", got)
	}
	// A query that has not been answered is not a completed query yet. The
	// counter that matters to an operator (Queries) must therefore still say
	// nothing about it, or the rate graph counts work that is still running.
	if got := h.srv.ByTransport()["tcp"].Queries; got != 0 {
		t.Errorf("tcp.Queries = %d while the only query is still in flight, want 0", got)
	}

	release()
	res, err := co.ReadMsg()
	if err != nil {
		t.Fatalf("read: %v", err)
	}
	if a := answerAddr(res); a != "203.0.113.9" {
		t.Errorf("A = %q, want 203.0.113.9", a)
	}
	waitFor(t, 2*time.Second, "the in-flight count to fall back to zero", func() bool {
		ts := h.srv.ByTransport()["tcp"]
		return h.srv.Stats().Inflight == 0 && ts.Queries == 1 && ts.Sent == 1
	})
	ts := h.srv.ByTransport()["tcp"]
	if ts.Queries != 1 || ts.Sent != 1 {
		t.Errorf("tcp counters after the query finished = %+v, want queries=1 sent=1", ts)
	}
}

// TestByTransportAlwaysReportsTheSameFourTransports: the console's transport
// tab iterates whatever /status returns. A row that disappears when a listener
// is disabled, or a missing key when nothing has been served yet, turns a
// zero into a rendering bug at exactly the moment somebody is trying to read
// the numbers. All four are always present; a disabled or unused transport is
// a row of zeros.
func TestByTransportAlwaysReportsTheSameFourTransports(t *testing.T) {
	check := func(t *testing.T, got map[string]TransportStats) {
		t.Helper()
		if len(got) != 4 {
			t.Errorf("ByTransport() has %d entries (%+v), want 4", len(got), got)
		}
		for _, transport := range statTransports {
			ts, ok := got[transport]
			if !ok {
				t.Errorf("ByTransport() is missing %q: %+v", transport, got)
				continue
			}
			if ts.Transport != transport {
				t.Errorf("ByTransport()[%q].Transport = %q, want the key echoed back", transport, ts.Transport)
			}
		}
	}
	// A server that has never served a query and has no engine at all.
	fresh := New(Options{})
	if got := fresh.ByTransport(); len(got) != 4 {
		t.Errorf("a server that never started reports %+v, want four zeroed transports", got)
	}
	check(t, fresh.ByTransport())

	// A server with only the two plain listeners bound and one query served:
	// the two disabled transports are still there, at zero.
	h := newHarness(t, func(c *config.Config) {
		c.DNS.DoT = config.DoT{Enabled: false}
		c.DNS.DoH = config.DoH{Enabled: false}
	})
	h.ask("udp", "unused-transport.test.", dns.TypeA)
	got := h.srv.ByTransport()
	check(t, got)
	for _, transport := range []string{"dot", "doh"} {
		if got[transport] != (TransportStats{Transport: transport}) {
			t.Errorf("disabled transport %q = %+v, want an all-zero row", transport, got[transport])
		}
	}
	if got["udp"].Queries != 1 {
		t.Errorf("udp.Queries = %d, want 1", got["udp"].Queries)
	}
}