path: root/internal/cache/cache.go

package cache

import (
	"database/sql"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"codeberg.org/miekg/dns"
	_ "modernc.org/sqlite"
)

type Key struct {
	Name string
	Qtype uint16
	Class uint16
}

type entry struct {
	msg *dns.Msg
	storedAt time.Time
	ttl time.Duration
}

func (e *entry) expired() bool  {
	return time.Since(e.storedAt) >= e.ttl
}

func (e *entry) remaining() time.Duration {
	d := e.ttl - time.Since(e.storedAt)
	if d < 0 {
		return 0
	}
	return d
}

type Cache struct {
	mu sync.RWMutex
	entries map[Key]*entry
	maxSize int
	db *sql.DB
	dbCh chan dbWrite
	wg sync.WaitGroup
	hits int64
	misses int64
	evicted int64
	stopCh chan struct{}
}
type dbWrite struct {
	key Key
	e *entry
}

func NewCache(maxSize int, dbPath string) (*Cache, error) {
	if maxSize <= 0 {
		maxSize = 100000
	}
	c := &Cache{
		entries: make(map[Key]*entry),
		maxSize: maxSize,
		stopCh: make(chan struct{}),
	}

	if dbPath != "" {
		db, err := sql.Open("sqlite", dbPath)
		if err != nil {
			return nil, err
		}
		c.db = db

		db.Exec("PRAGMA journal_mode=WAL")
		db.Exec("PRAGMA synchronous=NORMAL")
		db.Exec("PRAGMA cache_size=-65536")

		if _, err := db.Exec(`
		CREATE TABLE IF NOT EXISTS cache (
			name TEXT NOT NULL,
			qtype INTEGER NOT NULL,
			class INTEGER NOT NULL,
			data BLOB NOT NULL,
			stored_at INTEGER NOT NULL,
			ttl_ns INTEGER NOT NULL,
			PRIMARY KEY (name, qtype, class)
		)
		`); err != nil {
			db.Close()
			return nil, err
		}

		c.loadFromDB()
		c.dbCh = make(chan dbWrite, 1024)
		c.wg.Add(1)
		go c.dbWriter()
	}
	go c.evictLoop()
	return c, nil
}

func (c *Cache) Stop() {
	close(c.stopCh)
	if c.db != nil {
		close(c.dbCh)
	}
	c.wg.Wait()
	if c.db != nil {
		c.db.Close()
	}
}

func (c *Cache) Get(key Key) (*dns.Msg, bool) {
	c.mu.RLock()
	e, ok := c.entries[key]
	if !ok || e.expired() {
		c.mu.RUnlock()
		atomic.AddInt64(&c.misses, 1)
		return nil, false
	}
	c.mu.RUnlock()

	atomic.AddInt64(&c.hits, 1)
	msg := deepCopyMsg(e.msg)
	remaining := e.remaining()
	adjustTTL(msg, e.ttl, remaining)
	return msg, true
}

func (c *Cache) Set(key Key, msg *dns.Msg, ttl time.Duration) {
	if ttl <= 0 {
		ttl = computeTTL(msg)
	}
	if ttl <= 0 {
		ttl = 60 * time.Second
	}

	e := &entry{
		msg: deepCopyMsg(msg),
		storedAt: time.Now(),
		ttl: ttl,
	}

	c.mu.Lock()
	if len(c.entries) >= c.maxSize {
		c.evictLocked()
	}
	c.entries[key] = e
	c.mu.Unlock()

	if c.db != nil {
		select {
		case c.dbCh <- dbWrite{key: key, e: e}:
		default:
		}
	}
}

func (c *Cache) Stats() (int64, int64, int64) {
	return atomic.LoadInt64(&c.hits),
	atomic.LoadInt64(&c.misses),
	atomic.LoadInt64(&c.evicted)
}

func (c *Cache) Len() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.entries)
}

func (c *Cache) evictLocked() {
	now := time.Now()
	for k, e := range c.entries {
		if now.After(e.storedAt.Add(e.ttl)) {
			delete(c.entries, k)
			atomic.AddInt64(&c.evicted, 1)
		}
	}
	if len(c.entries) < c.maxSize {
		return
	}

	var oldestKey Key
	var oldestTime time.Time
	first := true
	for k, e := range c.entries {
		if first || e.storedAt.Before(oldestTime) {
			oldestKey = k
			oldestTime = e.storedAt
			first = false
		}
	}
	if !first {
		delete(c.entries, oldestKey)
		atomic.AddInt64(&c.evicted, 1)
	}
}

func (c *Cache) evictLoop() {
	tk := time.NewTicker(1 * time.Minute)
	defer tk.Stop()
	for {
		select {
		case <-tk.C:
			now := time.Now()
			c.mu.RLock()
			var expired []Key
			for k, e := range c.entries {
				if now.After(e.storedAt.Add(e.ttl)) {
					expired = append(expired, k)
				}
			}
			c.mu.RUnlock()
			if len(expired) > 0 {
				c.mu.Lock()
				for _, k := range expired {
					if _, ok := c.entries[k]; ok {
						delete(c.entries, k)
						atomic.AddInt64(&c.evicted, 1)
					}
				}
				c.mu.Unlock()
			}
		case <-c.stopCh:
			return
		}
	}
}
func (c *Cache) writeToDB(key Key, e *entry) {
	if c.db == nil {
		return
	}
	err := e.msg.Pack()
	if err != nil {
		return
	}
	data := e.msg.Data
	_, err = c.db.Exec(
		`INSERT OR REPLACE INTO cache (name, qtype, class, data, stored_at, ttl_ns)
		VALUES (?, ?, ?, ?, ?, ?)`,
		key.Name, key.Qtype, key.Class, data,
		e.storedAt.UnixNano(), int64(e.ttl),
	)
	if err != nil {
		slog.Warn("cache write to db failed", "err", err)
	}
}

func (c *Cache) dbWriter() {
	defer c.wg.Done()
	for w := range c.dbCh {
		c.writeToDB(w.key, w.e)
	}
}
func (c *Cache) loadFromDB() {
	rows, err := c.db.Query(
		`SELECT name, qtype, class, data, stored_at, ttl_ns FROM cache`,
	)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var name string
		var qtype, class uint16
		var data []byte
		var storedAtNano, ttlNs int64

		if err := rows.Scan(&name, &qtype, &class, &data, &storedAtNano, &ttlNs); err != nil {
			continue
		}

		msg := new(dns.Msg)
		msg.Data = data
		if err := msg.Unpack(); err != nil {
			continue
		}

		e := &entry{
			msg:      msg,
			storedAt: time.Unix(0, storedAtNano),
			ttl:      time.Duration(ttlNs),
		}

		if !e.expired() {
			c.entries[Key{Name: name, Qtype: qtype, Class: class}] = e
		} else {
			c.db.Exec(`DELETE FROM cache WHERE name=? AND qtype=? AND class=?`, name, qtype, class)
		}
	}
}


func computeTTL(msg *dns.Msg) time.Duration {
	if msg.Rcode == dns.RcodeNameError && len(msg.Answer) == 0 {
		for _, rr := range msg.Ns {
			if soa, ok := rr.(*dns.SOA); ok {
				return time.Duration(soa.Minttl) * time.Second
			}
		}
	}

	var min uint32
	first := true
	for _, rr := range msg.Answer {
		if first || rr.Header().TTL < min {
			min = rr.Header().TTL
			first = false
		}
	}
	for _, rr := range msg.Ns {
		if first || rr.Header().TTL < min {
			min = rr.Header().TTL
			first = false
		}
	}
	for _, rr := range msg.Extra {
		if dns.RRToType(rr) == dns.TypeOPT {
			continue
		}
		if first || rr.Header().TTL < min {
			min = rr.Header().TTL
			first = false
		}
	}
	if !first {
		return time.Duration(min) * time.Second
	}
	for _, rr := range msg.Ns {
		if soa, ok := rr.(*dns.SOA); ok {
			return time.Duration(soa.Minttl) * time.Second
		}
	}
	return 0
}

func adjustTTL(msg *dns.Msg, originalTTL, remaining time.Duration) {
	ratio := float64(remaining) / float64(originalTTL)
	for _, rr := range msg.Answer {
		rr.Header().TTL = applyRatio(rr.Header().TTL, ratio)
	}
	for _, rr := range msg.Ns {
		if dns.RRToType(rr) == dns.TypeOPT {
			continue
		}
		rr.Header().TTL = applyRatio(rr.Header().TTL, ratio)
	}
	for _, rr := range msg.Extra {
		if dns.RRToType(rr) == dns.TypeOPT {
			continue
		}
		rr.Header().TTL = applyRatio(rr.Header().TTL, ratio)
	}
}

func applyRatio(ttl uint32, ratio float64) uint32 {
	return uint32(float64(ttl) * ratio)
}

func deepCopyMsg(msg *dns.Msg) *dns.Msg {
	cp := new(dns.Msg)
	cp.Response = msg.Response
	cp.ID = msg.ID
	cp.Opcode = msg.Opcode
	cp.Authoritative = msg.Authoritative
	cp.Truncated = msg.Truncated
	cp.RecursionDesired = msg.RecursionDesired
	cp.RecursionAvailable = msg.RecursionAvailable
	cp.Rcode = msg.Rcode
	cp.UDPSize = msg.UDPSize
	cp.Question = copyRRSlice(msg.Question)
	cp.Answer = copyRRSlice(msg.Answer)
	cp.Ns = copyRRSlice(msg.Ns)
	cp.Extra = copyRRSlice(msg.Extra)
	return cp
}

func copyRRSlice(src []dns.RR) []dns.RR {
	if src == nil {
		return nil
	}
	dst := make([]dns.RR, len(src))
	for i, rr := range src {
		dst[i] = copyRR(rr)
	}
	return dst
}

func copyRR(rr dns.RR) dns.RR {
	if rr == nil {
		return nil
	}
	switch v := rr.(type) {
	case *dns.A:
		cp := *v
		return &cp
	case *dns.AAAA:
		cp := *v
		return &cp
	case *dns.NS:
		cp := *v
		return &cp
	case *dns.CNAME:
		cp := *v
		return &cp
	case *dns.SOA:
		cp := *v
		return &cp
	case *dns.MX:
		cp := *v
		return &cp
	case *dns.TXT:
		cp := *v
		return &cp
	case *dns.SRV:
		cp := *v
		return &cp
	case *dns.PTR:
		cp := *v
		return &cp
	default:
		return rr
	}
}