fancy-dns/internal/engine/harness_test.go

package engine

import (
	"bytes"
	"context"
	"encoding/json"
	"net"
	"net/netip"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/accesslog"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/acl"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/blocklist"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/cache"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/localzone"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/metrics"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/resolver"
)

// fakeUpstream is a UDP resolver that answers every name under .test. with a
// fixed address, so a test can tell "resolved upstream" from "answered locally"
// without touching the internet. It records the queries it received, which is
// how the forwarded-message tests inspect what the engine put on the wire.
type fakeUpstream struct {
	*dns.Server
	addr string

	mu      sync.Mutex
	hits    map[string]int
	queries []*dns.Msg
	// writeErr is the last error the handler's write returned: a fake upstream
	// that silently fails to pack looks exactly like a client-side timeout.
	writeErr error

	// gate, when set, parks handlers until it is closed. arrived is signalled
	// before parking, so a test can wait for the query to be *inside* the
	// pipeline rather than guessing with a sleep.
	gate    chan struct{}
	arrived chan struct{}
	// release is closed at cleanup so a held gate can never wedge shutdown.
	release chan struct{}
}

func newFakeUpstream(t *testing.T) *fakeUpstream {
	t.Helper()
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	u := &fakeUpstream{
		addr:    pc.LocalAddr().String(),
		hits:    make(map[string]int),
		arrived: make(chan struct{}, 64),
		release: make(chan struct{}),
	}
	u.Server = &dns.Server{PacketConn: pc, Handler: dns.HandlerFunc(u.handle)}
	go func() { _ = u.ActivateAndServe() }()
	t.Cleanup(func() {
		close(u.release)
		_ = u.Shutdown()
	})
	return u
}

func (u *fakeUpstream) handle(w dns.ResponseWriter, q *dns.Msg) {
	u.mu.Lock()
	if len(q.Question) > 0 {
		u.hits[q.Question[0].Name]++
	}
	u.queries = append(u.queries, q.Copy())
	u.mu.Unlock()
	select {
	case u.arrived <- struct{}{}:
	default:
	}
	if g := u.gateCh(); g != nil {
		select {
		case <-g:
		case <-u.release:
		}
	}
	if len(q.Question) == 0 {
		return
	}
	m := new(dns.Msg)
	m.SetReply(q)
	name := q.Question[0].Name
	switch {
	case !strings.HasSuffix(strings.ToLower(name), ".test."):
		m.Rcode = dns.RcodeNameError
		m.Ns = []dns.RR{testSOA("test.", 60)}
	default:
		// A "tc*" name answers *and* sets TC, which is how a real upstream
		// reports an answer that did not fit; it must not be cached.
		if strings.HasPrefix(strings.ToLower(name), "tc.") {
			m.Truncated = true
		}
		switch q.Question[0].Qtype {
		case dns.TypeTXT:
			// "big*" names answer with a 4 KiB TXT so a 512-byte client
			// buffer has something to overflow.
			if strings.HasPrefix(strings.ToLower(name), "big") {
				// Built as a record rather than parsed, and as many short
				// strings rather than one long one: this version of miekg/dns
				// reports a character-string over 255 bytes as zero length, so a
				// single 4 KiB TXT never packs (a silent write failure that looks
				// like a client timeout). 15 x 200 lands just over 3 KiB, inside
				// the 4 KiB buffer the engine advertises, and far outside a
				// 512-byte client buffer.
				txt := make([]string, 15)
				for i := range txt {
					txt[i] = strings.Repeat("x", 200)
				}
				rr := &dns.TXT{
					Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 60},
					Txt: txt,
				}
				m.Answer = append(m.Answer, rr)
			} else {
				rr, _ := dns.NewRR(name + " 60 IN A 203.0.113.9")
				m.Answer = append(m.Answer, rr)
			}
		case dns.TypeA:
			rr, _ := dns.NewRR(name + " 60 IN A 203.0.113.7")
			m.Answer = append(m.Answer, rr)
		}
		m.Ns = []dns.RR{testSOA("test.", 60)}
	}
	if err := w.WriteMsg(m); err != nil {
		u.mu.Lock()
		u.writeErr = err
		u.mu.Unlock()
	}
}

// lastWriteErr reports the last write failure, or nil.
func (u *fakeUpstream) lastWriteErr() error {
	u.mu.Lock()
	defer u.mu.Unlock()
	return u.writeErr
}

func (u *fakeUpstream) gateCh() chan struct{} {
	u.mu.Lock()
	defer u.mu.Unlock()
	return u.gate
}

// hold parks every handler until the returned func is called. Handlers that
// arrive after it is released are not parked.
func (u *fakeUpstream) hold() func() {
	u.mu.Lock()
	defer u.mu.Unlock()
	if u.gate != nil {
		return func() {}
	}
	g := make(chan struct{})
	u.gate = g
	return func() { close(g) }
}

func (u *fakeUpstream) hitCount(name string) int {
	u.mu.Lock()
	defer u.mu.Unlock()
	return u.hits[dns.Fqdn(name)]
}

func (u *fakeUpstream) totalHits() int {
	u.mu.Lock()
	defer u.mu.Unlock()
	n := 0
	for _, c := range u.hits {
		n += c
	}
	return n
}

// received returns the queries the upstream saw, in order.
func (u *fakeUpstream) received() []*dns.Msg {
	u.mu.Lock()
	defer u.mu.Unlock()
	return append([]*dns.Msg(nil), u.queries...)
}

// waitArrived blocks until at least one query has reached the handler.
func (u *fakeUpstream) waitArrived(t *testing.T) {
	t.Helper()
	select {
	case <-u.arrived:
	case <-time.After(5 * time.Second):
		t.Fatal("upstream never received a query")
	}
}

// testSOA is the authority record the fake upstream attaches to every answer,
// which is what a real resolver does and what makes NXDOMAIN well-formed.
func testSOA(name string, ttl uint32) dns.RR {
	return &dns.SOA{
		Hdr:     dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: ttl},
		Ns:      "ns.test.",
		Mbox:    "hostmaster.test.",
		Serial:  1,
		Refresh: 3600,
		Retry:   600,
		Expire:  604800,
		Minttl:  60,
	}
}

// testClock is the injectable time source the cache manager uses. Advancing it
// is how a test makes an entry expire without sleeping.
type testClock struct {
	mu sync.Mutex
	t  time.Time
}

func newTestClock() *testClock { return &testClock{t: time.Now()} }

func (c *testClock) Now() time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.t
}

func (c *testClock) Advance(d time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.t = c.t.Add(d)
}

// hopt tunes a harness. The zero value is a fully working engine: fake
// upstream, local zone, cache, blocklist and query log.
type hopt struct {
	acl      func(*config.ACL)
	blockCfg func(*config.Block)
	server   func(*config.Server)
	cacheCfg func(*config.Cache)
	dnsCfg   func(*config.DNS)
	// rules are compiled into the blocklist set the engine holds.
	rules []blocklist.Rule
	// localZone is a master file; empty means no local zone at all.
	localZone string
	// upstreams replaces the default single fake upstream. A nil slice with
	// noPool false means "use the fake upstream".
	upstreams []config.UpstreamServer
	// noPool wires a nil pool, which is the "no upstream configured" case.
	noPool bool
	// serveStale enables the stale window and prefetching.
	serveStale time.Duration
	// clock installs the injectable clock into the cache manager.
	clock bool
	// timeout is the per-exchange upstream timeout.
	timeout time.Duration
}

type harness struct {
	t      *testing.T
	eng    *Engine
	cfg    *config.Config
	store  *config.Store
	pool   *resolver.Pool
	cache  *cache.Manager
	log    *accesslog.Logger
	lines  *lineSink
	up     *fakeUpstream
	clock  *testClock
	blocks *blocklist.Holder
	local  *localzone.Holder
	aclH   *acl.Holder
}

func newHarness(t *testing.T, o hopt) *harness {
	t.Helper()
	up := newFakeUpstream(t)
	timeout := o.timeout
	if timeout <= 0 {
		timeout = 5 * time.Second
	}

	cfg := config.Default()
	cfg.Cache.Persist.Enabled = false
	cfg.Server.Timeouts.Upstre = config.Duration(timeout)
	// A generous upstream buffer: an upstream cannot send us more than we
	// advertise, so the truncation tests need room for a 4 KiB answer to
	// arrive whole and then be trimmed for a small client.
	cfg.DNS.MaxUDPSize = 4096
	cfg.Cache.ServeStale = config.Duration(o.serveStale)
	cfg.Block.Lists = nil
	if o.acl != nil {
		o.acl(&cfg.ACL)
	}
	if o.blockCfg != nil {
		o.blockCfg(&cfg.Block)
	}
	if o.server != nil {
		o.server(&cfg.Server)
	}
	if o.cacheCfg != nil {
		o.cacheCfg(&cfg.Cache)
	}
	if o.dnsCfg != nil {
		o.dnsCfg(&cfg.DNS)
	}

	// The default pool is the fake upstream; a dead address is what the
	// upstream-failure test needs, and noPool needs no pool at all.
	if !o.noPool {
		servers := o.upstreams
		if servers == nil {
			servers = []config.UpstreamServer{{Addr: up.addr, Proto: "udp", Weight: 1}}
		}
		cfg.Upstream.Servers = servers
		cfg.Upstream.Policy = "parallel"
		cfg.Upstream.MaxAttempts = 1
		cfg.Upstream.HealthCheck = 0
		cfg.Upstream.BreakerThreshold = 5
		cfg.Upstream.BreakerCooldown = config.Duration(30 * time.Second)
	} else {
		cfg.Upstream.Servers = nil
	}

	// A local zone file, if asked for.
	var local *localzone.Store
	if o.localZone != "" {
		path := filepath.Join(t.TempDir(), "local.zone")
		if err := os.WriteFile(path, []byte(o.localZone), 0o600); err != nil {
			t.Fatalf("write zone: %v", err)
		}
		l, err := localzone.Load([]string{path}, localzone.Options{})
		if err != nil {
			t.Fatalf("load zone: %v", err)
		}
		local = l
	} else {
		l, err := localzone.Load(nil, localzone.Options{})
		if err != nil {
			t.Fatalf("load empty zone set: %v", err)
		}
		local = l
	}

	clock := newTestClock()
	cm := cache.NewManager(cache.Options{
		Size:       1000,
		Shards:     4,
		ServeStale: o.serveStale,
		Prefetch:   true,
		Now:        clock.Now,
	}, nil)

	lines := &lineSink{}
	// Every action is logged, so a test can count entries without reasoning
	// about the sampling policy.
	log := accesslog.New(accesslog.Options{
		Writer: lines, SampleRate: 1, LogBlocked: true, LogCached: true, LogErrors: true,
	})
	t.Cleanup(func() { _ = log.Close() })

	blocks := blocklist.NewHolder(buildSet(o.rules...))
	aclH := acl.NewHolder(acl.Compile(cfg.ACL))
	localH := localzone.NewHolder(local)
	store := config.NewStore(cfg, "")

	eng := New(Options{
		Config:     store,
		Logger:     log,
		Blocklists: blocks,
		Cache:      cm,
		ACL:        aclH,
		Local:      localH,
	})
	if !o.noPool {
		eng.pool.Store(resolver.NewPoolWithUDPSize(cfg.Upstream, cfg.DNS.MaxUDPSize))
	}

	h := &harness{
		t: t, eng: eng, cfg: cfg, store: store, cache: cm, log: log, lines: lines,
		up: up, clock: clock, blocks: blocks, local: localH, aclH: aclH,
	}
	if !o.noPool {
		h.pool = eng.Pool()
	}
	return h
}

func buildSet(rules ...blocklist.Rule) *blocklist.Set {
	b := blocklist.NewBuilder()
	for _, r := range rules {
		b.Add(r, "test", nil)
	}
	return b.Build()
}

// setRules swaps in a new compiled set, as a reload would.
func (h *harness) setRules(rules ...blocklist.Rule) {
	h.blocks.Set(buildSet(rules...))
}

func (h *harness) advance(d time.Duration) { h.clock.Advance(d) }

// request builds a Request the way the UDP listener would.
func (h *harness) request(name string, qtype uint16) Request {
	q := new(dns.Msg)
	q.SetQuestion(dns.Fqdn(name), qtype)
	q.RecursionDesired = true
	return Request{Msg: q, Remote: netip.MustParseAddr("127.0.0.1"), Transport: "udp", Start: time.Now()}
}

func (h *harness) ask(name string, qtype uint16) Response {
	h.t.Helper()
	return h.eng.Handle(context.Background(), h.request(name, qtype))
}

func (h *harness) askWith(req Request) Response {
	h.t.Helper()
	return h.eng.Handle(context.Background(), req)
}

// askEDNS asks with an EDNS0 buffer of the given size.
func (h *harness) askEDNS(name string, qtype uint16, size uint16) Response {
	h.t.Helper()
	req := h.request(name, qtype)
	req.Msg.SetEdns0(size, false)
	return h.askWith(req)
}

// askTCP asks with the TCP flag set, so the UDP size cap does not apply.
func (h *harness) askTCP(name string, qtype uint16) Response {
	h.t.Helper()
	req := h.request(name, qtype)
	req.TCP = true
	req.Transport = "tcp"
	return h.askWith(req)
}

// entries returns the access-log lines written so far, newest first.
func (h *harness) entries() []accesslog.Entry {
	h.log.Flush()
	return h.lines.entries()
}

func (h *harness) logged() int {
	h.log.Flush()
	return h.lines.count()
}

func (h *harness) totals() metrics.Totals { return metrics.TotalsNow() }

// answerA returns the first A record's address, or "" when the answer has no
// address record.
func answerA(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 ""
}

func firstTTL(m *dns.Msg) uint32 {
	if m == nil || len(m.Answer) == 0 {
		return 0
	}
	return m.Answer[0].Header().Ttl
}

// requireRcode fails unless the response carries rcode.
func requireRcode(t *testing.T, res Response, rcode int) {
	t.Helper()
	if res.Msg == nil {
		t.Fatalf("no response message (drop=%v)", res.Drop)
	}
	if res.Msg.Rcode != rcode {
		t.Fatalf("rcode = %s, want %s", dns.RcodeToString[res.Msg.Rcode], dns.RcodeToString[rcode])
	}
}

func requireAction(t *testing.T, res Response, want accesslog.Action) {
	t.Helper()
	if res.Action != want {
		t.Fatalf("action = %q, want %q", res.Action, want)
	}
}

func requireNoUpstreamHit(t *testing.T, h *harness, name string) {
	t.Helper()
	if n := h.up.hitCount(name); n != 0 {
		t.Fatalf("upstream saw %d queries for %s, want 0", n, name)
	}
}

// lineSink collects the access log's output so a test can read it back.
type lineSink struct {
	mu   sync.Mutex
	buf  []byte
	rows []accesslog.Entry
}

func (s *lineSink) Write(p []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.buf = append(s.buf, p...)
	for {
		i := bytes.IndexByte(s.buf, '\n')
		if i < 0 {
			break
		}
		line := s.buf[:i]
		s.buf = s.buf[i+1:]
		var raw struct {
			Client   string  `json:"client"`
			Group    string  `json:"client_group"`
			Name     string  `json:"name"`
			Action   string  `json:"action"`
			Protocol string  `json:"protocol"`
			List     string  `json:"list"`
			Upstream string  `json:"upstream"`
			Duration float64 `json:"duration_ms"`
		}
		if err := json.Unmarshal(line, &raw); err != nil {
			panic("access log line is not JSON: " + string(line))
		}
		s.rows = append(s.rows, accesslog.Entry{
			Client: raw.Client, Group: raw.Group, Name: raw.Name,
			Action: accesslog.Action(raw.Action), Protocol: raw.Protocol,
			List: raw.List, Upstream: raw.Upstream,
		})
	}
	return len(p), nil
}

func (s *lineSink) count() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.rows)
}

func (s *lineSink) entries() []accesslog.Entry {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]accesslog.Entry, len(s.rows))
	for i, r := range s.rows {
		out[len(s.rows)-1-i] = r
	}
	return out
}