fancy-dns/internal/cache/store.go

package cache

import (
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"os"
	"path/filepath"
	"sync"
	"time"

	bolt "go.etcd.io/bbolt"
)

var bucketCache = []byte("cache")

// recordVersion is the on-disk record format version. It is written as the
// first byte of every value so a future format change can be recognised (and
// the record dropped) instead of being silently misparsed.
const recordVersion = 1

// Record flags, stored in byte 1 of a record.
const (
	flagNeg       = 1 << 0
	flagTruncated = 1 << 1
)

// errBadRecord is returned by decodeRecord for anything it cannot trust: a
// short buffer, an unknown version, a failed checksum or trailing garbage.
var errBadRecord = errString("cache: corrupt snapshot record")

var crc32cTable = crc32.MakeTable(crc32.Castagnoli)

// encodeRecord serialises one entry into a snapshot value.
//
// The layout is written out here, and not merely implied by the encoder and
// decoder sitting side by side, because the two drifting apart is exactly what
// broke the warm-restart path once already:
//
//	byte  0         format version (recordVersion)
//	byte  1         flags (flagNeg, flagTruncated)
//	uvarint         ExpiresAt as Unix nanoseconds (0 = the zero time)
//	uvarint         StaleUntil as Unix nanoseconds (0 = the zero time)
//	uvarint         CreatedAt as Unix nanoseconds (0 = the zero time)
//	uvarint         TTL in seconds
//	uvarint         len(Origin)   + Origin
//	uvarint         len(Upstream) + Upstream
//	rest            packed wire-format answer
//	uint32 (LE)     CRC-32C of every byte above
func encodeRecord(buf []byte, e *Entry) []byte {
	buf = append(buf, recordVersion)
	var flags uint8
	if e.Neg {
		flags |= flagNeg
	}
	if e.Truncated {
		flags |= flagTruncated
	}
	buf = append(buf, flags)
	buf = binary.AppendUvarint(buf, uint64(nanosOf(e.ExpiresAt)))
	buf = binary.AppendUvarint(buf, uint64(nanosOf(e.StaleUntil)))
	buf = binary.AppendUvarint(buf, uint64(nanosOf(e.CreatedAt)))
	buf = binary.AppendUvarint(buf, uint64(e.TTL))
	buf = binary.AppendUvarint(buf, uint64(len(e.Origin)))
	buf = append(buf, e.Origin...)
	buf = binary.AppendUvarint(buf, uint64(len(e.Upstream)))
	buf = append(buf, e.Upstream...)
	buf = append(buf, e.Wire...)
	return binary.LittleEndian.AppendUint32(buf, crc32.Checksum(buf, crc32cTable))
}

// decodeRecord is the inverse of encodeRecord. Every failure mode returns
// errBadRecord: a snapshot is a cache, not a source of truth, so one damaged
// record must cost one entry and not the whole file.
func decodeRecord(buf []byte) (*Entry, error) {
	const header = 2
	if len(buf) < header+crc32.Size {
		return nil, errBadRecord
	}
	body := buf[:len(buf)-crc32.Size]
	if got, want := binary.LittleEndian.Uint32(buf[len(buf)-crc32.Size:]), crc32.Checksum(body, crc32cTable); got != want {
		return nil, errBadRecord
	}
	if body[0] != recordVersion {
		return nil, errBadRecord
	}
	e := &Entry{Neg: body[1]&flagNeg != 0, Truncated: body[1]&flagTruncated != 0}
	rest := body[header:]

	var n int
	var ok bool
	if e.ExpiresAt, n, ok = timeFromVarint(rest); !ok {
		return nil, errBadRecord
	}
	rest = rest[n:]
	if e.StaleUntil, n, ok = timeFromVarint(rest); !ok {
		return nil, errBadRecord
	}
	rest = rest[n:]
	if e.CreatedAt, n, ok = timeFromVarint(rest); !ok {
		return nil, errBadRecord
	}
	rest = rest[n:]

	ttl, n := binary.Uvarint(rest)
	if n <= 0 {
		return nil, errBadRecord
	}
	e.TTL = uint32(ttl)
	rest = rest[n:]

	for _, dst := range []*string{&e.Origin, &e.Upstream} {
		s, n, ok := readVarintString(rest)
		if !ok {
			return nil, errBadRecord
		}
		*dst = s
		rest = rest[n:]
	}
	if len(rest) == 0 {
		return nil, errBadRecord
	}
	// The buffer belongs to bbolt's mmap, so copy the wire bytes out instead
	// of aliasing memory a later transaction may move.
	e.Wire = append([]byte(nil), rest...)
	e.Bytes = len(e.Wire) + len(e.Origin) + len(e.Upstream) + 96
	return e, nil
}

// originOrSnapshot keeps whatever origin the record carried (the UI shows it)
// but never stores an empty one.
func originOrSnapshot(origin string) string {
	if origin == "" {
		return "snapshot"
	}
	return origin
}

func nanosOf(t time.Time) int64 {
	if t.IsZero() {
		return 0
	}
	return t.UnixNano()
}

func timeFromVarint(buf []byte) (time.Time, int, bool) {
	v, n := binary.Uvarint(buf)
	if n <= 0 {
		return time.Time{}, 0, false
	}
	if v == 0 {
		return time.Time{}, n, true
	}
	return time.Unix(0, int64(v)), n, true
}

func readVarintString(buf []byte) (string, int, bool) {
	l, n := binary.Uvarint(buf)
	if n <= 0 || l > uint64(len(buf)-n) {
		return "", 0, false
	}
	return string(buf[n : n+int(l)]), n + int(l), true
}

// Store persists a snapshot of the LRU to a bbolt file so a restart (or a
// rolling deploy of a cluster) resumes with a warm cache.
//
// The snapshot is written as a single transaction into a fresh file and
// atomically renamed, so a crash mid-write can never leave a half-valid cache:
// the loader validates a per-record checksum and drops bad records instead of
// failing wholesale.
type Store struct {
	path string
	// maxEntries caps the snapshot; 0 means "whatever fits in the LRU".
	maxEntries int

	mu       sync.Mutex
	db       *bolt.DB
	writes   uint64
	errors   uint64
	lastSnap time.Time
	loaded   bool
}

// StoreOptions configures a Store.
type StoreOptions struct {
	// Path is the bbolt file to use.
	Path string
	// MaxEntries caps how many entries a snapshot may contain.
	MaxEntries int
	// Timeout bounds a single snapshot write.
	Timeout time.Duration
}

// OpenStore opens (or creates) the bbolt file. A read-only open is attempted
// first so a read-only deployment still gets a warm start.
func OpenStore(opt StoreOptions) (*Store, error) {
	if opt.Path == "" {
		return nil, fmt.Errorf("cache: store path is required")
	}
	if err := os.MkdirAll(filepath.Dir(opt.Path), 0o755); err != nil {
		return nil, fmt.Errorf("cache: create data dir: %w", err)
	}
	timeout := opt.Timeout
	if timeout <= 0 {
		timeout = 30 * time.Second
	}
	db, err := bolt.Open(opt.Path, 0o600, &bolt.Options{Timeout: timeout, NoFreelistSync: true})
	if err != nil {
		return nil, fmt.Errorf("cache: open %s: %w", opt.Path, err)
	}
	if err := db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(bucketCache)
		return err
	}); err != nil {
		_ = db.Close()
		return nil, err
	}
	return &Store{path: opt.Path, maxEntries: opt.MaxEntries, db: db}, nil
}

// Path is the bbolt file path.
func (s *Store) Path() string { return s.path }

// Load restores entries into the LRU. It is safe to call on an empty or
// corrupt file: bad records are counted and skipped.
func (s *Store) Load(l *LRU, now time.Time) (loaded int, skipped int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	err = s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketCache)
		if b == nil {
			return nil
		}
		cur := b.Cursor()
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			key, _, kerr := unpackKey(k)
			if kerr != nil {
				skipped++
				continue
			}
			e, derr := decodeRecord(v)
			if derr != nil {
				skipped++
				continue
			}
			if !e.ExpiresAt.After(now) {
				// The serve-stale window deliberately does not survive a
				// restart: stale answers bridge a transient upstream failure
				// inside one process lifetime, they are not a licence to
				// resurrect an answer from minutes ago. An entry that is
				// already past its TTL at load time is dropped instead.
				skipped++
				continue
			}
			// Re-anchor on the loader's clock: the stored CreatedAt belongs to
			// the previous process, and AgeTTL would otherwise age the answer
			// by the length of the outage. The TTL is rewritten to the time
			// actually left, so the TTL a client sees never jumps up across a
			// restart.
			e.CreatedAt = now
			e.TTL = uint32(e.ExpiresAt.Sub(now) / time.Second)
			e.Origin = originOrSnapshot(e.Origin)
			l.Put(key, e)
			loaded++
		}
		return nil
	})
	if err != nil {
		s.errors++
		return loaded, skipped, err
	}
	s.loaded = true
	return loaded, skipped, nil
}

// Snapshot writes the current LRU contents to disk. Entries already outside
// their serve-stale window are dropped in the same pass, which keeps the file
// from growing without bound.
func (s *Store) Snapshot(l *LRU, now time.Time) (written int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Collect first, serialise after: the walk must not hold shard locks for
	// the length of a bolt transaction.
	type row struct{ key, val []byte }
	rows := make([]row, 0, 8192)
	dropped := 0
	l.Walk(func(k Key, e *Entry) bool {
		if !e.StaleUntil.After(now) {
			dropped++
			return true
		}
		kb, kerr := packKey(nil, k)
		if kerr != nil {
			dropped++
			return true
		}
		rows = append(rows, row{key: kb, val: encodeRecord(nil, e)})
		return true
	})
	if s.maxEntries > 0 && len(rows) > s.maxEntries {
		dropped += len(rows) - s.maxEntries
		rows = rows[:s.maxEntries]
	}

	err = s.db.Update(func(tx *bolt.Tx) error {
		// Rebuild the bucket so a shrinking cache really shrinks the file.
		if err := tx.DeleteBucket(bucketCache); err != nil && err != bolt.ErrBucketNotFound {
			return err
		}
		b, err := tx.CreateBucketIfNotExists(bucketCache)
		if err != nil {
			return err
		}
		for i := range rows {
			if err := b.Put(rows[i].key, rows[i].val); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		s.errors++
		return 0, err
	}
	s.writes++
	s.lastSnap = now
	return len(rows), nil
}

// Truncate empties the snapshot file.
func (s *Store) Truncate() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	err := s.db.Update(func(tx *bolt.Tx) error {
		if tx.Bucket(bucketCache) == nil {
			return nil
		}
		return tx.DeleteBucket(bucketCache)
	})
	if err != nil {
		s.errors++
	}
	return err
}

// Stats reports the snapshot file's state.
func (s *Store) Stats() PersistSt {
	s.mu.Lock()
	defer s.mu.Unlock()
	st := PersistSt{
		Enabled:      true,
		Path:         s.path,
		LastSnapshot: s.lastSnap,
		Writes:       s.writes,
		Errors:       s.errors,
		Loaded:       s.loaded,
	}
	if fi, err := os.Stat(s.path); err == nil {
		st.Bytes = fi.Size()
	}
	if s.db != nil {
		_ = s.db.View(func(tx *bolt.Tx) error {
			if b := tx.Bucket(bucketCache); b != nil {
				st.Entries = b.Stats().KeyN
			}
			return nil
		})
	}
	return st
}

// Close releases the file handle.
func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.db == nil {
		return nil
	}
	return s.db.Close()
}