package server import ( "bytes" "context" "crypto/tls" "io" "net" "runtime" "strconv" "strings" "sync" "testing" "time" "github.com/miekg/dns" "git.k8s-git-k6n.net/mats/fancy-dns/internal/config" ) // 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) } }) // DoT is the same shape as TCP with a TLS handshake in front of it, and // it is here because it is a *different* shutdown path in the product: // the stream loop takes an idle timeout from the config, so a DoT query // can be answered after its connection's deadline has been set, while a // plain TCP one cannot. t.Run("dot", func(t *testing.T) { h := newHarness(t, nil) release := holdOnce(t, h.fake) tc := h.clientTLS() tc.NextProtos = []string{"dot"} conn, err := tls.Dial("tcp", h.addr("dot"), tc) if err != nil { t.Fatalf("dial dot: %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) stopped := startStop(h.srv, 5*time.Second) waitPortRefused(t, "tcp", h.addr("dot"), "Stop closing the DoT listener") release() res, err := co.ReadMsg() if err != nil { t.Fatalf("the in-flight DoT 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) } if got := h.srv.Stats().Inflight; got != 0 { t.Errorf("Stats().Inflight = %d after Stop returned, want 0", got) } }) // UDP is the transport that cannot make the promise the other three keep. // There is no connection: the answer goes out on the very socket the // query arrived on, so a socket closed before the query has been answered // loses the datagram outright and the client has nothing but its own // retry. Stop therefore has to let the queries already in the pipeline // finish before it closes that socket, and the test pins that it does. t.Run("udp", func(t *testing.T) { h := newHarness(t, nil) release := holdOnce(t, h.fake) pc, err := net.ListenPacket("udp", "127.0.0.1:0") if err != nil { t.Fatalf("listen: %v", err) } t.Cleanup(func() { _ = pc.Close() }) wire, err := question("inflight-shutdown.test.", dns.TypeA).Pack() if err != nil { t.Fatalf("pack: %v", err) } addr, err := net.ResolveUDPAddr("udp", h.addr("udp")) if err != nil { t.Fatalf("resolve: %v", err) } if _, err := pc.WriteTo(wire, addr); err != nil { t.Fatalf("write: %v", err) } h.fake.waitArrived(t) // The probe is the opposite of the stream cases: a UDP port cannot be // "refused", but Stop closing the socket is observable as the port // becoming bindable. That also proves the drain finished — the close // is the last thing the listener's close func does, after the // in-flight packets have been answered. stopped := startStop(h.srv, 5*time.Second) waitPortFree(t, "udp", h.addr("udp"), "Stop draining and closing the UDP listener") release() _ = pc.SetReadDeadline(time.Now().Add(3 * time.Second)) buf := make([]byte, 4096) n, _, err := pc.ReadFrom(buf) if err != nil { t.Fatalf("the in-flight UDP query was lost across Stop: %v", err) } res := new(dns.Msg) if err := res.Unpack(buf[:n]); err != nil { t.Fatalf("unpack: %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 waited for the pipeline, so the app's cache snapshot that // follows it cannot race a query that is still writing to it. 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) } } // TestStopIsNotBlockedByAnIdleStreamConnection is the other half of shutdown: // a client that connects and says nothing must not be able to hold the // process open. // // The stream loop waits for its connections before Stop returns, and a // connection parked in ReadMsg with no deadline in front of it only ends when // the client goes away — the context is checked between reads, not during one. // So the grace period the caller gave Stop (15s in main) used to be spent // waiting for idle sockets, and Stop then returned context.DeadlineExceeded, // which the app logs as "shutdown completed with errors". A client that // opens a TCP connection and never sends a query is also trivially cheap to // make, so this is a shutdown DoS rather than an edge case. // // The DoT subtest sets a 60s idle timeout on purpose: a config value must not // be able to hold a shutdown hostage, and the harness default (2s) would hide // a regression behind a lucky tick. func TestStopIsNotBlockedByAnIdleStreamConnection(t *testing.T) { for _, transport := range []string{"tcp", "dot"} { t.Run(transport, func(t *testing.T) { h := newHarness(t, func(c *config.Config) { c.Server.Timeouts.TLS = config.Duration(60 * time.Second) }) conn := h.idleStream(t, transport) // One query has been exchanged, so the server has accepted the // connection and is inside ReadMsg waiting for a second one that // never comes. Nothing is written and nothing is read from here. // The bound is the caller's, not Stop's: a correct Stop returns // long before this fires, and a wrong one comes back with // context.DeadlineExceeded and an error message that names the // cause. ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() stopErr := h.srv.Stop(ctx) if stopErr != nil { t.Fatalf("Stop with an idle %s connection open = %v: a client that never sends a query must not be able to spend the whole shutdown grace period", transport, stopErr) } // Stop did not just give up on the connection, it ended it: the // client must see the socket go away rather than sit there until // its own deadline. A timeout error here is the tell — the server // is still holding the connection and nothing closed it. _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) buf := make([]byte, 1) if _, err := conn.Read(buf); err == nil { t.Fatalf("the server sent %d bytes on a connection with no query on it", len(buf)) } else if ne, ok := err.(net.Error); ok && ne.Timeout() { t.Errorf("the idle %s connection was still open 2s after Stop returned", transport) } }) } } // TestStopLeavesNoGoroutineRunningServerCode: Stop returning means the // listeners are done, and if a goroutine of ours is still inside the package // after that, the process is not. That matters twice over — the app snapshots // the cache and closes the store the moment Stop returns, so anything still // running can write to a closed store; and a goroutine blocked in a read on a // listener that is never closed again is a leak that a long-lived process // accumulates one connection at a time. // // The check counts goroutines whose stack still has a frame in this package // rather than comparing runtime.NumGoroutine() against a baseline: the test // binary shares its process with net/http, crypto/tls and the race detector's // own goroutines, so a raw count is noise. A frame in this package cannot // appear in any of those. func TestStopLeavesNoGoroutineRunningServerCode(t *testing.T) { h := newHarness(t, nil) // Every transport, plus the two shapes that keep goroutines alive after a // query is answered: a connection nobody has closed and an HTTP client // with a keep-alive connection in it. for _, transport := range []string{"udp", "tcp", "dot", "doh", "doh-post"} { if a := answerAddr(h.ask(transport, "leak."+transport+".test.", dns.TypeA)); a != "203.0.113.9" { t.Fatalf("%s query before the leak check returned %q", transport, a) } } // An accepted connection that then goes quiet is the goroutine a leak // check has to be able to see: it is the one a shutdown can strand. _ = h.idleStream(t, "tcp") if err := stopWithin(t, h.srv, 3*time.Second); err != nil { t.Fatalf("Stop: %v", err) } // Polled, not sampled once: a goroutine that is on its way out is not a // leak, and the drain and Shutdown paths hand work to other goroutines // (net/http's, the h2 server's) that finish just after Stop returns. Only // a goroutine that is *still* there after a second is a leak. deadline := time.Now().Add(3 * time.Second) var leaked []string for { leaked = serverGoroutines() if len(leaked) == 0 { break } if time.Now().After(deadline) { for _, g := range leaked { t.Errorf("still running 3s after Stop returned:\n%s", g) } t.FailNow() } time.Sleep(25 * time.Millisecond) } } // serverGoroutines returns the stack of every goroutine that has a frame in // this package, except the caller's own: the test function is itself in this // package, so without the exclusion every run would find exactly one. func serverGoroutines() []string { self := goroutineID() buf := make([]byte, 64<<10) for { n := runtime.Stack(buf, true) if n < len(buf) { buf = buf[:n] break } buf = make([]byte, 2*len(buf)) } var out []string for _, g := range strings.Split(string(buf), "\n\n") { if !strings.Contains(g, "fancy-dns/internal/server.") { continue } // "goroutine 775 [running]:" — the id is the second field, and the // caller's own block is the one this scan must not report. if fields := strings.Fields(g); len(fields) >= 2 { if n, err := strconv.ParseUint(fields[1], 10, 64); err == nil && n == self { continue } } if id, _, ok := strings.Cut(g, " "); ok { if n, err := strconv.ParseUint(strings.TrimPrefix(id, "goroutine"), 10, 64); err == nil && n == self { continue } } out = append(out, g) } return out } // goroutineID returns the id of the calling goroutine. runtime has no API for // it, but every stack dump starts with "goroutine N [state]:", and the // current goroutine's dump with all=false is the one place that is printed. func goroutineID() uint64 { var buf [64]byte n := runtime.Stack(buf[:], false) fields := strings.Fields(string(buf[:n])) if len(fields) < 2 { return 0 } id, err := strconv.ParseUint(fields[1], 10, 64) if err != nil { return 0 } return id }