fancy-dns/internal/app/harness_test.go

package app

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/miekg/dns"

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

// testUpstream is a fake recursive resolver: it answers A queries for names
// under test. and NXDOMAIN for everything else, so a test can tell "resolved
// upstream" from "blocked locally" without touching the internet.
type testUpstream struct {
	*dns.Server
	addr string
	mu   sync.Mutex
	hits map[string]int
}

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

func (u *testUpstream) handle(w dns.ResponseWriter, q *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(q)
	name := q.Question[0].Name
	u.mu.Lock()
	u.hits[name]++
	u.mu.Unlock()
	if strings.HasSuffix(name, ".test.") {
		if q.Question[0].Qtype == dns.TypeA {
			rr, _ := dns.NewRR(name + " 60 IN A 203.0.113.7")
			m.Answer = append(m.Answer, rr)
		}
		soa, _ := dns.NewRR("test. 60 IN SOA ns.test. hostmaster.test. 1 3600 600 604800 60")
		m.Ns = append(m.Ns, soa)
	} else {
		m.Rcode = dns.RcodeNameError
		soa, _ := dns.NewRR("test. 60 IN SOA ns.test. hostmaster.test. 1 3600 600 604800 60")
		m.Ns = append(m.Ns, soa)
	}
	_ = w.WriteMsg(m)
}

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

// harness is a fully wired App on ephemeral ports, with a fake upstream.
type harness struct {
	t    *testing.T
	app  *App
	addr string // host:port for plain DNS
	api  string // host:port for the admin API
	key  string
	up   *testUpstream
	dir  string
	// config is the document the harness wrote to disk, so a test can make a
	// realistic edit rather than inventing a minimal stub.
	config string
}

// OriginalConfig returns the document the harness wrote to disk, so a test
// can make a realistic single-value edit instead of inventing a minimal stub.
func (h *harness) OriginalConfig() string { return h.config }

func newHarness(t *testing.T, mutate func(*config.Config), opts ...func(*Options)) *harness {
	t.Helper()
	dir := t.TempDir()
	up := newTestUpstream(t)

	hosts := filepath.Join(dir, "hosts")
	if err := os.WriteFile(hosts, []byte(
		"0.0.0.0 ads.example.com\n"+
			"0.0.0.0 tracker.example.net\n"+
			"0.0.0.0 shared.example.com\n"+
			// Under .test. so the fake upstream resolves it: NOERROR then means
			// "the blocklist let it through" and NXDOMAIN means "blocked", with
			// no third possibility to confuse an assertion.
			"0.0.0.0 blocked.example.test\n"+
			"# comment\n"+
			"127.0.0.1 localhost\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	zone := filepath.Join(dir, "local.zone")
	if err := os.WriteFile(zone, []byte(
		"$TTL 60\n"+
			"local. SOA ns.local. hostmaster.local. 1 3600 600 604800 60\n"+
			"local. NS ns.local.\n"+
			"ns.local. A 192.168.1.2\n"+
			"*.wild.local. A 192.168.1.50\n"+
			"deep.svc.local. A 192.168.1.60\n"+
			"*.deep.svc.local. A 192.168.1.61\n"), 0o644); err != nil {
		t.Fatal(err)
	}

	// Pick a free port up front so the client can address the listeners and so
	// a port clash fails the test rather than the app's start-up.
	dnsPort := freePort(t)
	dnsAddr := fmt.Sprintf("127.0.0.1:%d", dnsPort)

	cfg := config.Default()
	cfg.WorkingDirOrSet(dir)
	cfg.DNS.UDP = []string{dnsAddr}
	cfg.DNS.TCP = []string{dnsAddr}
	cfg.DNS.DoT.Enabled = false
	cfg.DNS.DoH.Enabled = false
	cfg.API.Listen = []string{"127.0.0.1:0"}
	cfg.API.Auth.Mode = "apikey"
	cfg.API.Auth.APIKeys = []string{"test-key"}
	cfg.API.Auth.TrustLoopback = false
	cfg.API.UI = false
	cfg.Upstream.Policy = "parallel"
	cfg.Upstream.MaxAttempts = 1
	cfg.Upstream.HealthCheck = 0
	cfg.Upstream.Servers = []config.UpstreamServer{{Addr: up.addr, Proto: "udp", Weight: 1, HealthName: "test"}}
	cfg.Cache.Persist.Path = filepath.Join(dir, "cache.db")
	cfg.Block.Lists = []config.ListSource{{
		Name: "test-hosts", Format: "hosts", Path: hosts, Enabled: true, Refresh: 0,
	}}
	cfg.Block.LocalZones = []string{zone}
	cfg.Block.BlockEntries = []string{"inline-blocked.example.com", "*.inline-sub.example.com"}
	cfg.Block.AllowEntries = []string{"allow.example.com"}
	if mutate != nil {
		mutate(cfg)
	}
	path := filepath.Join(dir, "config.yaml")
	// Snapshot the document before the app starts: the server may rewrite this
	// file (every API mutation does), so a test must not read it concurrently
	// with a live instance's own writes.
	doc := string(mustYAML(t, cfg))
	if err := os.WriteFile(path, []byte(doc), 0o600); err != nil {
		t.Fatal(err)
	}

	// The access log goes to stdout by design; keep it out of the test output
	// unless a test captures it.
	opt := Options{ConfigPath: path, Version: "test", AccessLogOut: io.Discard}
	for _, o := range opts {
		o(&opt)
	}
	a, err := New(opt)
	if err != nil {
		t.Fatalf("new app: %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := a.Start(ctx); err != nil {
		cancel()
		t.Fatalf("start: %v", err)
	}
	t.Cleanup(func() {
		sctx, scancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = a.Stop(sctx)
		scancel()
		cancel()
	})
	return &harness{t: t, app: a, addr: dnsAddr, api: a.Addr(), key: "test-key", up: up, dir: dir, config: doc}
}

// freePort returns a currently-free TCP port on the loopback interface.
func freePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("free port: %v", err)
	}
	port := ln.Addr().(*net.TCPAddr).Port
	_ = ln.Close()
	return port
}

func (h *harness) query(name string, qtype uint16) *dns.Msg {
	h.t.Helper()
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(name), qtype)
	m.RecursionDesired = true
	c := &dns.Client{Net: "udp", Timeout: 5 * time.Second}
	res, _, err := c.Exchange(m, h.addr)
	if err != nil {
		h.t.Fatalf("query %s: %v", name, err)
	}
	return res
}

func (h *harness) get(path string) (int, []byte) {
	h.t.Helper()
	req, _ := http.NewRequest("GET", "http://"+h.api+path, nil)
	req.Header.Set("X-API-Key", h.key)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		h.t.Fatalf("GET %s: %v", path, err)
	}
	defer res.Body.Close()
	b, _ := io.ReadAll(res.Body)
	return res.StatusCode, b
}

func (h *harness) do(method, path, body string) (int, []byte) {
	h.t.Helper()
	var rdr io.Reader
	if body != "" {
		rdr = strings.NewReader(body)
	}
	req, _ := http.NewRequest(method, "http://"+h.api+path, rdr)
	req.Header.Set("X-API-Key", h.key)
	req.Header.Set("Content-Type", "application/json")
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		h.t.Fatalf("%s %s: %v", method, path, err)
	}
	defer res.Body.Close()
	b, _ := io.ReadAll(res.Body)
	return res.StatusCode, b
}

func mustYAML(t *testing.T, c *config.Config) []byte {
	t.Helper()
	b, err := config.Marshal(c)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

var _ = base64.StdEncoding
var _ = tls.VersionTLS13