fancy-dns/internal/qlog/qlog.go

// Package qlog records DNS queries for the admin console: a live in-memory
// ring for the SSE stream, a bbolt-backed ring for the query log, and rolling
// per-domain counters that back the "top blocked" and "top queried" panels.
//
// The write path is deliberately allocation-light: entries are value types, the
// live ring is a fixed array, and bbolt writes are batched on a background
// goroutine so a burst of queries never blocks a DNS response.
package qlog

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"
	bolt "go.etcd.io/bbolt"
)

// Action is the outcome classification shown in the log.
type Action string

const (
	// ActionAllow is a normal resolution.
	ActionAllow Action = "allow"
	// ActionBlock is a blocklist hit.
	ActionBlock Action = "block"
	// ActionRedirect is a blocklist rewrite to an address or CNAME target.
	ActionRedirect Action = "redirect"
	// ActionCached is a cache hit.
	ActionCached Action = "cached"
	// ActionError is a SERVFAIL or other failure.
	ActionError Action = "error"
	// ActionDrop is a query answered with silence.
	ActionDrop Action = "drop"
	// ActionRefused is an ACL refusal.
	ActionRefused Action = "refused"
)

// Entry is one logged query.
type Entry struct {
	ID       uint64    `json:"id"`
	TS       time.Time `json:"ts"`
	Client   string    `json:"client"`
	Group    string    `json:"client_group,omitempty"`
	Name     string    `json:"name"`
	QType    uint16    `json:"qtype"`
	QTypeStr string    `json:"qtype_name"`
	RCode    int       `json:"rcode"`
	Action   Action    `json:"action"`
	Response string    `json:"response,omitempty"`
	Duration int64     `json:"duration_ms"`
	Upstream string    `json:"upstream,omitempty"`
	List     string    `json:"list,omitempty"`
	Protocol string    `json:"protocol,omitempty"`
	// Blocked is a convenience mirror of Action == ActionBlock so the UI can
	// filter without parsing the action string.
	Blocked bool `json:"blocked"`
}

// FromMsg fills the question-derived fields from a DNS message.
func (e *Entry) FromMsg(m *dns.Msg) {
	if len(m.Question) > 0 {
		q := m.Question[0]
		e.Name = q.Name
		e.QType = q.Qtype
		e.QTypeStr = dns.Type(q.Qtype).String()
	}
	e.RCode = m.Rcode
}

// RCodeName renders the RCODE as its mnemonic.
func (e Entry) RCodeName() string {
	if s, ok := dns.RcodeToString[e.RCode]; ok {
		return s
	}
	return fmt.Sprintf("RCODE%d", e.RCode)
}

// liveEntry is a decoded copy kept in memory for the SSE stream.
type liveEntry struct {
	entry Entry
}

// Logger batches query logging to bbolt and maintains the live ring.
type Logger struct {
	opts Options

	// seq assigns monotonic ids so the UI can detect drops.
	seq atomic.Uint64

	// live is a fixed-size ring of the most recent entries. Each slot is a
	// pointer, so a writer on the query path never races a reader walking the
	// ring: publishing is one atomic word, and a reader sees either the whole
	// entry or the one it replaced. The slot an entry lands in comes from its
	// id, so concurrent writers cannot leave a hole or reorder the stream.
	live     []atomic.Pointer[liveEntry]
	liveMu   sync.RWMutex
	liveCond *sync.Cond
	liveFull atomic.Bool

	// pending batches entries between bbolt transactions.
	pendingMu sync.Mutex
	pending   []Entry

	// closed is set by Close; queries can still be logged after shutdown, so
	// Log has to notice rather than write to a closed bbolt handle.
	closed atomic.Bool

	// counts are rolling per-domain counters.
	counts *domainCounts

	// logErrors counts persistence failures, which the API surfaces so a
	// silently-failing query log never looks healthy.
	logErrors atomic.Uint64

	// subscribers is the set of live SSE consumers.
	subMu  sync.Mutex
	subs   map[int]chan Entry
	subSeq int

	db   *bolt.DB
	stop chan struct{}
	done chan struct{}
	wg   sync.WaitGroup
}

// Options configures a Logger.
type Options struct {
	// Path is the bbolt file; empty disables persistence.
	Path string
	// MaxEntries bounds the persisted log ring.
	MaxEntries int
	// LiveBuffer is the in-memory ring size for the SSE stream.
	LiveBuffer int
	// FlushInterval is the bbolt batch window.
	FlushInterval time.Duration
	// SampleRate is the fraction of non-blocked, non-error queries to keep.
	SampleRate float64
	// LogBlocked, LogCached, LogErrors are forced categories.
	LogBlocked bool
	LogCached  bool
	LogErrors  bool
	// TopDomains is the size of the rolling counter set.
	TopDomains int
	// TopTTL is how long a counter survives without being touched.
	TopTTL time.Duration
	// Now is injectable for tests.
	Now func() time.Time
}

// DefaultOptions returns sane defaults.
func DefaultOptions() Options {
	return Options{
		Path:          "data/queries.db",
		MaxEntries:    200000,
		LiveBuffer:    2000,
		FlushInterval: time.Second,
		SampleRate:    1,
		LogBlocked:    true,
		LogCached:     true,
		LogErrors:     true,
		TopDomains:    5000,
		TopTTL:        5 * time.Minute,
	}
}

// New opens a Logger, creating the bbolt file when a path is given.
func New(opt Options) (*Logger, error) {
	if opt.Now == nil {
		opt.Now = time.Now
	}
	if opt.LiveBuffer <= 0 {
		opt.LiveBuffer = 2000
	}
	if opt.FlushInterval <= 0 {
		opt.FlushInterval = time.Second
	}
	if opt.SampleRate <= 0 {
		opt.SampleRate = 1
	}
	if opt.SampleRate > 1 {
		opt.SampleRate = 1
	}
	if opt.TopDomains <= 0 {
		opt.TopDomains = 5000
	}
	if opt.TopTTL <= 0 {
		opt.TopTTL = 5 * time.Minute
	}
	l := &Logger{
		opts:    opt,
		live:    make([]atomic.Pointer[liveEntry], opt.LiveBuffer),
		counts:  newDomainCounts(opt.TopDomains, opt.TopTTL),
		subs:    make(map[int]chan Entry),
		stop:    make(chan struct{}),
		done:    make(chan struct{}),
		pending: make([]Entry, 0, 4096),
	}
	l.liveCond = sync.NewCond(&l.liveMu)
	if opt.Path != "" {
		if err := os.MkdirAll(filepath.Dir(opt.Path), 0o755); err != nil {
			return nil, fmt.Errorf("qlog: create data dir: %w", err)
		}
		db, err := bolt.Open(opt.Path, 0o600, &bolt.Options{Timeout: 5 * time.Second, NoFreelistSync: true})
		if err != nil {
			return nil, fmt.Errorf("qlog: open %s: %w", opt.Path, err)
		}
		if err := db.Update(func(tx *bolt.Tx) error {
			_, err := tx.CreateBucketIfNotExists(bucketLog)
			if err != nil {
				return err
			}
			meta, err := tx.CreateBucketIfNotExists(bucketMeta)
			if err != nil {
				return err
			}
			// Resume the id sequence where the previous process stopped. The
			// bbolt keys are that same sequence, so ids stay unique across a
			// restart: a repeated id would make the SSE stream and the UI show
			// a live entry as an older one carrying the same id.
			if v := meta.Get(keyNextID); len(v) == 8 {
				// next_id is the id the next entry will take, while seq
				// counts ids already handed out.
				l.seq.Store(beUint64(v) - 1)
			}
			return nil
		}); err != nil {
			_ = db.Close()
			return nil, err
		}
		l.db = db
		l.wg.Add(1)
		go l.flushLoop()
	} else {
		close(l.done)
	}
	return l, nil
}

var (
	bucketLog  = []byte("log")
	bucketMeta = []byte("meta")
	keyNextID  = []byte("next_id")
)

// ShouldLog reports whether an entry with this action passes the sampling
// policy. Blocked, redirected, cached and error entries are always kept when
// their category flag is set; ordinary resolutions are sampled.
func (l *Logger) ShouldLog(a Action) bool {
	switch a {
	case ActionBlock, ActionRedirect:
		return l.opts.LogBlocked
	case ActionCached:
		return l.opts.LogCached
	case ActionError, ActionDrop:
		return l.opts.LogErrors
	default:
		if l.opts.SampleRate >= 1 {
			return true
		}
		return randFloat() < l.opts.SampleRate
	}
}

// Log records a query. It never blocks on bbolt: the entry goes to the live
// ring and a batch, and a background goroutine does the disk write.
func (l *Logger) Log(e Entry) {
	if l == nil || l.closed.Load() {
		return
	}
	if e.TS.IsZero() {
		e.TS = l.opts.Now()
	}
	if e.QTypeStr == "" && e.QType != 0 {
		e.QTypeStr = dns.Type(e.QType).String()
	}
	e.Blocked = e.Action == ActionBlock
	e.ID = l.seq.Add(1)
	l.counts.observe(e)

	// Live ring. The slot is published with a single atomic store; a reader
	// walking the ring concurrently gets either the old or the new entry,
	// never a half-written one.
	// The index is the entry id rather than a second counter, so a burst from
	// several resolver goroutines still lands gap free and in id order.
	l.live[(e.ID-1)%uint64(len(l.live))].Store(&liveEntry{entry: e})

	// Fan out to SSE subscribers (non-blocking: a slow consumer drops entries
	// rather than stalling the query path).
	l.subMu.Lock()
	for _, ch := range l.subs {
		select {
		case ch <- e:
		default:
		}
	}
	l.subMu.Unlock()

	if l.db == nil {
		return
	}
	l.pendingMu.Lock()
	l.pending = append(l.pending, e)
	n := len(l.pending)
	l.pendingMu.Unlock()
	if n >= 2048 {
		go l.Flush()
	}
}

// flushLoop writes pending batches on a timer.
func (l *Logger) flushLoop() {
	defer l.wg.Done()
	t := time.NewTicker(l.opts.FlushInterval)
	defer t.Stop()
	for {
		select {
		case <-l.stop:
			l.Flush()
			return
		case <-t.C:
			l.Flush()
		}
	}
}

// Flush writes any pending entries to bbolt and trims the ring.
func (l *Logger) Flush() {
	if l == nil || l.db == nil {
		return
	}
	l.pendingMu.Lock()
	if len(l.pending) == 0 {
		l.pendingMu.Unlock()
		return
	}
	batch := l.pending
	// Not pending[:0]: that would hand the next Log call the array this batch is
	// still being read from.
	l.pending = nil
	l.pendingMu.Unlock()

	err := l.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketLog)
		if b == nil {
			return nil
		}
		meta := tx.Bucket(bucketMeta)
		nextID := uint64(1)
		if meta != nil {
			if v := meta.Get(keyNextID); len(v) == 8 {
				nextID = beUint64(v)
			}
		}
		for _, e := range batch {
			data, err := json.Marshal(e)
			if err != nil {
				continue
			}
			if err := b.Put(be64b(nextID), data); err != nil {
				return err
			}
			nextID++
		}
		if meta != nil {
			meta.Put(keyNextID, be64b(nextID))
		}
		return l.trim(b, nextID)
	})
	if err != nil {
		// A logging failure must never take the resolver down; surface it via
		// the error counter the API reads.
		l.logErrors.Add(1)
	}
}

// trim enforces the entry cap by deleting the oldest ids.
func (l *Logger) trim(b *bolt.Bucket, nextID uint64) error {
	if l.opts.MaxEntries <= 0 {
		return nil
	}
	// The keys are a dense id range, so the count is the span from the oldest
	// surviving key up to nextID. Bucket.Stats cannot answer that here: it
	// walks the pages as they sit on disk, and a write transaction has not
	// spilled this batch yet, so it always under-counts and the cap is never
	// enforced.
	cur := b.Cursor()
	first, _ := cur.First()
	if first == nil {
		return nil
	}
	oldest := beUint64(first)
	span := nextID - oldest
	if span <= uint64(l.opts.MaxEntries) {
		return nil
	}
	excess := span - uint64(l.opts.MaxEntries)
	// Collect the doomed keys before deleting any: bbolt rewrites a leaf
	// page in place once it drops under a quarter full, which invalidates the
	// cursor's page stack, so deleting while iterating stops early and the
	// cap is only half enforced.
	keys := make([][]byte, 0, excess)
	for k, _ := cur.First(); k != nil && len(keys) < int(excess); k, _ = cur.Next() {
		keys = append(keys, append([]byte(nil), k...))
	}
	for _, k := range keys {
		if err := b.Delete(k); err != nil {
			return err
		}
	}
	return nil
}

// ErrorCount reports how many batch writes have failed.
func (l *Logger) ErrorCount() uint64 { return l.logErrors.Load() }

// be64b encodes a uint64 big-endian for use as a bbolt key.
func be64b(v uint64) []byte {
	return []byte{
		byte(v >> 56), byte(v >> 48), byte(v >> 40), byte(v >> 32),
		byte(v >> 24), byte(v >> 16), byte(v >> 8), byte(v),
	}
}

// beUint64 decodes what be64b wrote.
func beUint64(b []byte) uint64 {
	var v uint64
	for _, c := range b {
		v = v<<8 | uint64(c)
	}
	return v
}

// Query filters the persisted log.
type Query struct {
	// Limit is the page size.
	Limit int
	// Offset skips the newest Offset entries.
	Offset int
	// Text matches name, client or list.
	Text string
	// Action filters by classification.
	Action string
	// Client filters by client address.
	Client string
	// RCode filters by numeric rcode (0 = no filter).
	RCode int
	// BlockedOnly keeps only blocklist hits.
	BlockedOnly bool
	// Since keeps only entries at or after this time.
	Since time.Time
}

// Page is a page of log entries plus the total match count.
type Page struct {
	Items  []Entry `json:"items"`
	Total  int     `json:"total"`
	Offset int     `json:"offset"`
	Limit  int     `json:"limit"`
}

// Recent returns the newest entries from the in-memory ring, newest first.
// It powers the SSE initial burst so a client connecting mid-stream sees
// recent history immediately.
func (l *Logger) Recent(limit int) []Entry {
	if l == nil {
		return nil
	}
	if limit <= 0 {
		limit = 200
	}
	if limit > len(l.live) {
		limit = len(l.live)
	}
	out := make([]Entry, 0, limit)
	size := uint64(len(l.live))
	pos := l.seq.Load() // the newest id doubles as the newest ring position
	total := pos
	if total > size {
		total = size
	}
	for i := uint64(0); i < total; i++ {
		le := l.live[(pos-1-i)%size].Load()
		if le == nil {
			// Never written (or already aged out of the window).
			continue
		}
		out = append(out, le.entry)
		if len(out) >= limit {
			break
		}
	}
	return out
}

// Search queries the persisted log.
func (l *Logger) Search(q Query) (Page, error) {
	if l == nil {
		return Page{Offset: q.Offset, Limit: q.Limit}, nil
	}
	l.Flush()
	page := Page{Offset: q.Offset, Limit: q.Limit}
	if q.Limit <= 0 {
		q.Limit = 200
	}
	if q.Limit > 1000 {
		q.Limit = 1000
	}
	if l.db == nil {
		// No persistence: serve from the live ring.
		items := l.Recent(1000)
		filtered := l.filter(items, q)
		page.Total = len(filtered)
		page.Items = paginate(filtered, q.Offset, q.Limit)
		return page, nil
	}
	var (
		matched []Entry
		total   int
	)
	err := l.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketLog)
		if b == nil {
			return nil
		}
		// Walk newest-first.
		c := b.Cursor()
		for k, v := c.Last(); k != nil; k, v = c.Prev() {
			var e Entry
			if err := json.Unmarshal(v, &e); err != nil {
				continue
			}
			if !matchEntry(e, q) {
				continue
			}
			total++
			if total > q.Offset && len(matched) < q.Limit {
				matched = append(matched, e)
			}
			if total >= q.Offset+q.Limit+1 && q.Offset+q.Limit > 0 {
				// We have enough to fill the page and know the total is at
				// least this high; keep scanning only if we still need an
				// exact total. For the UI, a "has more" hint is enough.
				break
			}
		}
		return nil
	})
	if err != nil {
		return page, err
	}
	page.Items = matched
	page.Total = total
	return page, nil
}

func (l *Logger) filter(items []Entry, q Query) []Entry {
	out := make([]Entry, 0, len(items))
	for _, e := range items {
		if matchEntry(e, q) {
			out = append(out, e)
		}
	}
	return out
}

func matchEntry(e Entry, q Query) bool {
	if q.BlockedOnly && !e.Blocked {
		return false
	}
	if q.Action != "" && string(e.Action) != q.Action {
		return false
	}
	if q.Client != "" && e.Client != q.Client {
		return false
	}
	if q.RCode != 0 && e.RCode != q.RCode {
		return false
	}
	if !q.Since.IsZero() && e.TS.Before(q.Since) {
		return false
	}
	if q.Text != "" {
		t := q.Text
		if !containsFold(e.Name, t) && !containsFold(e.Client, t) && !containsFold(e.List, t) {
			return false
		}
	}
	return true
}

func paginate(items []Entry, offset, limit int) []Entry {
	if offset >= len(items) {
		return []Entry{}
	}
	end := offset + limit
	if end > len(items) {
		end = len(items)
	}
	return items[offset:end]
}

func containsFold(hay, needle string) bool {
	if len(needle) == 0 {
		return true
	}
	return indexFold(hay, needle) >= 0
}

func indexFold(hay, needle string) int {
	n := len(needle)
	if n == 0 {
		return 0
	}
	if n > len(hay) {
		return -1
	}
	for i := 0; i+n <= len(hay); i++ {
		match := true
		for j := 0; j < n; j++ {
			if lower(hay[i+j]) != lower(needle[j]) {
				match = false
				break
			}
		}
		if match {
			return i
		}
	}
	return -1
}

func lower(b byte) byte {
	if b >= 'A' && b <= 'Z' {
		return b + 32
	}
	return b
}

// Subscribe registers a live-stream consumer. The returned cancel func must
// be called. Entries are delivered best-effort: a full channel drops rather
// than blocking the DNS path.
func (l *Logger) Subscribe() (<-chan Entry, func()) {
	l.subMu.Lock()
	defer l.subMu.Unlock()
	l.subSeq++
	id := l.subSeq
	ch := make(chan Entry, 256)
	l.subs[id] = ch
	return ch, func() {
		l.subMu.Lock()
		if c, ok := l.subs[id]; ok {
			delete(l.subs, id)
			close(c)
		}
		l.subMu.Unlock()
	}
}

// SubscriberCount reports how many live streams are attached.
func (l *Logger) SubscriberCount() int {
	l.subMu.Lock()
	defer l.subMu.Unlock()
	return len(l.subs)
}

// Clear removes every persisted entry.
func (l *Logger) Clear() error {
	if l == nil || l.db == nil {
		return nil
	}
	err := l.db.Update(func(tx *bolt.Tx) error {
		if tx.Bucket(bucketLog) == nil {
			_, err := tx.CreateBucket(bucketLog)
			return err
		}
		if err := tx.DeleteBucket(bucketLog); err != nil {
			return err
		}
		// Recreate it in the same transaction: the flush path assumes the
		// bucket is there, and a logger that silently stopped persisting after
		// a clear would look like a healthy, empty log.
		_, err := tx.CreateBucket(bucketLog)
		return err
	})
	if err != nil {
		l.logErrors.Add(1)
	}
	return err
}

// Count is the number of persisted entries.
func (l *Logger) Count() int {
	if l == nil {
		return 0
	}
	if l.db == nil {
		return int(l.seq.Load())
	}
	n := 0
	_ = l.db.View(func(tx *bolt.Tx) error {
		if b := tx.Bucket(bucketLog); b != nil {
			n = b.Stats().KeyN
		}
		return nil
	})
	return n
}

// TopDomains returns the hottest domains, optionally only blocked ones.
type DomainCount struct {
	Name    string `json:"domain"`
	Count   uint64 `json:"count"`
	Blocked uint64 `json:"blocked"`
	QType   uint16 `json:"qtype"`
	QTypeS  string `json:"qtype_name"`
}

// TopDomains returns the top n domains by query count.
func (l *Logger) TopDomains(n int, blockedOnly bool) []DomainCount {
	return l.counts.top(n, blockedOnly)
}

// TopBlocked returns the top blocked domains.
func (l *Logger) TopBlocked(n int) []DomainCount { return l.counts.top(n, true) }

// TopQueried returns the top queried domains.
func (l *Logger) TopQueried(n int) []DomainCount { return l.counts.top(n, false) }

// Stats is a JSON-friendly summary.
type Stats struct {
	Enabled     bool    `json:"enabled"`
	Path        string  `json:"path"`
	Entries     int     `json:"entries"`
	MaxEntries  int     `json:"max_entries"`
	LiveBuffer  int     `json:"live_buffer"`
	Subscribers int     `json:"subscribers"`
	SampleRate  float64 `json:"sample_rate"`
	TotalLogged uint64  `json:"total_logged"`
	Pending     int     `json:"pending"`
	Errors      uint64  `json:"errors"`
}

// Stats snapshots the logger.
func (l *Logger) Stats() Stats {
	l.pendingMu.Lock()
	pending := len(l.pending)
	l.pendingMu.Unlock()
	return Stats{
		Enabled:     l.db != nil,
		Path:        l.opts.Path,
		Entries:     l.Count(),
		MaxEntries:  l.opts.MaxEntries,
		LiveBuffer:  l.opts.LiveBuffer,
		Subscribers: l.SubscriberCount(),
		SampleRate:  l.opts.SampleRate,
		TotalLogged: l.seq.Load(),
		Pending:     pending,
		Errors:      l.logErrors.Load(),
	}
}

// Close flushes and releases resources.
func (l *Logger) Close() error {
	if l == nil {
		return nil
	}
	l.closed.Store(true)
	select {
	case <-l.stop:
	default:
		close(l.stop)
	}
	l.wg.Wait()
	if l.db != nil {
		l.Flush()
		return l.db.Close()
	}
	return nil
}