fancy-dns/internal/server/harness_test.go

package server

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"io"
	"log/slog"
	"math/big"
	"net"
	"net/http"
	"net/netip"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/miekg/dns"

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

// testVersion is what the fake server reports in the DoH Server header.
const testVersion = "9.9.9-test"

// recorded is one request as the handler saw it. The question is copied out
// of the *dns.Msg because the engine is allowed to mutate the message, and a
// test that asserts on a pointer the pipeline may rewrite is a flake waiting
// to happen.
type recorded struct {
	transport string
	name      string
	qtype     uint16
	remote    netip.Addr
	tcp       bool
	names     []string
	tags      []string
	dohPath   string
	ednsSize  uint16
}

// fakeHandler is the one Handler every transport funnels into. It records what
// it was asked and answers from a fixed recipe, so a test can tell "the UDP
// listener delivered this" from "the engine produced this".
type fakeHandler struct {
	mu   sync.Mutex
	reqs []recorded

	inflight atomic.Int64
	flushes  atomic.Uint64
	// packErr is the last failure to pack a response the handler built. A
	// response that will not pack is silently dropped on the write path and
	// presents as a client timeout, so it is recorded here rather than being
	// discovered three layers away.
	packErr error

	// arrived is signalled on every Handle entry, so a test can wait for a
	// query to be *inside* the pipeline instead of sleeping and hoping.
	arrived chan struct{}

	gate    chan struct{} // non-nil while queries are parked
	release chan struct{} // closed at cleanup so a held gate cannot wedge shutdown

	muGate sync.Mutex
}

func newFakeHandler() *fakeHandler {
	return &fakeHandler{
		arrived: make(chan struct{}, 64),
		release: make(chan struct{}),
	}
}

// Handle records the request and answers from the recipe below.
func (f *fakeHandler) Handle(ctx context.Context, req engine.Request) engine.Response {
	f.inflight.Add(1)
	defer f.inflight.Add(-1)

	rec := recorded{
		transport: req.Transport,
		remote:    req.Remote,
		tcp:       req.TCP,
		names:     append([]string(nil), req.ClientNames...),
		tags:      append([]string(nil), req.ClientTags...),
		dohPath:   req.DoHPath,
	}
	if req.Msg != nil && len(req.Msg.Question) > 0 {
		rec.name = req.Msg.Question[0].Name
		rec.qtype = req.Msg.Question[0].Qtype
		if o := req.Msg.IsEdns0(); o != nil {
			rec.ednsSize = o.UDPSize()
		}
	}
	f.mu.Lock()
	f.reqs = append(f.reqs, rec)
	f.mu.Unlock()

	select {
	case f.arrived <- struct{}{}:
	default:
	}

	if g := f.gateCh(); g != nil {
		select {
		case <-g:
		case <-f.release:
		case <-ctx.Done():
		}
	}

	res := f.answer(req.Msg, rec)
	if res.Msg != nil {
		if _, err := res.Msg.Pack(); err != nil {
			f.mu.Lock()
			f.packErr = err
			f.mu.Unlock()
		}
	}
	return res
}

// answer is the canned response. The question name selects the outcome, so a
// test drives a blocked/cached/dropped/big answer by asking for it by name
// instead of reaching into the server. The reply is built from a copy of the
// query: SetReply carries the query's ID across, and a reply that invents its
// own ID is discarded by every real client.
func (f *fakeHandler) answer(q *dns.Msg, rec recorded) engine.Response {
	m := newReply(q)
	switch {
	case strings.HasPrefix(rec.name, "block."):
		return engine.Response{Msg: blockedReply(m), Action: "block", BlockSource: "test"}
	case strings.HasPrefix(rec.name, "cached."):
		return engine.Response{Msg: answerA(m), Action: "cached", CacheHit: true, CacheKind: "hit"}
	case strings.HasPrefix(rec.name, "fail."):
		m.Rcode = dns.RcodeServerFailure
		return engine.Response{Msg: m, Action: "error"}
	case strings.HasPrefix(rec.name, "drop."):
		return engine.Response{Drop: true, Action: "drop"}
	case strings.HasPrefix(rec.name, "big."):
		return engine.Response{Msg: bigAnswer(m), Action: "allow"}
	case len(rec.name) == 0:
		m.Rcode = dns.RcodeFormatError
		return engine.Response{Msg: m, Action: "error"}
	}
	return engine.Response{Msg: answerA(m), Action: "allow"}
}

func (f *fakeHandler) Inflight() int64 { return f.inflight.Load() }
func (f *fakeHandler) Flush()          { f.flushes.Add(1) }

// flushes returns how many times the admin flush action reached the engine.
func (f *fakeHandler) flushCount() uint64 { return f.flushes.Load() }

// hold parks every handler that arrives from now on until the returned func is
// called. It is how a test gets a query that is reliably in flight.
func (f *fakeHandler) hold() func() {
	f.muGate.Lock()
	defer f.muGate.Unlock()
	if f.gate != nil {
		return func() {}
	}
	g := make(chan struct{})
	f.gate = g
	return func() { close(g) }
}

func (f *fakeHandler) gateCh() chan struct{} {
	f.muGate.Lock()
	defer f.muGate.Unlock()
	return f.gate
}

// waitArrived blocks until a request has reached the handler.
func (f *fakeHandler) waitArrived(t *testing.T) {
	t.Helper()
	select {
	case <-f.arrived:
	case <-time.After(5 * time.Second):
		t.Fatal("no query reached the handler")
	}
}

// lastPackErr reports the last response that failed to pack.
func (f *fakeHandler) lastPackErr() error {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.packErr
}

// seen returns a copy of everything recorded so far.
func (f *fakeHandler) seen() []recorded {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]recorded(nil), f.reqs...)
}

// only asserts that exactly one request was recorded and returns it.
func (f *fakeHandler) only(t *testing.T) recorded {
	t.Helper()
	got := f.seen()
	if len(got) != 1 {
		t.Fatalf("handler saw %d requests, want exactly 1: %+v", len(got), got)
	}
	return got[0]
}

// last returns the most recent recorded request.
func (f *fakeHandler) last() recorded {
	got := f.seen()
	if len(got) == 0 {
		return recorded{}
	}
	return got[len(got)-1]
}

// reset drops the recording, so one server can serve several independent
// questions.
func (f *fakeHandler) reset() {
	f.mu.Lock()
	f.reqs = nil
	f.mu.Unlock()
}

// newReply builds the shell of a reply: the query's ID, opcode, question and
// OPT. A real engine's finalize step does the EDNS0 handshake and SetReply
// carries neither the OPT nor a made-up ID, so a fake that skipped this would
// be the only thing in the path dropping them.
func newReply(q *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(q.Copy())
	m.Authoritative = false
	if o := q.IsEdns0(); o != nil {
		m.SetEdns0(o.UDPSize(), o.Do())
	}
	return m
}

// answerA adds the canned A record to a prepared reply.
func answerA(m *dns.Msg) *dns.Msg {
	if len(m.Question) > 0 && m.Question[0].Qtype == dns.TypeA {
		m.Answer = append(m.Answer, &dns.A{
			Hdr: dns.RR_Header{Name: m.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
			A:   net.ParseIP("203.0.113.9"),
		})
	}
	return m
}

// blockedReply adds the 0.0.0.0 sinkhole answer.
func blockedReply(m *dns.Msg) *dns.Msg {
	if len(m.Question) > 0 && m.Question[0].Qtype == dns.TypeA {
		m.Answer = append(m.Answer, &dns.A{
			Hdr: dns.RR_Header{Name: m.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
			A:   net.ParseIP("0.0.0.0"),
		})
	}
	return m
}

// bigAnswer builds a multi-kilobyte answer out of many short TXT strings.
// One long string does NOT work: miekg/dns v1.1.73 reports a character-string
// over 255 bytes as zero length, so Msg.Pack fails with "dns: string exceeded
// 255 bytes" and, on a server write path, fails silently. 15 x 200 lands just
// over 3 KiB: far outside a 512-byte client buffer, inside the 4096 EDNS0
// buffer the server accepts.
func bigAnswer(m *dns.Msg) *dns.Msg {
	txt := make([]string, 15)
	for i := range txt {
		txt[i] = strings.Repeat("x", 200)
	}
	name := "."
	if len(m.Question) > 0 {
		name = m.Question[0].Name
	}
	m.Answer = append(m.Answer, &dns.TXT{
		Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 60},
		Txt: txt,
	})
	return m
}

// harness is a started Server with all four listeners on loopback:0, a fake
// engine behind them and a self-signed CA the test clients pin.
type harness struct {
	t     *testing.T
	cfg   *config.Config
	store *config.Store
	fake  *fakeHandler
	srv   *Server
	pool  *x509.CertPool
	cert  tls.Certificate
}

func newHarness(t *testing.T, mutate func(*config.Config)) *harness {
	t.Helper()
	certFile, keyFile := writeTestCert(t)
	pool := testCertPool(t, certFile)

	cfg := config.Default()
	cfg.DNS.UDP = []string{"127.0.0.1:0"}
	cfg.DNS.TCP = []string{"127.0.0.1:0"}
	cfg.DNS.DoT = config.DoT{Enabled: true, Listen: []string{"127.0.0.1:0"}, CertFile: certFile, KeyFile: keyFile, MinVersion: "1.2"}
	cfg.DNS.DoH = config.DoH{Enabled: true, Listen: []string{"127.0.0.1:0"}, CertFile: certFile, KeyFile: keyFile, BasePath: "/dns-query", MinVersion: "1.2", AllowPost: true}
	cfg.DNS.DoH.Transport = config.Transport{Mode: "auto"}
	// A short stream idle timeout keeps a test that leaves a connection open
	// from holding a goroutine for the four-second default.
	cfg.Server.Timeouts.TLS = config.Duration(2 * time.Second)
	if mutate != nil {
		mutate(cfg)
	}

	store := config.NewStore(cfg, "")
	fake := newFakeHandler()
	srv := New(Options{
		Engine:  fake,
		Config:  store,
		Logger:  slog.New(slog.NewTextHandler(io.Discard, nil)),
		Version: testVersion,
	})
	if err := srv.Start(context.Background()); err != nil {
		t.Fatalf("start: %v", err)
	}
	t.Cleanup(func() {
		close(fake.release)
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = srv.Stop(ctx)
	})
	return &harness{t: t, cfg: cfg, store: store, fake: fake, srv: srv, pool: pool}
}

// addr returns the bound address of one transport, read back from the
// listener snapshot rather than guessed.
func (h *harness) addr(transport string) string {
	h.t.Helper()
	for _, l := range h.srv.Listeners() {
		if l.Transport == transport {
			return l.Addr
		}
	}
	h.t.Fatalf("no %s listener; have %+v", transport, h.srv.Listeners())
	return ""
}

func (h *harness) transports() []string {
	var out []string
	for _, l := range h.srv.Listeners() {
		out = append(out, l.Transport)
	}
	return out
}

// base returns the normalised DoH base path.
func (h *harness) base() string { return normaliseBase(h.cfg.DNS.DoH.BasePath) }

func (h *harness) clientTLS() *tls.Config {
	return &tls.Config{RootCAs: h.pool, MinVersion: tls.VersionTLS12}
}

// writeTestCert generates a self-signed certificate for 127.0.0.1 into
// t.TempDir and returns the two paths. Generating it in-test keeps the
// encrypted-listener tests self-contained: no testdata to expire.
func writeTestCert(t *testing.T) (certFile, keyFile string) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("key: %v", err)
	}
	tmpl := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "fancy-dns test"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(24 * time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:                  true,
		DNSNames:              []string{"localhost"},
		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		t.Fatalf("certificate: %v", err)
	}
	kb, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		t.Fatalf("marshal key: %v", err)
	}
	dir := t.TempDir()
	certFile = filepath.Join(dir, "cert.pem")
	keyFile = filepath.Join(dir, "key.pem")
	if err := os.WriteFile(certFile, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}), 0o600); err != nil {
		t.Fatalf("write cert: %v", err)
	}
	if err := os.WriteFile(keyFile, pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: kb}), 0o600); err != nil {
		t.Fatalf("write key: %v", err)
	}
	return certFile, keyFile
}

// testCertPool returns a pool trusting only the generated certificate, so a
// client that skips verification is a deliberate choice rather than a habit.
func testCertPool(t *testing.T, certFile string) *x509.CertPool {
	t.Helper()
	pemBytes, err := os.ReadFile(certFile)
	if err != nil {
		t.Fatalf("read cert: %v", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(pemBytes) {
		t.Fatal("test certificate did not parse into a pool")
	}
	return pool
}

// ask builds a query and exchanges it over the named transport. The client
// timeouts are short on purpose: a hung listener must fail the test in
// seconds, not at the package timeout.
func (h *harness) ask(transport string, name string, qtype uint16) *dns.Msg {
	h.t.Helper()
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(name), qtype)
	m.RecursionDesired = true
	switch transport {
	case "udp":
		return h.askUDP(m)
	case "tcp":
		return h.askTCP(m)
	case "dot":
		return h.askDoT(m)
	case "doh":
		return h.askDoHGET(m)
	case "doh-post":
		return h.askDoHPOST(m)
	}
	h.t.Fatalf("unknown transport %q", transport)
	return nil
}

func (h *harness) askUDP(m *dns.Msg) *dns.Msg {
	h.t.Helper()
	c := &dns.Client{Net: "udp", Timeout: 3 * time.Second}
	res, _, err := c.Exchange(m, h.addr("udp"))
	if err != nil {
		h.t.Fatalf("udp exchange: %v", err)
	}
	return res
}

func (h *harness) askTCP(m *dns.Msg) *dns.Msg {
	h.t.Helper()
	c := &dns.Client{Net: "tcp", Timeout: 3 * time.Second}
	res, _, err := c.Exchange(m, h.addr("tcp"))
	if err != nil {
		h.t.Fatalf("tcp exchange: %v", err)
	}
	return res
}

func (h *harness) askDoT(m *dns.Msg) *dns.Msg {
	h.t.Helper()
	tc := h.clientTLS()
	tc.NextProtos = []string{"dot"}
	c := &dns.Client{Net: "tcp-tls", Timeout: 3 * time.Second, TLSConfig: tc}
	res, _, err := c.Exchange(m, h.addr("dot"))
	if err != nil {
		h.t.Fatalf("dot exchange: %v", err)
	}
	return res
}

func (h *harness) dohURL(path string) string {
	return "https://" + h.addr("doh") + path
}

func (h *harness) dohClient(forceH2 bool) *http.Client {
	tr := &http.Transport{TLSClientConfig: h.clientTLS(), ForceAttemptHTTP2: forceH2}
	h.t.Cleanup(tr.CloseIdleConnections)
	return &http.Client{Transport: tr, Timeout: 5 * time.Second}
}

// askDoHGET issues the RFC 8484 GET profile with the wire message in ?dns=.
func (h *harness) askDoHGET(m *dns.Msg) *dns.Msg {
	h.t.Helper()
	res, body := h.doHGET(m)
	if res.StatusCode != http.StatusOK {
		h.t.Fatalf("doh GET status %d: %s", res.StatusCode, body)
	}
	return unpackDoH(h.t, body)
}

// askDoHPOST issues the RFC 8484 POST profile.
func (h *harness) askDoHPOST(m *dns.Msg) *dns.Msg {
	h.t.Helper()
	res, body := h.doHPOST(m)
	if res.StatusCode != http.StatusOK {
		h.t.Fatalf("doh POST status %d: %s", res.StatusCode, body)
	}
	return unpackDoH(h.t, body)
}

func (h *harness) doHGET(m *dns.Msg) (*http.Response, []byte) {
	h.t.Helper()
	wire, err := m.Pack()
	if err != nil {
		h.t.Fatalf("pack query: %v", err)
	}
	u := h.dohURL(h.base() + "?dns=" + base64.RawURLEncoding.EncodeToString(wire))
	res, err := h.dohClient(false).Get(u)
	if err != nil {
		h.t.Fatalf("doh GET: %v", err)
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(res.Body)
	return res, body
}

func (h *harness) doHPOST(m *dns.Msg) (*http.Response, []byte) {
	h.t.Helper()
	return h.doHPostWith(m, contentTypeDNS)
}

func (h *harness) doHPostWith(m *dns.Msg, contentType string) (*http.Response, []byte) {
	h.t.Helper()
	wire, err := m.Pack()
	if err != nil {
		h.t.Fatalf("pack query: %v", err)
	}
	res, err := h.dohClient(false).Post(h.dohURL(h.base()), contentType, strings.NewReader(string(wire)))
	if err != nil {
		h.t.Fatalf("doh POST: %v", err)
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(res.Body)
	return res, body
}

// idleStream hands back a stream connection the server has definitely
// accepted, and that the caller then leaves alone: one query has been
// exchanged on it, so the server is inside ReadMsg waiting for the next one.
//
// Waiting for a real exchange is the point. A test that only dials races the
// accept loop — Stop can close the listener while the connection is still in
// the backlog, the server never accepts it, and a shutdown that hangs on idle
// connections passes by accident.
func (h *harness) idleStream(t *testing.T, transport string) net.Conn {
	t.Helper()
	var (
		conn net.Conn
		err  error
	)
	if transport == "dot" {
		tc := h.clientTLS()
		tc.NextProtos = []string{"dot"}
		conn, err = tls.Dial("tcp", h.addr("dot"), tc)
	} else {
		conn, err = net.DialTimeout("tcp", h.addr("tcp"), 2*time.Second)
	}
	if err != nil {
		t.Fatalf("dial %s: %v", transport, err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
	co := &dns.Conn{Conn: conn}
	name := "accepted-" + transport + ".test."
	if err := co.WriteMsg(question(name, dns.TypeA)); err != nil {
		t.Fatalf("write on the %s connection: %v", transport, err)
	}
	res, err := co.ReadMsg()
	if err != nil {
		t.Fatalf("the %s connection was never served: %v", transport, err)
	}
	if a := answerAddr(res); a != "203.0.113.9" {
		t.Fatalf("the %s connection answered %q, want 203.0.113.9", transport, a)
	}
	// The deadline is dropped again so a later read blocks for real: this
	// connection must not end the test on its own.
	_ = conn.SetDeadline(time.Time{})
	return conn
}

func unpackDoH(t *testing.T, body []byte) *dns.Msg {
	t.Helper()
	m := new(dns.Msg)
	if err := m.Unpack(body); err != nil {
		t.Fatalf("unpack doh body (%d bytes): %v", len(body), err)
	}
	return m
}

// waitFor polls cond until it holds or the deadline passes. Every wait in this
// suite is bounded: a hang looks exactly like a suite failure.
func waitFor(t *testing.T, timeout time.Duration, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for {
		if cond() {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("timed out after %s waiting for %s", timeout, what)
		}
		time.Sleep(2 * time.Millisecond)
	}
}

// question builds a query the fake's recipe understands.
func question(name string, qtype uint16) *dns.Msg {
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(name), qtype)
	m.RecursionDesired = true
	return m
}

var _ = engine.Request{}