initial commit

2026-08-16 21:18:45 -05:00
commit 1e05a01bcf
122 changed files with 29178 additions and 0 deletions
@@ -0,0 +1,690 @@
+// 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
+}