fancy-dns/internal/cache/manager.go

package cache

import (
	"math"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"
)

var (
	errKeyTooLong = errString("cache: name too long to serialise")
	errBadKey     = errString("cache: corrupt persisted key")
)

type errString string

func (e errString) Error() string { return string(e) }

// Options configures a Manager.
type Options struct {
	// Size is the global entry budget for the in-memory LRU.
	Size int
	// MaxBytes bounds the same LRU by accounted bytes. The entry budget alone
	// does not bound the footprint: one cached ANY or TXT answer can be tens
	// of kilobytes, so 500k entries is anywhere from 150MB to several GB
	// depending entirely on what the traffic asked for. 0 means unbounded.
	MaxBytes int64
	// Shards is the LRU shard count (rounded to a power of two).
	Shards int
	// MinTTL and MaxTTL clamp the TTLs taken from upstream answers.
	MinTTL time.Duration
	MaxTTL time.Duration
	// NegativeTTL bounds cached NXDOMAIN/NODATA answers.
	NegativeTTL time.Duration
	// ServeStale keeps expired entries answerable for this long.
	ServeStale time.Duration
	// Prefetch refreshes stale entries in the background when a hook is set.
	Prefetch bool
	// Now is injectable for tests.
	Now func() time.Time
}

// Manager is the cache front door: TTL policy, serve-stale, prefetching and
// the persisted snapshot all live here, so the DNS handler stays a straight
// pipeline.
type Manager struct {
	lru *LRU
	opt Options
	now func() time.Time

	// set holds a snapshot-aware view: the LRU plus the persisted store.
	store *Store

	// prefetchFn is called (off the hot path) when a stale entry is served.
	// It is set once at wiring time via SetPrefetcher.
	prefetchFn atomic.Pointer[prefetcher]

	// top holds a rolling per-domain hit counter for the "top cached" UI.
	top *topDomains

	// Blocked counts blocklist hits served from the matcher, kept here so the
	// dashboard can show block volume without a second counter package.
	blockedSince atomic.Uint64

	mu       sync.Mutex
	flushing bool
}

type prefetcher func(Key)

// NewManager builds a Manager. store may be nil (persistence disabled).
func NewManager(opt Options, store *Store) *Manager {
	if opt.Now == nil {
		opt.Now = time.Now
	}
	if opt.Size <= 0 {
		opt.Size = 200000
	}
	m := &Manager{
		lru:   NewLRU(opt.Size, opt.Shards),
		opt:   opt,
		now:   opt.Now,
		store: store,
		top:   newTopDomains(5000),
	}
	// Applied after construction rather than through NewLRU so the LRU keeps
	// its two-argument signature; a 0 budget means unbounded.
	m.lru.SetMaxBytes(opt.MaxBytes)
	return m
}

// SetPrefetcher installs the callback invoked when a stale answer is served.
func (m *Manager) SetPrefetcher(fn func(Key)) {
	if fn == nil {
		m.prefetchFn.Store(nil)
		return
	}
	m.prefetchFn.Store((*prefetcher)(&fn))
}

// LRU exposes the underlying LRU for stats and tests.
func (m *Manager) LRU() *LRU { return m.lru }

// Now returns the manager's clock.
func (m *Manager) Now() time.Time { return m.now() }

// Result describes a cache hit.
type Result struct {
	// Kind is Fresh, Stale or Miss.
	Kind HitKind
	// Msg is a private unpacked copy for Fresh and Stale hits.
	Msg *dns.Msg
	// TTL is the TTL the answer was served with (already aged).
	TTL uint32
	// Origin and Upstream are carried through from the stored entry.
	Origin   string
	Upstream string
	// Age is how long ago the entry was stored.
	Age time.Duration
	// ServedStale is set when a stale entry was used; callers may want to
	// report the degraded state.
	ServedStale bool
}

// Get looks up a key. On a Fresh hit it returns a private, TTL-aged copy.
// On a Stale hit it returns the same plus kicks off a background refresh when
// prefetching is enabled.
func (m *Manager) Get(k Key) Result {
	now := m.now()
	e, kind := m.lru.Get(k, now)
	if e == nil {
		return Result{Kind: Miss}
	}
	msg, err := e.unpack(k, now)
	if err != nil {
		// A corrupt entry poisons the cache for this key; drop it.
		m.lru.Delete(k)
		return Result{Kind: Miss}
	}
	m.top.observe(k, e)
	res := Result{
		Kind:        kind,
		Msg:         msg,
		TTL:         e.AgeTTL(now),
		Origin:      e.Origin,
		Upstream:    e.Upstream,
		Age:         now.Sub(e.CreatedAt),
		ServedStale: kind == Stale,
	}
	if kind == Stale {
		res.TTL = 0
		if p := m.prefetchFn.Load(); p != nil && m.opt.Prefetch && m.lru.TryPrefetch(k) {
			fn := *p
			go func() {
				defer m.lru.ReleasePrefetch(k)
				fn(k)
			}()
		}
	}
	return res
}

// unpack rebuilds a private message from stored wire bytes and ages the TTLs
// down to what is left of the entry's lifetime.
func (e *Entry) unpack(k Key, now time.Time) (*dns.Msg, error) {
	msg := new(dns.Msg)
	if err := msg.Unpack(e.Wire); err != nil {
		return nil, err
	}
	remaining := e.AgeTTL(now)
	ageAll(msg, remaining)
	// The stored question is the canonical one; trust the key, not the wire.
	if len(msg.Question) == 0 || !equalName(msg.Question[0].Name, k.Name) {
		msg.Question = []dns.Question{{
			Name:   dns.Fqdn(k.Name),
			Qtype:  k.Type,
			Qclass: k.Class,
		}}
	}
	return msg, nil
}

func equalName(a, b string) bool {
	if a == b {
		return true
	}
	if n := len(a); n > 0 && a[n-1] == '.' {
		a = a[:n-1]
	}
	if n := len(b); n > 0 && b[n-1] == '.' {
		b = b[:n-1]
	}
	return a == b
}

// ageAll sets the TTL of every record in msg to ttl.
func ageAll(msg *dns.Msg, ttl uint32) {
	for i := range msg.Answer {
		msg.Answer[i].Header().Ttl = ttl
	}
	for i := range msg.Ns {
		msg.Ns[i].Header().Ttl = ttl
	}
	for i := range msg.Extra {
		// OPT records carry no TTL semantics; leave the flags alone.
		if _, ok := msg.Extra[i].(*dns.OPT); ok {
			continue
		}
		msg.Extra[i].Header().Ttl = ttl
	}
}

// PutResponse stores an upstream response under k, deriving the TTL from the
// answer itself and clamping it to the configured bounds.
//
// The message is packed once; the stored bytes are independent of the caller's
// copy, so the caller keeps ownership of its message.
func (m *Manager) PutResponse(k Key, msg *dns.Msg, upstream string) {
	ttl, neg := m.TTLFor(msg)
	if ttl <= 0 {
		// A zero TTL means "do not cache"; respect that rather than pinning an
		// uncacheable answer in memory.
		return
	}
	wire, err := msg.Pack()
	if err != nil {
		return
	}
	m.putWire(k, wire, ttl, "upstream", upstream, neg)
}

// PutLocal stores a synthesised answer (local zone, blocklist response).
func (m *Manager) PutLocal(k Key, msg *dns.Msg, origin string) {
	ttl, neg := m.TTLFor(msg)
	if ttl <= 0 {
		ttl = 60
	}
	wire, err := msg.Pack()
	if err != nil {
		return
	}
	m.putWire(k, wire, ttl, origin, "", neg)
}

// PutRaw stores pre-packed wire bytes with an explicit TTL. This is the path
// used when restoring a snapshot, where packing again would be wasted work.
func (m *Manager) PutRaw(k Key, wire []byte, ttl uint32, origin, upstream string, neg bool) {
	m.putWire(k, wire, ttl, origin, upstream, neg)
}

func (m *Manager) putWire(k Key, wire []byte, ttl uint32, origin, upstream string, neg bool) {
	if m == nil || len(wire) == 0 || ttl == 0 {
		return
	}
	now := m.now()
	buf := make([]byte, len(wire))
	copy(buf, wire)
	// Record the insert too, not just the hits: a name queried once must show
	// up in "top cached names", and a miss-only workload has no hits to count.
	m.top.observe(k, &Entry{TTL: ttl, CreatedAt: now})
	// Key bytes, the wire body and the entry struct: account for all three so
	// the LRU budget reflects real memory.
	e := &Entry{
		Wire:       buf,
		TTL:        ttl,
		Origin:     origin,
		Upstream:   upstream,
		CreatedAt:  now,
		ExpiresAt:  now.Add(time.Duration(ttl) * time.Second),
		StaleUntil: now.Add(time.Duration(ttl)*time.Second + m.serveStaleWindow()),
		Bytes:      len(wire) + len(k.Name) + 96,
		Neg:        neg,
	}
	m.lru.Put(k, e)
}

func (m *Manager) serveStaleWindow() time.Duration {
	if m.opt.ServeStale < 0 {
		return 0
	}
	return m.opt.ServeStale
}

// TTLFor derives a cache TTL from a response, honouring the negative-cache
// policy. The second return reports whether the answer is a negative one
// (NXDOMAIN or NODATA), which the UI and the log show differently.
func (m *Manager) TTLFor(msg *dns.Msg) (ttl uint32, neg bool) {
	if msg == nil {
		return 0, false
	}
	if msg.Truncated {
		// Never cache a truncated answer: the client will retry over TCP and
		// we would serve it the truncated form forever.
		return 0, false
	}
	switch msg.Rcode {
	case dns.RcodeNameError:
		return m.negativeTTL(msg), true
	case dns.RcodeSuccess, dns.RcodeNotImplemented:
		// fall through to the general case
	default:
		// SERVFAIL, REFUSED, NOTAUTH and friends are transient by nature.
		return 0, true
	}

	if len(msg.Answer) == 0 {
		// NODATA: the authority section's SOA carries the negative TTL.
		return m.negativeTTL(msg), true
	}
	minTTL := uint32(math.MaxUint32)
	found := false
	for _, rr := range msg.Answer {
		switch rr.(type) {
		case *dns.OPT:
			continue
		}
		if t := rr.Header().Ttl; t < minTTL {
			minTTL = t
			found = true
		}
	}
	if !found {
		return 0, true
	}
	return m.clamp(minTTL), false
}

// soaOf finds the zone's start-of-authority record, which is where a
// negative answer's TTL comes from (RFC 2308).
func soaOf(msg *dns.Msg) *dns.SOA {
	for _, rr := range msg.Ns {
		if soa, ok := rr.(*dns.SOA); ok {
			return soa
		}
	}
	for _, rr := range msg.Answer {
		if soa, ok := rr.(*dns.SOA); ok {
			return soa
		}
	}
	return nil
}

func (m *Manager) negativeTTL(msg *dns.Msg) uint32 {
	soa := soaOf(msg)
	ttl := uint32(0)
	if soa != nil {
		ttl = soa.Hdr.Ttl
		if soa.Minttl < ttl {
			ttl = soa.Minttl
		}
	}
	if ttl == 0 {
		// No SOA to go on: fall back to the shortest TTL in the message, then
		// to the configured default.
		for _, rr := range msg.Ns {
			if t := rr.Header().Ttl; t > 0 && t < ttl {
				ttl = t
			}
		}
		if ttl == 0 {
			ttl = uint32(m.opt.NegativeTTL / time.Second)
		}
	}
	return m.clamp(ttl)
}

func (m *Manager) clamp(ttl uint32) uint32 {
	if min := uint32(m.opt.MinTTL / time.Second); min > 0 && ttl < min {
		ttl = min
	}
	if max := uint32(m.opt.MaxTTL / time.Second); max > 0 && ttl > max {
		ttl = max
	}
	return ttl
}

// Delete drops one key.
func (m *Manager) Delete(k Key) { m.lru.Delete(k) }

// Flush empties the cache and, when persistence is enabled, removes the
// snapshot so a restart does not resurrect it.
func (m *Manager) Flush() {
	m.lru.Flush()
	m.top.reset()
	if m.store != nil {
		_ = m.store.Truncate()
	}
}

// DeleteName drops every entry for a name, across all types. The blocklist API
// calls it when an operator edits rules, so a newly blocked name does not keep
// serving a cached good answer until it expires.
func (m *Manager) DeleteName(name string) int {
	want := canonical(name)
	return m.lru.DeleteWhere(func(k Key, _ *Entry) bool { return k.Name == want })
}

// DeleteSubtree drops every entry at or below a name suffix. Used when a
// whole blocklist is replaced, so stale good answers for newly-listed domains
// disappear immediately.
func (m *Manager) DeleteSubtree(suffix string) int {
	want := canonical(suffix)
	return m.lru.DeleteWhere(func(k Key, _ *Entry) bool {
		return k.Name == want || hasSuffix(k.Name, want)
	})
}

func hasSuffix(name, suffix string) bool {
	if len(name) < len(suffix)+1 {
		return false
	}
	return name[len(name)-len(suffix)-1:] == "."+suffix
}

// Resize applies a new capacity, e.g. after a config reload.
func (m *Manager) Resize(size, shards int) {
	m.opt.Size = size
	m.lru.SetCapacity(size, shards)
}

// Options returns the manager's current options.
func (m *Manager) Options() Options { return m.opt }

// TopEntries returns the hottest names in the cache for the UI.
func (m *Manager) TopEntries(limit int) []TopEntry {
	return m.top.top(limit)
}

// Stats is the JSON-friendly cache snapshot.
type Stats struct {
	Enabled     bool      `json:"enabled"`
	Size        int       `json:"size"`
	Entries     int       `json:"entries"`
	Shards      int       `json:"shards"`
	MinTTLMS    int64     `json:"min_ttl_ms"`
	MaxTTLMS    int64     `json:"max_ttl_ms"`
	ServeStaleM int64     `json:"serve_stale_ms"`
	NegativeMS  int64     `json:"negative_ttl_ms"`
	HitRate     float64   `json:"hit_rate"`
	Hits        uint64    `json:"hits"`
	Misses      uint64    `json:"misses"`
	StaleHits   uint64    `json:"stale_hits"`
	Inserts     uint64    `json:"inserts"`
	Evictions   uint64    `json:"evictions"`
	Expired     uint64    `json:"expired"`
	Bytes       int64     `json:"bytes"`
	MemoryHuman string    `json:"bytes_human"`
	Persist     PersistSt `json:"persisted"`
}

// PersistSt describes the snapshot file.
type PersistSt struct {
	Enabled      bool      `json:"enabled"`
	Path         string    `json:"path"`
	Bytes        int64     `json:"bytes"`
	Entries      int       `json:"entries"`
	LastSnapshot time.Time `json:"last_snapshot"`
	Writes       uint64    `json:"writes"`
	Errors       uint64    `json:"errors"`
	Loaded       bool      `json:"loaded"`
}

// Stats snapshots the manager.
func (m *Manager) Stats() Stats {
	s := m.lru.Stats()
	st := Stats{
		Enabled:     true,
		Size:        m.opt.Size,
		Entries:     s.Entries,
		Shards:      s.Shards,
		MinTTLMS:    m.opt.MinTTL.Milliseconds(),
		MaxTTLMS:    m.opt.MaxTTL.Milliseconds(),
		ServeStaleM: m.opt.ServeStale.Milliseconds(),
		NegativeMS:  m.opt.NegativeTTL.Milliseconds(),
		HitRate:     s.HitRate,
		Hits:        s.Hits,
		Misses:      s.Misses,
		StaleHits:   s.StaleHits,
		Inserts:     s.Inserts,
		Evictions:   s.Evictions,
		Expired:     s.Expired,
		Bytes:       s.Bytes,
	}
	st.MemoryHuman = HumanBytes(s.Bytes)
	if m.store != nil {
		st.Persist = m.store.Stats()
	} else {
		st.Persist = PersistSt{}
	}
	return st
}

// HumanBytes formats a byte count the way the UI shows it.
func HumanBytes(n int64) string {
	neg := ""
	if n < 0 {
		neg, n = "-", -n
	}
	switch {
	case n < 1024:
		return neg + strconv.FormatInt(n, 10) + " B"
	case n < 1024*1024:
		return neg + strconv.FormatFloat(float64(n)/1024, 'f', 1, 64) + " KiB"
	case n < 1024*1024*1024:
		return neg + strconv.FormatFloat(float64(n)/(1024*1024), 'f', 1, 64) + " MiB"
	default:
		return neg + strconv.FormatFloat(float64(n)/(1024*1024*1024), 'f', 2, 64) + " GiB"
	}
}

// topDomains is a bounded, periodic top-N counter. It uses a compact
// min-heap-free approach: a plain map plus periodic pruning, which is far
// cheaper than maintaining a sorted structure on every query.
type topDomains struct {
	mu     sync.Mutex
	counts map[string]*topCount
	limit  int
}

type topCount struct {
	name  string
	qtype uint16
	hits  uint64
	bytes int64
}

func newTopDomains(limit int) *topDomains {
	return &topDomains{counts: make(map[string]*topCount, 4096), limit: limit}
}

func (t *topDomains) observe(k Key, e *Entry) {
	if t == nil {
		return
	}
	id := k.Name + "|" + strconv.Itoa(int(k.Type))
	t.mu.Lock()
	c, ok := t.counts[id]
	if !ok {
		if len(t.counts) > t.limit*2 {
			t.pruneLocked()
		}
		c = &topCount{name: k.Name, qtype: k.Type}
		t.counts[id] = c
	}
	c.hits++
	c.bytes += int64(len(e.Wire))
	t.mu.Unlock()
}

func (t *topDomains) pruneLocked() {
	// Drop the bottom half so pruning is amortised.
	type kv struct {
		id string
		c  *topCount
	}
	all := make([]kv, 0, len(t.counts))
	for id, c := range t.counts {
		all = append(all, kv{id, c})
	}
	sort.Slice(all, func(i, j int) bool { return all[i].c.hits < all[j].c.hits })
	keep := len(all) / 2
	for _, e := range all[keep:] {
		delete(t.counts, e.id)
	}
}

func (t *topDomains) reset() {
	t.mu.Lock()
	t.counts = make(map[string]*topCount, 4096)
	t.mu.Unlock()
}

// TopEntry is one row of the "top cached names" table.
type TopEntry struct {
	Name  string `json:"name"`
	Type  uint16 `json:"qtype"`
	Hits  uint64 `json:"hits"`
	Bytes int64  `json:"bytes"`
}

func (t *topDomains) top(limit int) []TopEntry {
	if t == nil {
		return nil
	}
	if limit <= 0 || limit > 500 {
		limit = 50
	}
	t.mu.Lock()
	out := make([]TopEntry, 0, limit)
	for _, c := range t.counts {
		out = append(out, TopEntry{Name: c.name, Type: c.qtype, Hits: c.hits, Bytes: c.bytes})
	}
	t.mu.Unlock()
	sort.Slice(out, func(i, j int) bool { return out[i].Hits > out[j].Hits })
	if len(out) > limit {
		out = out[:limit]
	}
	return out
}