fancy-dns/internal/blocklist/set.go

// Package blocklist compiles heterogeneous block/allow sources (hosts files,
// adblock lists, plain domain lists and RPZ policy zones) into a single
// immutable Set that answers "what should I do with this name?" in a handful
// of map lookups.
//
// A Set is never mutated after construction; Manager swaps a whole new Set in
// atomically when lists are refreshed, so a query either sees the old rules or
// the new ones, never a half-applied mixture.
package blocklist

import (
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"io"
	"strings"
	"sync/atomic"
	"time"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

// Action is what to do with a matched name.
type Action uint8

const (
	// ActionBlock uses the manager's default response (nxdomain/zero/...).
	ActionBlock Action = iota
	// ActionAllow whitelists the name: blocklists are skipped entirely.
	ActionAllow
	// ActionPassthrough resolves the name upstream but ignores blocklists
	// (RPZ rpz-passthru).
	ActionPassthrough
	// ActionDrop answers nothing at all (RPZ rpz-drop).
	ActionDrop
	// ActionNXDomain answers NXDOMAIN regardless of the default response.
	ActionNXDomain
	// ActionZero answers the configured sinkhole address.
	ActionZero
	// ActionCNAME answers a CNAME to Response.
	ActionCNAME
	// ActionRedirect answers A/AAAA for Response as an address literal.
	ActionRedirect
)

func (a Action) String() string {
	switch a {
	case ActionBlock:
		return "block"
	case ActionAllow:
		return "allow"
	case ActionPassthrough:
		return "passthrough"
	case ActionDrop:
		return "drop"
	case ActionNXDomain:
		return "nxdomain"
	case ActionZero:
		return "zero"
	case ActionCNAME:
		return "cname"
	case ActionRedirect:
		return "redirect"
	}
	return "unknown"
}

// ParseAction maps configuration and UI strings onto an Action.
func ParseAction(s string) (Action, bool) {
	switch strings.ToLower(strings.TrimSpace(s)) {
	case "", "block":
		return ActionBlock, true
	case "allow", "whitelist":
		return ActionAllow, true
	case "passthrough", "pass":
		return ActionPassthrough, true
	case "drop":
		return ActionDrop, true
	case "nxdomain":
		return ActionNXDomain, true
	case "zero", "sinkhole", "null":
		return ActionZero, true
	case "cname":
		return ActionCNAME, true
	case "redirect":
		return ActionRedirect, true
	}
	return ActionBlock, false
}

// Entry is one compiled rule. It is shared by every name that maps to it, so
// keep it small: the strings are interned per source.
//
// The hit counter is a pointer so an Entry can be copied by value (the API
// returns copies) without copying a mutex or tripping go vet's lock check.
type Entry struct {
	// Name is the rule's pattern, normalised (lowercase, no trailing dot).
	Name   string
	Action Action
	// Response is the sinkhole address or CNAME target for the actions that
	// need one; empty otherwise.
	Response string
	// Source is the list or file the rule came from (shown in the UI).
	Source string
	// Categories are the list's tags, copied per rule.
	Categories []string
	// Priority orders allow-over-block conflicts: higher wins, and ties break
	// towards the first rule seen.
	Priority int
	// hits counts matches since the set was built. A pointer keeps Entry
	// copyable by value, which the admin API relies on.
	hits *atomic.Uint64
	// TTL overrides the manager's response TTL for this rule (RPZ records
	// carry their own TTL).
	TTL uint32
}

// Hits returns how many times this rule matched.
func (e *Entry) Hits() uint64 {
	if e == nil || e.hits == nil {
		return 0
	}
	return e.hits.Load()
}

func (e *Entry) markHit() {
	if e == nil || e.hits == nil {
		return
	}
	e.hits.Add(1)
}

// IsAllow reports whether the rule whitelists the name.
func (e *Entry) IsAllow() bool {
	return e.Action == ActionAllow || e.Action == ActionPassthrough
}

// Rule is a pattern plus the action it implies, as produced by the parsers.
type Rule struct {
	// Name is the pattern; may start with "*." (subdomains only) or
	// "**." (subdomains and the apex itself).
	Name   string
	Action Action
	// Response is an address literal (for redirect) or a target name (cname).
	Response string
	// TTL is a per-rule TTL when the source provides one.
	TTL uint32
	// Category tags the rule.
	Category string
	// Priority disambiguates allow/block conflicts.
	Priority int
}

// Set is an immutable compiled rule table.
//
// Three maps cover the pattern space:
//
//	exact[name]        "ads.example.com"          the name and nothing else
//	sub[example.com]   "*.example.com"            strict subdomains
//	root[example.com]  "**.example.com"           subdomains and the apex
//
// Lookup is one hash for the common case (exact hit or "no wildcards at all"),
// and at most one hash per label otherwise — no per-query allocation.
type Set struct {
	exact map[string]*Entry
	sub   map[string]*Entry
	root  map[string]*Entry
	// hasWild lets Match skip the suffix walk for sets of pure exact rules,
	// which is the common hosts-file case.
	hasWild bool
	// meta is per-source bookkeeping for the UI.
	meta map[string]SourceMeta
	// counts is the number of compiled rules, per kind.
	counts RuleCounts
	// total is the number of rules dropped as duplicates or malformed.
	duplicates int
	// builtAt is when this set was compiled.
	builtAt time.Time
	// digest covers the compiled patterns, so two instances can compare rules
	// without shipping the rules around.
	digest string
	// sources is the ordered list of source names that contributed.
	sources []string
	// maxDepth is the longest pattern label count, used to bound the walk.
	maxDepth int
}

// RuleCounts breaks down compiled rules by kind.
type RuleCounts struct {
	Exact    int
	Sub      int
	Root     int
	Allow    int
	Block    int
	Passthru int
	Drop     int
}

// SourceMeta describes one contributing source.
type SourceMeta struct {
	Name      string    `json:"name"`
	Format    string    `json:"format"`
	Rules     int       `json:"rules"`
	Bytes     int64     `json:"bytes"`
	ETag      string    `json:"etag"`
	LastMod   string    `json:"last_mod"`
	UpdatedAt time.Time `json:"updated_at"`
	// Duration is a duration STRING on the wire, like every other duration the
	// API serves: a raw time.Duration marshals as a nanosecond count, and the
	// console was dividing by 1e6 to compensate for a number nobody else on
	// the wire sends.
	Duration   config.Duration `json:"duration"`
	Err        string          `json:"err"`
	Enabled    bool            `json:"enabled"`
	URL        string          `json:"url"`
	Path       string          `json:"path"`
	Stale      bool            `json:"stale"`
	HTTPStatus int             `json:"http_status"`
}

// Builder accumulates rules and produces a Set.
type Builder struct {
	exact map[string]*Entry
	sub   map[string]*Entry
	root  map[string]*Entry
	meta  map[string]SourceMeta
	order []string
	// ordered mirrors order, so a source is listed once even when both a rule
	// and its metadata mention it.
	ordered map[string]bool
	dup     int
	counts  RuleCounts
}

// NewBuilder returns an empty Builder.
func NewBuilder() *Builder {
	return &Builder{
		exact: make(map[string]*Entry, 1<<12),
		sub:   make(map[string]*Entry, 1<<10),
		root:  make(map[string]*Entry, 1<<10),
		meta:  make(map[string]SourceMeta, 8),
	}
}

// Add compiles one rule. Names are normalised here, so callers may pass
// user-typed domain names.
func (b *Builder) Add(r Rule, source string, categories []string) {
	name, kind := splitPattern(r.Name)
	if name == "" {
		return
	}
	prio := r.Priority
	if prio == 0 {
		// Allow rules outrank block rules by default so an explicit
		// whitelist always wins over a list entry.
		if r.Action == ActionAllow || r.Action == ActionPassthrough {
			prio = 10
		} else {
			prio = 1
		}
	}
	e := &Entry{
		hits:       new(atomic.Uint64),
		Name:       name,
		Action:     r.Action,
		Response:   r.Response,
		Source:     source,
		Categories: categories,
		Priority:   prio,
		TTL:        r.TTL,
	}
	if r.Category != "" {
		e.Categories = append(e.Categories, r.Category)
	}
	var m map[string]*Entry
	switch kind {
	case patternExact:
		m = b.exact
	case patternSub:
		m = b.sub
	case patternRoot:
		m = b.root
	}
	if existing, ok := m[name]; ok {
		b.dup++
		// Higher priority wins; on a tie the existing rule stays, so the
		// order lists appear in the config is meaningful.
		if prio <= existing.Priority {
			return
		}
		b.uncount(existing)
		m[name] = e
		b.countAction(e, 1)
		return
	}
	m[name] = e
	b.countAction(e, 1)
	b.noteOrder(source)
	// A source that only ever contributed rules (an inline blocklist entry, or
	// a list whose bookkeeping was never noted) is still a source: the UI
	// lists sources, and a rule with a source nobody has heard of would show
	// up as a blank column.
	if source != "" {
		if _, ok := b.meta[source]; !ok {
			b.meta[source] = SourceMeta{Name: source}
		}
	}
}

// metaFor reports whether a source already contributed to this build.
func (b *Builder) metaFor(name string) (SourceMeta, bool) {
	m, ok := b.meta[name]
	return m, ok
}

func (b *Builder) countAction(e *Entry, delta int) {
	switch e.Action {
	case ActionAllow:
		b.counts.Allow += delta
	case ActionPassthrough:
		b.counts.Passthru += delta
	case ActionDrop:
		b.counts.Drop += delta
	default:
		b.counts.Block += delta
	}
}

// uncount removes an entry that a higher-priority rule is replacing.
func (b *Builder) uncount(e *Entry) {
	switch e.Action {
	case ActionAllow:
		b.counts.Allow--
	case ActionPassthrough:
		b.counts.Passthru--
	case ActionDrop:
		b.counts.Drop--
	default:
		b.counts.Block--
	}
}

// NoteSource records (or updates) bookkeeping for a contributing source.
func (b *Builder) NoteSource(meta SourceMeta) {
	b.noteOrder(meta.Name)
	b.meta[meta.Name] = meta
}

// noteOrder appends a source to the build order exactly once.
func (b *Builder) noteOrder(name string) {
	if name == "" {
		return
	}
	if b.ordered == nil {
		b.ordered = make(map[string]bool, 16)
	}
	if b.ordered[name] {
		return
	}
	b.ordered[name] = true
	b.order = append(b.order, name)
}

// SourceCount is the number of contributing sources.
func (b *Builder) SourceCount() int { return len(b.meta) }

// Build freezes the accumulated rules into an immutable Set.
func (b *Builder) Build() *Set {
	s := &Set{
		exact:      b.exact,
		sub:        b.sub,
		root:       b.root,
		meta:       b.meta,
		sources:    b.order,
		builtAt:    time.Now(),
		duplicates: b.dup,
		hasWild:    len(b.sub) > 0 || len(b.root) > 0,
	}
	s.counts = b.counts
	s.counts.Exact = len(b.exact)
	s.counts.Sub = len(b.sub)
	s.counts.Root = len(b.root)
	s.digest = b.digest()
	for k := range b.exact {
		if d := labelCount(k); d > s.maxDepth {
			s.maxDepth = d
		}
	}
	for k := range b.sub {
		if d := labelCount(k); d > s.maxDepth {
			s.maxDepth = d
		}
	}
	for k := range b.root {
		if d := labelCount(k); d > s.maxDepth {
			s.maxDepth = d
		}
	}
	return s
}

// digest hashes every compiled pattern together with its action. The map walk
// makes it order-independent, but stable for a given rule set, which is exactly
// what cluster convergence checking needs.
func (b *Builder) digest() string {
	h := sha256.New()
	var buf [3]byte
	write := func(kind byte, name string, e *Entry) {
		buf[0] = kind
		h.Write(buf[:1])
		io.WriteString(h, name)
		buf[1] = byte(e.Action)
		binary.LittleEndian.PutUint16(buf[1:], uint16(e.Priority))
		h.Write(buf[:3])
	}
	for name, e := range b.exact {
		write(0, name, e)
	}
	for name, e := range b.sub {
		write(1, name, e)
	}
	for name, e := range b.root {
		write(2, name, e)
	}
	return hex.EncodeToString(h.Sum(nil)[:8])
}

// Digest returns the compiled content's digest.
func (s *Set) Digest() string {
	if s == nil {
		return ""
	}
	return s.digest
}

type patternKind uint8

const (
	patternExact patternKind = iota
	patternSub
	patternRoot
)

// splitPattern normalises a rule name and reports which map it belongs in.
//
//	"Example.COM." -> exact "example.com"
//	"*.example.com" -> sub   "example.com"
//	"**.example.com" -> root  "example.com"
//	"." (the whole internet) is dropped: it is never a useful rule.
func splitPattern(in string) (string, patternKind) {
	name := strings.ToLower(strings.TrimSpace(in))
	if name == "" {
		return "", patternExact
	}
	kind := patternExact
	switch {
	case strings.HasPrefix(name, "**."):
		kind = patternRoot
		name = name[3:]
	case strings.HasPrefix(name, "*."):
		kind = patternSub
		name = name[2:]
	case strings.HasPrefix(name, "."):
		// A leading dot in hosts/adblock convention means "and subdomains".
		kind = patternRoot
		name = name[1:]
	}
	name = strings.TrimSuffix(name, ".")
	name = strings.Trim(name, ".")
	if name == "" || name == "." {
		return "", patternExact
	}
	// A name with an interior "*" is not something we can match; treat the
	// first label as a fixed prefix and drop the rest, which degrades to
	// "closest sane behaviour" instead of silently not matching.
	if i := strings.IndexByte(name, '*'); i >= 0 {
		name = name[:i]
		name = strings.TrimSuffix(name, ".")
		if name == "" {
			return "", patternExact
		}
	}
	return name, kind
}

// NormalizeName canonicalises a queried name for matching: lowercased, no
// trailing dot. It is allocation-free for the common case.
func NormalizeName(name string) string {
	if len(name) > 0 && name[len(name)-1] == '.' {
		name = name[:len(name)-1]
	}
	return strings.ToLower(name)
}

// Match returns the rule that applies to name, or nil. name may be fqdn or
// relative; the caller is responsible for lowercasing only if it has not
// already.
func (s *Set) Match(name string) *Entry {
	if s == nil {
		return nil
	}
	name = NormalizeName(name)
	if name == "" {
		return nil
	}
	if e, ok := s.exact[name]; ok {
		e.markHit()
		return e
	}
	if !s.hasWild {
		return nil
	}
	// The apex pattern "**.x" also covers the name itself.
	if e, ok := s.root[name]; ok {
		e.markHit()
		return e
	}
	// Walk strictly shorter suffixes, most specific first.
	for off := labelEnd(name); off < len(name); off = labelEnd(name[off:]) + off {
		suffix := name[off:]
		if e, ok := s.root[suffix]; ok {
			e.markHit()
			return e
		}
		if e, ok := s.sub[suffix]; ok {
			e.markHit()
			return e
		}
	}
	return nil
}

// labelEnd returns the index just past the first label of s, or len(s) when s
// has no dot.
func labelEnd(s string) int {
	if i := strings.IndexByte(s, '.'); i >= 0 {
		return i + 1
	}
	return len(s)
}

func labelCount(s string) int {
	n := 1
	for i := 0; i < len(s); i++ {
		if s[i] == '.' {
			n++
		}
	}
	return n
}

// Len reports the total number of compiled rules.
func (s *Set) Len() int {
	if s == nil {
		return 0
	}
	return s.counts.Exact + s.counts.Sub + s.counts.Root
}

// Counts returns the per-kind rule breakdown.
func (s *Set) Counts() RuleCounts {
	if s == nil {
		return RuleCounts{}
	}
	return s.counts
}

// HasWildcards reports whether any subdomain rules exist.
func (s *Set) HasWildcards() bool { return s != nil && s.hasWild }

// Duplicates is the number of rules dropped because a higher- or equal-priority
// rule already claimed the name.
func (s *Set) Duplicates() int { return s.DuplicatesN() }

// DuplicatesN is Duplicates with a nil guard.
func (s *Set) DuplicatesN() int {
	if s == nil {
		return 0
	}
	return s.duplicates
}

// BuiltAt is when the set was compiled.
func (s *Set) BuiltAt() time.Time {
	if s == nil {
		return time.Time{}
	}
	return s.builtAt
}

// Sources lists the contributing source names in order.
func (s *Set) Sources() []string {
	if s == nil {
		return nil
	}
	return s.sources
}

// SourceMeta returns bookkeeping for one source.
func (s *Set) SourceMeta(name string) (SourceMeta, bool) {
	if s == nil {
		return SourceMeta{}, false
	}
	m, ok := s.meta[name]
	return m, ok
}

// AllSourceMeta returns the per-source bookkeeping map.
func (s *Set) AllSourceMeta() map[string]SourceMeta {
	if s == nil {
		return nil
	}
	return s.meta
}

// PatternKind says which of the three maps a rule lives in.
type PatternKind uint8

const (
	// PatternExact is "ads.example.com": the name and nothing else.
	PatternExact PatternKind = iota
	// PatternSub is "*.example.com": strict subdomains.
	PatternSub
	// PatternRoot is "**.example.com": subdomains and the apex.
	PatternRoot
)

// String renders the kind for the UI.
func (k PatternKind) String() string {
	switch k {
	case PatternSub:
		return "subdomain"
	case PatternRoot:
		return "domain"
	default:
		return "exact"
	}
}

// Pattern renders a stored name for display.
//
// The wildcard prefix is *not* re-added: a Set stores the prefix-free name and
// the kind already carries the wildcard semantics, so the stored form is the
// form a caller wrote. Use WireName when the prefixed form is needed — to hand
// a rule back to the API, for instance.
func (k PatternKind) Pattern(name string) string {
	return name
}

// WireName renders a stored name with the wildcard prefix this kind implies.
// It is the round-trip form: WireName(k, k.Pattern(n)) == n, and feeding it
// back into Add reproduces the original rule.
func (k PatternKind) WireName(name string) string {
	switch k {
	case PatternSub:
		return "*." + name
	case PatternRoot:
		return "**." + name
	default:
		return name
	}
}

// Each visits every rule. It exists for the admin API's rule browser; the DNS
// hot path uses Match and never pays this cost.
func (s *Set) Each(fn func(pattern string, e *Entry, kind PatternKind) bool) {
	if s == nil {
		return
	}
	for name, e := range s.exact {
		if !fn(name, e, PatternExact) {
			return
		}
	}
	for name, e := range s.sub {
		if !fn(name, e, PatternSub) {
			return
		}
	}
	for name, e := range s.root {
		if !fn(name, e, PatternRoot) {
			return
		}
	}
}

// MatchPattern returns the kind of rule that matched a name, and the stored
// pattern (without its wildcard prefix). It reports PatternExact when only an
// exact rule matched.
func (s *Set) MatchPattern(name string) PatternKind {
	if s == nil {
		return PatternExact
	}
	name = NormalizeName(name)
	if _, ok := s.exact[name]; ok {
		return PatternExact
	}
	if _, ok := s.root[name]; ok {
		return PatternRoot
	}
	for off := labelEnd(name); off < len(name); off = labelEnd(name[off:]) + off {
		suffix := name[off:]
		if _, ok := s.root[suffix]; ok {
			return PatternRoot
		}
		if _, ok := s.sub[suffix]; ok {
			return PatternSub
		}
	}
	return PatternExact
}

// Lookup returns the exact rule for a name without walking wildcards. It is
// used by the admin API to explain a match.
func (s *Set) Lookup(name string) *Entry {
	if s == nil {
		return nil
	}
	return s.exact[NormalizeName(name)]
}

// Stats is a JSON-friendly snapshot for the API and UI.
type Stats struct {
	Total      int                   `json:"total"`
	Exact      int                   `json:"exact"`
	Subdomains int                   `json:"subdomain_rules"`
	RootRules  int                   `json:"root_rules"`
	Allow      int                   `json:"allow"`
	Block      int                   `json:"block"`
	Passthru   int                   `json:"passthrough"`
	Drop       int                   `json:"drop"`
	Duplicates int                   `json:"duplicates"`
	Sources    int                   `json:"sources"`
	BuiltAt    time.Time             `json:"built_at"`
	PerSource  map[string]SourceMeta `json:"sources_detail"`
}

// Stats summarises the set for the API.
func (s *Set) Stats() Stats {
	if s == nil {
		return Stats{PerSource: map[string]SourceMeta{}}
	}
	return Stats{
		Total:      s.Len(),
		Exact:      s.counts.Exact,
		Subdomains: s.counts.Sub,
		RootRules:  s.counts.Root,
		Allow:      s.counts.Allow,
		Block:      s.counts.Block,
		Passthru:   s.counts.Passthru,
		Drop:       s.counts.Drop,
		Duplicates: s.duplicates,
		Sources:    len(s.meta),
		BuiltAt:    s.builtAt,
		PerSource:  s.meta,
	}
}

// Holder is a swappable Set, so hot paths can grab a pointer once per query
// and be immune to a concurrent refresh.
type Holder struct {
	cur atomic.Pointer[Set]
}

// NewHolder wraps an initial set (may be nil).
func NewHolder(s *Set) *Holder {
	h := &Holder{}
	if s != nil {
		h.cur.Store(s)
	}
	return h
}

// Get returns the current set, or nil when nothing has been compiled yet. The
// returned set is immutable and is shared with every other reader, so a caller
// must not modify it.
func (h *Holder) Get() *Set {
	if h == nil {
		return nil
	}
	return h.cur.Load()
}

// Set installs a new set.
func (h *Holder) Set(s *Set) { h.cur.Store(s) }