vibedns/internal/cache/cache.go

// Package cache implements the resolver cache: a sharded, LRU-bounded store of
// DNS responses with TTL decay, negative caching, stale serving and prefetch.
//
// The cache lives entirely in memory. It is never persisted, because a cache
// that survives a restart would serve answers whose TTLs it can no longer
// reason about.
package cache

import (
	"container/list"
	"hash/fnv"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"
)

// shardCount must be a power of two.
const shardCount = 64

// Config controls cache behaviour. It is swapped in wholesale on change.
type Config struct {
	Enabled         bool
	MaxEntries      int
	MinTTL          uint32
	MaxTTL          uint32
	NegativeTTL     uint32
	ServeStale      bool
	StaleTTL        uint32
	Prefetch        bool
	PrefetchPercent int
}

// Key identifies a cached response. The DO bit is part of the key because a
// DNSSEC-aware answer carries RRSIG records that a non-DO client must not see.
type Key struct {
	Name  string // lowercase FQDN
	Type  uint16
	Class uint16
	DO    bool
}

// String renders the key in the form shown in the cache browser.
func (k Key) String() string {
	s := k.Name + " " + dns.TypeToString[k.Type]
	if k.DO {
		s += " +dnssec"
	}
	return s
}

func (k Key) hash() uint32 {
	h := fnv.New32a()
	_, _ = h.Write([]byte(k.Name))
	_, _ = h.Write([]byte{byte(k.Type >> 8), byte(k.Type), byte(k.Class >> 8), byte(k.Class)})
	if k.DO {
		_, _ = h.Write([]byte{1})
	}
	return h.Sum32()
}

// KeyFor builds a cache key from a question.
func KeyFor(q dns.Question, do bool) Key {
	return Key{Name: strings.ToLower(dns.Fqdn(q.Name)), Type: q.Qtype, Class: q.Qclass, DO: do}
}

// entry is one cached response.
type entry struct {
	key      Key
	msg      *dns.Msg // stored with original TTLs
	stored   time.Time
	ttl      uint32 // seconds the answer is fresh for
	origTTL  uint32 // TTL at insertion, used for the prefetch threshold
	rcode    int
	size     int
	elem     *list.Element // position in the shard LRU
	negative bool
}

// expiresAt returns the instant the entry stops being fresh.
func (e *entry) expiresAt() time.Time {
	return e.stored.Add(time.Duration(e.ttl) * time.Second)
}

type shard struct {
	mu      sync.RWMutex
	entries map[Key]*entry
	lru     *list.List // front = most recently used
	bytes   int64
}

// Cache is the resolver cache.
type Cache struct {
	shards [shardCount]*shard

	cfgMu sync.RWMutex
	cfg   Config

	hits       atomic.Int64
	misses     atomic.Int64
	staleHits  atomic.Int64
	insertions atomic.Int64
	evictions  atomic.Int64
	expiries   atomic.Int64

	// prefetch is invoked asynchronously when a fresh-but-ageing entry is hit.
	prefetchMu sync.RWMutex
	prefetch   func(Key)
	inflight   sync.Map // Key -> struct{}, dedupes prefetch requests
}

// New creates a cache with the given configuration.
func New(cfg Config) *Cache {
	c := &Cache{}
	for i := range c.shards {
		c.shards[i] = &shard{entries: map[Key]*entry{}, lru: list.New()}
	}
	c.SetConfig(cfg)
	return c
}

// SetConfig replaces the cache configuration. Shrinking MaxEntries evicts down
// to the new bound, and disabling the cache flushes it.
func (c *Cache) SetConfig(cfg Config) {
	if cfg.MaxEntries <= 0 {
		cfg.MaxEntries = 10_000
	}
	if cfg.MaxTTL == 0 {
		cfg.MaxTTL = 86400
	}
	if cfg.PrefetchPercent <= 0 || cfg.PrefetchPercent >= 100 {
		cfg.PrefetchPercent = 10
	}
	c.cfgMu.Lock()
	c.cfg = cfg
	c.cfgMu.Unlock()

	if !cfg.Enabled {
		c.Flush()
		return
	}
	c.enforceBound()
}

// Config returns the current configuration.
func (c *Cache) Config() Config {
	c.cfgMu.RLock()
	defer c.cfgMu.RUnlock()
	return c.cfg
}

// SetPrefetcher registers the callback used to refresh ageing entries.
func (c *Cache) SetPrefetcher(fn func(Key)) {
	c.prefetchMu.Lock()
	c.prefetch = fn
	c.prefetchMu.Unlock()
}

func (c *Cache) shardFor(k Key) *shard {
	return c.shards[k.hash()&(shardCount-1)]
}

// Result describes a cache lookup outcome.
type Result struct {
	Msg    *dns.Msg
	Hit    bool
	Stale  bool
	Age    time.Duration
	Expiry time.Time
}

// Get looks up a response. The returned message is a copy with TTLs decayed by
// the time the entry has spent in the cache, so clients never see a TTL that
// stands still.
func (c *Cache) Get(k Key, req *dns.Msg) Result {
	cfg := c.Config()
	if !cfg.Enabled {
		return Result{}
	}
	sh := c.shardFor(k)

	sh.mu.RLock()
	e, ok := sh.entries[k]
	if !ok {
		sh.mu.RUnlock()
		c.misses.Add(1)
		return Result{}
	}
	stored, ttl, origTTL := e.stored, e.ttl, e.origTTL
	msg := e.msg
	sh.mu.RUnlock()

	age := time.Since(stored)
	elapsed := uint32(age / time.Second)

	switch {
	case elapsed < ttl:
		remaining := ttl - elapsed
		out := decayed(msg, req, remaining, elapsed)
		c.hits.Add(1)
		c.touch(sh, k)
		if cfg.Prefetch && shouldPrefetch(remaining, origTTL, cfg.PrefetchPercent) {
			c.triggerPrefetch(k)
		}
		return Result{Msg: out, Hit: true, Age: age, Expiry: stored.Add(time.Duration(ttl) * time.Second)}

	case cfg.ServeStale && cfg.StaleTTL > 0 && elapsed < ttl+cfg.StaleTTL:
		// RFC 8767: serve the expired answer with a short TTL while a fresh one
		// is fetched, rather than failing the client outright.
		const staleClientTTL = 30
		out := decayed(msg, req, staleClientTTL, elapsed)
		c.staleHits.Add(1)
		c.hits.Add(1)
		c.triggerPrefetch(k)
		return Result{Msg: out, Hit: true, Stale: true, Age: age,
			Expiry: stored.Add(time.Duration(ttl) * time.Second)}

	default:
		c.remove(sh, k)
		c.expiries.Add(1)
		c.misses.Add(1)
		return Result{}
	}
}

func shouldPrefetch(remaining, orig uint32, percent int) bool {
	if orig == 0 {
		return false
	}
	threshold := orig * uint32(percent) / 100
	if threshold < 1 {
		threshold = 1
	}
	return remaining <= threshold
}

func (c *Cache) triggerPrefetch(k Key) {
	c.prefetchMu.RLock()
	fn := c.prefetch
	c.prefetchMu.RUnlock()
	if fn == nil {
		return
	}
	if _, loaded := c.inflight.LoadOrStore(k, struct{}{}); loaded {
		return
	}
	go func() {
		defer c.inflight.Delete(k)
		fn(k)
	}()
}

// decayed copies a stored message for a specific request, reducing every TTL by
// the number of seconds the entry has been cached.
func decayed(stored *dns.Msg, req *dns.Msg, remaining, elapsed uint32) *dns.Msg {
	out := stored.Copy()
	if req != nil {
		out.Id = req.Id
		out.Question = req.Question
		out.RecursionDesired = req.RecursionDesired
	}
	adjust := func(rrs []dns.RR) {
		for _, rr := range rrs {
			if rr.Header().Rrtype == dns.TypeOPT {
				continue
			}
			t := rr.Header().Ttl
			if t <= elapsed {
				rr.Header().Ttl = remaining
				continue
			}
			nt := t - elapsed
			if nt < 1 {
				nt = 1
			}
			rr.Header().Ttl = nt
		}
	}
	adjust(out.Answer)
	adjust(out.Ns)
	adjust(out.Extra)
	return out
}

// Put stores a response. It returns the TTL the entry was stored with, or 0 if
// the response was not cacheable.
func (c *Cache) Put(k Key, msg *dns.Msg) uint32 {
	cfg := c.Config()
	if !cfg.Enabled || msg == nil {
		return 0
	}
	if !cacheable(msg) {
		return 0
	}

	negative := isNegative(msg)
	ttl := responseTTL(msg, negative, cfg)
	if ttl == 0 {
		return 0
	}

	stored := msg.Copy()
	// The OPT record describes the transport of one exchange, not the data, so
	// it must not be replayed to a different client.
	stored.Extra = stripOPT(stored.Extra)
	stored.Id = 0

	e := &entry{
		key:      k,
		msg:      stored,
		stored:   time.Now(),
		ttl:      ttl,
		origTTL:  ttl,
		rcode:    msg.Rcode,
		size:     estimateSize(k, stored),
		negative: negative,
	}

	sh := c.shardFor(k)
	sh.mu.Lock()
	if old, ok := sh.entries[k]; ok {
		sh.lru.Remove(old.elem)
		sh.bytes -= int64(old.size)
	}
	e.elem = sh.lru.PushFront(k)
	sh.entries[k] = e
	sh.bytes += int64(e.size)
	sh.mu.Unlock()

	c.insertions.Add(1)
	c.enforceBound()
	return ttl
}

// cacheable rejects responses that must never be reused.
func cacheable(msg *dns.Msg) bool {
	if msg.Truncated {
		return false
	}
	switch msg.Rcode {
	case dns.RcodeSuccess, dns.RcodeNameError:
		return true
	default:
		// SERVFAIL, REFUSED and friends are transient or client specific.
		return false
	}
}

func isNegative(msg *dns.Msg) bool {
	return msg.Rcode == dns.RcodeNameError || len(msg.Answer) == 0
}

// responseTTL derives the cache lifetime from the response, clamped to the
// configured bounds. Negative answers use the SOA MINIMUM per RFC 2308.
func responseTTL(msg *dns.Msg, negative bool, cfg Config) uint32 {
	if negative {
		ttl := cfg.NegativeTTL
		if soa := findSOA(msg.Ns); soa != nil {
			t := soa.Minttl
			if soa.Hdr.Ttl < t {
				t = soa.Hdr.Ttl
			}
			if t < ttl || ttl == 0 {
				ttl = t
			}
		}
		if ttl == 0 {
			return 0
		}
		return clampTTL(ttl, cfg)
	}

	ttl := uint32(0)
	first := true
	for _, section := range [][]dns.RR{msg.Answer, msg.Ns} {
		for _, rr := range section {
			if rr.Header().Rrtype == dns.TypeOPT {
				continue
			}
			t := rr.Header().Ttl
			if first || t < ttl {
				ttl = t
				first = false
			}
		}
	}
	if first {
		return 0 // nothing with a TTL to key off
	}
	return clampTTL(ttl, cfg)
}

func clampTTL(ttl uint32, cfg Config) uint32 {
	if cfg.MinTTL > 0 && ttl < cfg.MinTTL {
		ttl = cfg.MinTTL
	}
	if cfg.MaxTTL > 0 && ttl > cfg.MaxTTL {
		ttl = cfg.MaxTTL
	}
	return ttl
}

func findSOA(rrs []dns.RR) *dns.SOA {
	for _, rr := range rrs {
		if soa, ok := rr.(*dns.SOA); ok {
			return soa
		}
	}
	return nil
}

func stripOPT(rrs []dns.RR) []dns.RR {
	out := rrs[:0]
	for _, rr := range rrs {
		if rr.Header().Rrtype == dns.TypeOPT {
			continue
		}
		out = append(out, rr)
	}
	return out
}

// estimateSize approximates the heap cost of an entry, for the memory readout.
func estimateSize(k Key, msg *dns.Msg) int {
	const entryOverhead = 160 // struct, map bucket and list element
	return entryOverhead + len(k.Name) + msg.Len()
}

func (c *Cache) touch(sh *shard, k Key) {
	sh.mu.Lock()
	if e, ok := sh.entries[k]; ok && e.elem != nil {
		sh.lru.MoveToFront(e.elem)
	}
	sh.mu.Unlock()
}

func (c *Cache) remove(sh *shard, k Key) {
	sh.mu.Lock()
	if e, ok := sh.entries[k]; ok {
		if e.elem != nil {
			sh.lru.Remove(e.elem)
		}
		sh.bytes -= int64(e.size)
		delete(sh.entries, k)
	}
	sh.mu.Unlock()
}

// Delete removes one entry. It reports whether the entry was present.
func (c *Cache) Delete(k Key) bool {
	sh := c.shardFor(k)
	sh.mu.Lock()
	defer sh.mu.Unlock()
	e, ok := sh.entries[k]
	if !ok {
		return false
	}
	if e.elem != nil {
		sh.lru.Remove(e.elem)
	}
	sh.bytes -= int64(e.size)
	delete(sh.entries, k)
	return true
}

// Flush empties the cache and returns how many entries were dropped.
func (c *Cache) Flush() int {
	total := 0
	for _, sh := range c.shards {
		sh.mu.Lock()
		total += len(sh.entries)
		sh.entries = map[Key]*entry{}
		sh.lru.Init()
		sh.bytes = 0
		sh.mu.Unlock()
	}
	return total
}

// FlushName removes every entry for one name, across all types.
func (c *Cache) FlushName(name string) int {
	name = strings.ToLower(dns.Fqdn(name))
	removed := 0
	for _, sh := range c.shards {
		sh.mu.Lock()
		for k, e := range sh.entries {
			if k.Name == name {
				if e.elem != nil {
					sh.lru.Remove(e.elem)
				}
				sh.bytes -= int64(e.size)
				delete(sh.entries, k)
				removed++
			}
		}
		sh.mu.Unlock()
	}
	return removed
}

// enforceBound evicts least-recently-used entries until the cache fits.
//
// The bound is applied per shard so that eviction never has to lock the whole
// cache at once.
func (c *Cache) enforceBound() {
	cfg := c.Config()
	if cfg.MaxEntries <= 0 {
		return
	}
	perShard := cfg.MaxEntries / shardCount
	if perShard < 1 {
		perShard = 1
	}
	for _, sh := range c.shards {
		sh.mu.Lock()
		for len(sh.entries) > perShard {
			back := sh.lru.Back()
			if back == nil {
				break
			}
			k := back.Value.(Key)
			if e, ok := sh.entries[k]; ok {
				sh.bytes -= int64(e.size)
				delete(sh.entries, k)
			}
			sh.lru.Remove(back)
			c.evictions.Add(1)
		}
		sh.mu.Unlock()
	}
}

// Cleanup drops entries that are past both their TTL and their stale window.
// It returns the number removed.
func (c *Cache) Cleanup() int {
	cfg := c.Config()
	grace := time.Duration(0)
	if cfg.ServeStale {
		grace = time.Duration(cfg.StaleTTL) * time.Second
	}
	now := time.Now()
	removed := 0
	for _, sh := range c.shards {
		sh.mu.Lock()
		for k, e := range sh.entries {
			if now.After(e.expiresAt().Add(grace)) {
				if e.elem != nil {
					sh.lru.Remove(e.elem)
				}
				sh.bytes -= int64(e.size)
				delete(sh.entries, k)
				removed++
			}
		}
		sh.mu.Unlock()
	}
	c.expiries.Add(int64(removed))
	return removed
}

// Run starts the periodic cleanup loop. It returns when done is closed.
func (c *Cache) Run(done <-chan struct{}, interval func() time.Duration) {
	for {
		d := interval()
		if d <= 0 {
			d = time.Minute
		}
		select {
		case <-done:
			return
		case <-time.After(d):
			c.Cleanup()
		}
	}
}

// Stats is a snapshot of cache counters for the dashboard and metrics.
type Stats struct {
	Enabled     bool    `json:"enabled"`
	Entries     int     `json:"entries"`
	MaxEntries  int     `json:"max_entries"`
	Bytes       int64   `json:"bytes"`
	Hits        int64   `json:"hits"`
	Misses      int64   `json:"misses"`
	StaleHits   int64   `json:"stale_hits"`
	Insertions  int64   `json:"insertions"`
	Evictions   int64   `json:"evictions"`
	Expirations int64   `json:"expirations"`
	HitRate     float64 `json:"hit_rate"`
}

// Stats returns the current counters.
func (c *Cache) Stats() Stats {
	s := Stats{
		Enabled:     c.Config().Enabled,
		MaxEntries:  c.Config().MaxEntries,
		Hits:        c.hits.Load(),
		Misses:      c.misses.Load(),
		StaleHits:   c.staleHits.Load(),
		Insertions:  c.insertions.Load(),
		Evictions:   c.evictions.Load(),
		Expirations: c.expiries.Load(),
	}
	for _, sh := range c.shards {
		sh.mu.RLock()
		s.Entries += len(sh.entries)
		s.Bytes += sh.bytes
		sh.mu.RUnlock()
	}
	if total := s.Hits + s.Misses; total > 0 {
		s.HitRate = float64(s.Hits) / float64(total) * 100
	}
	return s
}

// ResetStats zeroes the counters without touching the cached data.
func (c *Cache) ResetStats() {
	c.hits.Store(0)
	c.misses.Store(0)
	c.staleHits.Store(0)
	c.insertions.Store(0)
	c.evictions.Store(0)
	c.expiries.Store(0)
}

// EntryView describes one cached entry for the cache browser.
type EntryView struct {
	Name     string    `json:"name"`
	Type     string    `json:"type"`
	DO       bool      `json:"dnssec"`
	Rcode    string    `json:"rcode"`
	Answers  int       `json:"answers"`
	TTL      int64     `json:"ttl"`
	Stored   time.Time `json:"stored"`
	Expires  time.Time `json:"expires"`
	Size     int       `json:"size"`
	Negative bool      `json:"negative"`
	Stale    bool      `json:"stale"`
}

// Entries returns cached entries matching a substring of the name, newest
// first, capped at limit. It also returns the total number of matches.
func (c *Cache) Entries(search string, limit, offset int) ([]EntryView, int) {
	search = strings.ToLower(strings.TrimSpace(search))
	now := time.Now()

	var out []EntryView
	for _, sh := range c.shards {
		sh.mu.RLock()
		for k, e := range sh.entries {
			if search != "" && !strings.Contains(k.Name, search) {
				continue
			}
			expires := e.expiresAt()
			out = append(out, EntryView{
				Name:     k.Name,
				Type:     dns.TypeToString[k.Type],
				DO:       k.DO,
				Rcode:    dns.RcodeToString[e.rcode],
				Answers:  len(e.msg.Answer),
				TTL:      int64(expires.Sub(now) / time.Second),
				Stored:   e.stored,
				Expires:  expires,
				Size:     e.size,
				Negative: e.negative,
				Stale:    now.After(expires),
			})
		}
		sh.mu.RUnlock()
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Name != out[j].Name {
			return out[i].Name < out[j].Name
		}
		return out[i].Type < out[j].Type
	})

	total := len(out)
	if offset > total {
		offset = total
	}
	out = out[offset:]
	if limit > 0 && len(out) > limit {
		out = out[:limit]
	}
	return out, total
}