package server

import (
	"io"
	"net"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
)

// TestAllFourListenersShareOneHandler is the package's headline claim: the
// four transports are thin shells around a single engine. A bug that made one
// of them bypass the handler (or answer from its own state) would show up here
// as a missing or duplicated call, so every transport is driven with a real
// client and checked for exactly one Handle, the right transport label, the
// right client address and the right answer on the wire.
func TestAllFourListenersShareOneHandler(t *testing.T) {
	cases := []struct {
		transport string
		// tcp marks the transports that set engine.Request.TCP, i.e. the ones
		// allowed to answer with a response larger than a UDP datagram.
		wantTCP bool
	}{
		{transport: "udp", wantTCP: false},
		{transport: "tcp", wantTCP: true},
		{transport: "dot", wantTCP: true},
		{transport: "doh", wantTCP: true},
		{transport: "doh-post", wantTCP: true},
	}
	for _, tc := range cases {
		t.Run(tc.transport, func(t *testing.T) {
			h := newHarness(t, nil)
			name := "host." + tc.transport + ".test."
			res := h.ask(tc.transport, name, dns.TypeA)

			got := h.fake.only(t)
			wantTransport := tc.transport
			if wantTransport == "doh-post" {
				wantTransport = "doh"
			}
			if got.transport != wantTransport {
				t.Errorf("Request.Transport = %q, want %q", got.transport, wantTransport)
			}
			if got.name != dns.Fqdn(name) {
				t.Errorf("question = %q, want %q", got.name, dns.Fqdn(name))
			}
			if got.qtype != dns.TypeA {
				t.Errorf("qtype = %d, want %d", got.qtype, dns.TypeA)
			}
			// The client is always loopback here, whatever the transport: that
			// is what "the client address reached the engine" has to mean for a
			// test that binds 127.0.0.1.
			if got.remote.String() != "127.0.0.1" {
				t.Errorf("Request.Remote = %v, want 127.0.0.1", got.remote)
			}
			if got.tcp != tc.wantTCP {
				t.Errorf("Request.TCP = %v, want %v", got.tcp, tc.wantTCP)
			}
			// And the answer came back over the wire, not just out of the
			// handler: same ID, the question echoed, one A record.
			if len(res.Question) != 1 || res.Question[0].Name != dns.Fqdn(name) {
				t.Errorf("reply question = %+v, want %s", res.Question, dns.Fqdn(name))
			}
			if a := answerAddr(res); a != "203.0.113.9" {
				t.Errorf("reply A = %q, want 203.0.113.9 (reply: %v)", a, res)
			}
		})
	}
}

// TestAllFourListenersAreUpAndReported pins the Listeners() contract the admin
// API and the cluster node state read: one entry per bound listener, with the
// real address read back from the socket rather than the configured one (the
// tests bind port 0, so the two differ).
func TestAllFourListenersAreUpAndReported(t *testing.T) {
	h := newHarness(t, nil)
	seen := map[string]string{}
	for _, l := range h.srv.Listeners() {
		if !l.Enabled || !l.Running {
			t.Errorf("listener %+v reported as not enabled/running after Start", l)
		}
		if l.Err != "" {
			t.Errorf("listener %s reported error %q after a clean start", l.Transport, l.Err)
		}
		seen[l.Transport] = l.Addr
	}
	for _, want := range []string{"udp", "tcp", "dot", "doh"} {
		addr, ok := seen[want]
		if !ok {
			t.Fatalf("no %s listener in %v", want, seen)
		}
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			t.Errorf("listener %s address %q: %v", want, addr, err)
			continue
		}
		if host != "127.0.0.1" || port == "0" {
			t.Errorf("listener %s address = %q, want a real 127.0.0.1 port", want, addr)
		}
	}
	if got := len(h.srv.Listeners()); got != 4 {
		t.Errorf("Listeners() has %d entries, want 4", got)
	}
}

// TestUDPListenerAnswersEveryDatagram is the raw-packet counterpart of the
// client-based test: it uses a net.PacketConn so a reply that the DNS client
// would have silently retried still has to arrive. The client address is the
// socket's local port, which is what makes it a different source address per
// exchange and therefore a check that the server writes back to the peer it
// read from rather than to whoever it read from first.
func TestUDPListenerAnswersEveryDatagram(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)
	}
	for i := 0; i < 3; i++ {
		q := question("raw.test.", dns.TypeA)
		q.Id = uint16(0x1000 + i)
		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, from, err := pc.ReadFrom(buf)
		if err != nil {
			t.Fatalf("read %d: %v", i, err)
		}
		if from.String() != addr.String() {
			t.Errorf("reply from %v, want %v", from, addr)
		}
		res := new(dns.Msg)
		if err := res.Unpack(buf[:n]); err != nil {
			t.Fatalf("unpack %d: %v", i, err)
		}
		if res.Id != q.Id {
			t.Errorf("reply id = %d, want %d", res.Id, q.Id)
		}
		if a := answerAddr(res); a != "203.0.113.9" {
			t.Errorf("reply A = %q, want 203.0.113.9", a)
		}
	}
	waitFor(t, 2*time.Second, "three handler calls", func() bool {
		return len(h.fake.seen()) == 3
	})
	for _, rec := range h.fake.seen() {
		if rec.transport != "udp" {
			t.Errorf("transport = %q, want udp", rec.transport)
		}
	}
}

// TestDoHOverHTTP2 exercises the h2c wrapper over a TLS listener. The DoH
// listener advertises h2 in ALPN, so a real client (curl --http2, a browser)
// negotiates it; if the h2 path were broken the request would be answered over
// HTTP/1.1 or dropped, and neither is what an operator configured.
func TestDoHOverHTTP2(t *testing.T) {
	h := newHarness(t, nil)
	wire, err := question("h2.test.", dns.TypeA).Pack()
	if err != nil {
		t.Fatalf("pack: %v", err)
	}
	res, err := h.dohClient(true).Post(h.dohURL(h.base()), contentTypeDNS, strings.NewReader(string(wire)))
	if err != nil {
		t.Fatalf("POST: %v", err)
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(res.Body)
	if res.Proto != "HTTP/2.0" {
		t.Errorf("proto = %q, want HTTP/2.0", res.Proto)
	}
	if res.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200: %s", res.StatusCode, body)
	}
	if a := answerAddr(unpackDoH(t, body)); a != "203.0.113.9" {
		t.Errorf("A = %q, want 203.0.113.9", a)
	}
	got := h.fake.only(t)
	if got.transport != "doh" {
		t.Errorf("transport = %q, want doh", got.transport)
	}
	if got.dohPath != h.base() {
		t.Errorf("DoHPath = %q, want %q", got.dohPath, h.base())
	}
}

// TestDoHHealthz is the endpoint an Ingress or a pod probe hits. It is
// registered on the same mux as the DNS handler, so a base path that shadows
// "/healthz" would be a real misconfiguration worth catching.
func TestDoHHealthz(t *testing.T) {
	h := newHarness(t, nil)
	res, err := h.dohClient(false).Get(h.dohURL("/healthz"))
	if err != nil {
		t.Fatalf("GET /healthz: %v", err)
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		t.Errorf("status = %d, want 200", res.StatusCode)
	}
	if len(h.fake.seen()) != 0 {
		t.Errorf("healthz reached the engine: %+v", h.fake.seen())
	}
}

// TestTCPPipelining pins RFC 1035 §4.2.2: one connection carries several
// queries with length prefixes, and a resolver that has a keep-alive
// connection open must be able to use it. A transport that closed after one
// query would still pass every single-query test.
func TestTCPPipelining(t *testing.T) {
	h := newHarness(t, nil)
	conn, err := net.DialTimeout("tcp", h.addr("tcp"), 3*time.Second)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
	co := &dns.Conn{Conn: conn}
	for i := 0; i < 3; i++ {
		q := question("pipelined.test.", dns.TypeA)
		q.Id = uint16(0x2000 + i)
		if err := co.WriteMsg(q); err != nil {
			t.Fatalf("write %d: %v", i, err)
		}
		res, err := co.ReadMsg()
		if err != nil {
			t.Fatalf("read %d: %v", i, err)
		}
		if res.Id != q.Id {
			t.Errorf("reply %d id = %d, want %d", i, res.Id, q.Id)
		}
		if a := answerAddr(res); a != "203.0.113.9" {
			t.Errorf("reply %d A = %q, want 203.0.113.9", i, a)
		}
	}
	waitFor(t, 2*time.Second, "three pipelined queries", func() bool {
		return len(h.fake.seen()) == 3
	})
}

// answerAddr returns the first A record's address, or "".
func answerAddr(m *dns.Msg) string {
	if m == nil {
		return ""
	}
	for _, rr := range m.Answer {
		if a, ok := rr.(*dns.A); ok {
			return a.A.String()
		}
	}
	return ""
}