package server

import (
	"bytes"
	"context"
	"io"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/miekg/dns"
)

// Shutdown is the only part of this package that can lose a client's query,
// so it is tested as behaviour rather than as a sequence: the socket is
// released, Stop is idempotent, and a query that is already inside the
// pipeline is still answered. The last one is the invariant that matters — a
// resolver that answers "connection closed" to a query it had already accepted
// makes every client retry against somewhere else.

// holdOnce parks every handler that arrives from now on and returns a release
// func that is safe to call twice, so a failing test can unblock the gate in
// its cleanup without a double-close panic.
func holdOnce(t *testing.T, f *fakeHandler) func() {
	t.Helper()
	release := f.hold()
	var once sync.Once
	out := func() { once.Do(release) }
	t.Cleanup(out)
	return out
}

// startStop calls srv.Stop in the background and returns the channel it will
// report on. stopWithin is that plus a bounded wait; the in-flight test needs
// the gap, because it has to watch the shutdown overlap a parked query instead
// of waiting for it to finish.
func startStop(srv *Server, ctxTimeout time.Duration) <-chan error {
	done := make(chan error, 1)
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), ctxTimeout)
		defer cancel()
		done <- srv.Stop(ctx)
	}()
	return done
}

// awaitStop waits for a Stop started by startStop.
func awaitStop(t *testing.T, done <-chan error, d time.Duration) error {
	t.Helper()
	select {
	case err := <-done:
		return err
	case <-time.After(d):
		t.Fatalf("Stop did not return within %s", d)
		return nil
	}
}

// stopWithin calls srv.Stop and fails the test if it has not returned within d.
// A hang inside a test binary is reported as a package timeout, which says
// nothing about which test caused it; bounding every Stop call here makes the
// failure name the test.
func stopWithin(t *testing.T, srv *Server, d time.Duration) error {
	t.Helper()
	return awaitStop(t, startStop(srv, 5*time.Second), d)
}

// bindFree tries to bind network/addr and closes the socket again. It reports
// the error rather than failing so a caller can retry: a socket that was
// closed a moment ago can still be in TIME_WAIT, and the point of every use
// here is that a *second* bind of the port works, not that the first one does.
func bindFree(network, addr string) error {
	if network == "udp" {
		pc, err := net.ListenPacket("udp", addr)
		if err != nil {
			return err
		}
		return pc.Close()
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return ln.Close()
}

// waitPortFree blocks until network/addr can be bound. Use it only when nothing
// else holds the port: a connected client's own local socket is on that same
// port, so a bind probe can never succeed while a connection is open. The
// in-flight test uses waitPortRefused for exactly that reason.
func waitPortFree(t *testing.T, network, addr, what string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for {
		err := bindFree(network, addr)
		if err == nil {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("%s: %s %s is still bound 2s after Stop was called (%v)", what, network, addr, err)
		}
		time.Sleep(2 * time.Millisecond)
	}
}

// waitPortRefused blocks until a dial to addr is refused, which is how the
// in-flight test observes that Stop has reached the point where it closes the
// listeners — before it waits for its goroutines. A refused dial is the right
// probe there: the parked query holds a *connection* to that port, so binding
// it would fail for a reason that has nothing to do with Stop.
func waitPortRefused(t *testing.T, network, addr, what string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for {
		conn, err := net.DialTimeout(network, addr, 200*time.Millisecond)
		if err != nil {
			return
		}
		conn.Close()
		if time.Now().After(deadline) {
			t.Fatalf("%s: %s %s still accepts connections 2s after Stop was called", what, network, addr)
		}
		time.Sleep(2 * time.Millisecond)
	}
}

// TestStopReleasesEveryListenerPort: an operator who restarts fancy-dns — or
// hands the port to something else — must be able to bind it again. It also
// empties the listener snapshot, because that is what the admin API and the
// cluster node state read: a listener still reported as running after the
// process stopped serving on it sends the health check the wrong answer.
func TestStopReleasesEveryListenerPort(t *testing.T) {
	h := newHarness(t, nil)
	addrs := map[string]string{}
	for _, l := range h.srv.Listeners() {
		addrs[l.Transport] = l.Addr
	}
	if len(addrs) != 4 {
		t.Fatalf("fixture bound %d listeners, want 4: %+v", len(addrs), addrs)
	}
	if err := stopWithin(t, h.srv, 2*time.Second); err != nil {
		t.Fatalf("Stop: %v", err)
	}
	for transport, addr := range addrs {
		network := "tcp"
		if transport == "udp" {
			network = "udp"
		}
		// One retry: TIME_WAIT is a property of the socket, not a listener
		// still being held, and failing on it would make this a flake.
		var err error
		for i := 0; i < 2; i++ {
			if err = bindFree(network, addr); err == nil {
				break
			}
			time.Sleep(20 * time.Millisecond)
		}
		if err != nil {
			t.Errorf("%s %s could not be bound after Stop: %v", network, addr, err)
		}
	}
	if got := h.srv.Listeners(); len(got) != 0 {
		t.Errorf("Listeners() = %+v after Stop, want empty", got)
	}
}

// TestStopIsIdempotent: Stop is called from signal handling, from an admin
// endpoint and from a test cleanup, so more than one of them can reach a
// server that is already down. A second call must be a no-op returning nil —
// not a panic on a half-drained listener list, and not a wait on goroutines
// that are already gone. The same is true of a server that never started at
// all, which is what a config error produces.
func TestStopIsIdempotent(t *testing.T) {
	h := newHarness(t, nil)
	if err := stopWithin(t, h.srv, 2*time.Second); err != nil {
		t.Fatalf("first Stop: %v", err)
	}
	// The code clears started before it closes anything, so the second call
	// short-circuits on that flag and never re-closes a listener or re-waits
	// on the wait group. Pinned: a second Stop that returned the first Stop's
	// error, or waited again, would be a bug in its own right.
	if err := stopWithin(t, h.srv, 2*time.Second); err != nil {
		t.Errorf("second Stop = %v, want nil", err)
	}
	if err := stopWithin(t, h.srv, 2*time.Second); err != nil {
		t.Errorf("third Stop = %v, want nil", err)
	}
	// A server that never started: nothing to close, and nothing to say.
	if err := stopWithin(t, New(Options{}), 2*time.Second); err != nil {
		t.Errorf("Stop on a server that never started = %v, want nil", err)
	}
}

// TestStopWithAnAlreadyCancelledContextReturnsPromptly: the context bounds how
// long Stop waits, and a caller that has already given up must not be made to
// wait for a listener to close. Only the bound is asserted — never a
// duration. The return *value* is deliberately not asserted: Stop races its
// own wait group against ctx.Done() and takes whichever is ready first, so the
// same call can return nil or context.Canceled depending on how the shutdown
// went. Asserting either would pin the scheduler, not the code.
//
// What IS asserted is that the listeners are released even then: closing them
// is not the part the context is allowed to interrupt, because that is the
// work the next process start needs.
func TestStopWithAnAlreadyCancelledContextReturnsPromptly(t *testing.T) {
	h := newHarness(t, nil)
	udp := h.addr("udp")
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	done := make(chan error, 1)
	go func() { done <- h.srv.Stop(ctx) }()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Stop with an already-cancelled context did not return within 2s")
	}
	waitPortFree(t, "udp", udp, "a Stop whose context was already cancelled")
}

// TestInFlightQueryIsAnsweredAcrossStop is the invariant in this file. A
// query that has been accepted and is inside the pipeline belongs to the
// client: closing the listener must not make it vanish. The client here has
// its query parked in the handler while Stop runs, so if the shutdown path
// abandoned in-flight work the client would see a closed connection (TCP) or a
// 5xx/EOF (DoH) and every client would retry elsewhere.
//
// The test proves Stop is really in flight before the handler is released: it
// waits for the listener's port to become bindable, which Stop does before it
// waits for its goroutines. Without that ordering the test would pass even if
// the shutdown never overlapped the query.
func TestInFlightQueryIsAnsweredAcrossStop(t *testing.T) {
	t.Run("tcp", func(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-shutdown.test.", dns.TypeA)); err != nil {
			t.Fatalf("write: %v", err)
		}
		h.fake.waitArrived(t)

		// Stop runs in the background and the test watches the overlap: it
		// closes the listener before it waits for its goroutines, so a refused
		// dial is proof that the shutdown is running *while* the query is
		// still parked in the handler. Without that ordering the test would
		// pass even if the shutdown never overlapped the query at all.
		stopped := startStop(h.srv, 5*time.Second)
		waitPortRefused(t, "tcp", h.addr("tcp"), "Stop closing the TCP listener")

		release()
		res, err := co.ReadMsg()
		if err != nil {
			t.Fatalf("the in-flight query was lost across Stop: %v", err)
		}
		if a := answerAddr(res); a != "203.0.113.9" {
			t.Errorf("A = %q, want 203.0.113.9 (reply: %v)", a, res)
		}
		if err := awaitStop(t, stopped, 3*time.Second); err != nil {
			t.Errorf("Stop: %v", err)
		}
		// Stop waits for the connection goroutines, so by the time it returns
		// nothing is still inside the engine. A shutdown that returned while a
		// query was in flight would leave the process exiting with work
		// running.
		if got := h.srv.Stats().Inflight; got != 0 {
			t.Errorf("Stats().Inflight = %d after Stop returned, want 0", got)
		}
	})

	t.Run("doh", func(t *testing.T) {
		h := newHarness(t, nil)
		release := holdOnce(t, h.fake)
		cl := h.dohClient(false)
		type reply struct {
			status int
			body   []byte
			err    error
		}
		got := make(chan reply, 1)
		// The request is issued from its own goroutine: it does not come back
		// until the handler is released, which happens after Stop has started.
		// h.doHPOST is not usable here because it calls t.Fatalf, and a
		// Fatalf from a non-test goroutine does not stop the test.
		go func() {
			wire, err := question("inflight-shutdown.test.", dns.TypeA).Pack()
			if err != nil {
				got <- reply{err: err}
				return
			}
			res, err := cl.Post(h.dohURL(h.base()), contentTypeDNS, bytes.NewReader(wire))
			if err != nil {
				got <- reply{err: err}
				return
			}
			defer res.Body.Close()
			body, _ := io.ReadAll(res.Body)
			got <- reply{status: res.StatusCode, body: body}
		}()
		h.fake.waitArrived(t)

		// The DoH path is the interesting one: http.Server.Shutdown waits for
		// the in-flight request rather than cancelling it, so Stop has to
		// block until the answer has been written. It must still return, and
		// the client must still get its answer. The probe is the same as on
		// the TCP side: the listener is closed before Stop waits.
		stopped := startStop(h.srv, 5*time.Second)
		waitPortRefused(t, "tcp", h.addr("doh"), "Stop closing the DoH listener")

		release()
		select {
		case r := <-got:
			if r.err != nil {
				t.Fatalf("the in-flight DoH request failed across Stop: %v", r.err)
			}
			if r.status != 200 {
				t.Fatalf("status = %d, want 200: %s", r.status, r.body)
			}
			if a := answerAddr(unpackDoH(t, r.body)); a != "203.0.113.9" {
				t.Errorf("A = %q, want 203.0.113.9", a)
			}
		case <-time.After(3 * time.Second):
			t.Fatal("the in-flight DoH request never came back")
		}
		if err := awaitStop(t, stopped, 3*time.Second); err != nil {
			t.Errorf("Stop: %v", err)
		}
		if got := h.srv.Stats().Inflight; got != 0 {
			t.Errorf("Stats().Inflight = %d after Stop returned, want 0", got)
		}
	})
}

// TestStopWithNoQueryInFlightReturnsAndStartsAgain: an idle Stop has nothing
// to wait for, so it returns immediately — the bound here is generous because
// the point is "does not hang", not "is fast". Then the process has to be able
// to serve again: the same Server is started a second time, binds four fresh
// listeners and answers. The counters deliberately survive the restart: /status
// is a lifetime view, and a counter that reset on restart would make the
// console's time series jump to zero every time the pod rolled.
func TestStopWithNoQueryInFlightReturnsAndStartsAgain(t *testing.T) {
	h := newHarness(t, nil)
	before := h.addr("udp")
	if err := stopWithin(t, h.srv, 2*time.Second); err != nil {
		t.Fatalf("Stop: %v", err)
	}
	if err := h.srv.Start(context.Background()); err != nil {
		t.Fatalf("Start after Stop: %v", err)
	}
	if got := len(h.srv.Listeners()); got != 4 {
		t.Errorf("Listeners() after the restart has %d entries, want 4", got)
	}
	res := h.ask("udp", "after-restart.test.", dns.TypeA)
	if a := answerAddr(res); a != "203.0.113.9" {
		t.Errorf("A after the restart = %q, want 203.0.113.9", a)
	}
	if got := h.fake.last().name; got != "after-restart.test." {
		t.Errorf("the restarted server did not reach the engine: last question %q", got)
	}
	if got := h.srv.ByTransport()["udp"].Queries; got != 1 {
		t.Errorf("udp.Queries = %d after the restart, want 1 (counters are a lifetime total)", got)
	}
	// The new listeners are different sockets: the harness binds port 0, so a
	// restart is only possible because Stop gave the old ones back. Asserting
	// it also keeps the test honest about which server answered.
	if after := h.addr("udp"); after == before {
		t.Logf("udp rebound the same address %s; port 0 makes this possible but unlikely", after)
	}
	// And the harness's own cleanup still stops the restarted server without
	// complaint, which is the t.Cleanup path every test in this package takes.
	if got := h.srv.Stats().Inflight; got != 0 {
		t.Errorf("Stats().Inflight = %d, want 0", got)
	}
}