fancy-dns/internal/blocklist/manager.go

package blocklist

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

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

// Source is a compiled list plus the bookkeeping the UI shows.
type Source struct {
	Name   string     `json:"name"`
	Format Format     `json:"format"`
	Meta   SourceMeta `json:"meta"`
	Stats  ParseStats `json:"stats"`
	// Digest identifies the compiled content, so instances can tell whether
	// they hold the same rules without shipping the rules around.
	Digest string `json:"digest"`
}

// State is the compiled blocklist plus its per-source detail.
type State struct {
	Set     *Set
	Sources []Source
	// Duration is how long the last full compile took.
	Duration time.Duration
	// Errors are per-source failures from the last refresh. A source that
	// fails keeps its previous content, so the errors are warnings, not
	// fatal.
	Errors map[string]string
	// Digest covers every source, in order.
	Digest string
}

// Manager owns the blocklist lifecycle: fetching sources, compiling them, and
// atomically swapping the compiled Set into the holder.
type Manager struct {
	cfg   *config.Store
	log   *slog.Logger
	hold  *Holder
	http  *http.Client
	stats *StateHolder

	mu         sync.Mutex
	running    map[string]bool // in-flight refreshes, by source name
	managedMu  sync.RWMutex
	managedRPZ string

	// trigger lets the cluster ask one instance to refresh (leader election).
	trigger chan string
	stop    chan struct{}
	wg      sync.WaitGroup
}

// StateHolder is a swappable State for the API.
type StateHolder struct{ cur atomic.Pointer[State] }

// NewStateHolder wraps an initial state.
func NewStateHolder(s *State) *StateHolder {
	h := &StateHolder{}
	if s != nil {
		h.cur.Store(s)
	}
	return h
}

// Get returns the current state.
func (h *StateHolder) Get() *State { return h.cur.Load() }

// ManagerOptions configures a Manager.
type ManagerOptions struct {
	// Config is the hot-reloadable config store.
	Config *config.Store
	// Logger is the structured logger.
	Logger *slog.Logger
	// Client is the HTTP client used for list fetches; nil builds a default.
	Client *http.Client
}

// NewManager builds a Manager.
func NewManager(opt ManagerOptions) *Manager {
	log := opt.Logger
	if log == nil {
		log = slog.Default()
	}
	client := opt.Client
	if client == nil {
		client = &http.Client{
			Timeout: 120 * time.Second,
			Transport: &http.Transport{
				MaxIdleConns:        16,
				MaxIdleConnsPerHost: 4,
				IdleConnTimeout:     60 * time.Second,
			},
		}
	}
	m := &Manager{
		cfg:     opt.Config,
		log:     log.With("component", "blocklist"),
		hold:    NewHolder(nil),
		http:    client,
		running: make(map[string]bool),
		trigger: make(chan string, 16),
		stop:    make(chan struct{}),
	}
	m.stats = NewStateHolder(&State{Errors: map[string]string{}})
	return m
}

// Holder returns the compiled-set holder the engine reads.
func (m *Manager) Holder() *Holder { return m.hold }

// State returns the detailed state the API serves.
func (m *Manager) State() *State { return m.stats.Get() }

// SetManagedRPZ replaces the shared UI-managed address override source.
func (m *Manager) SetManagedRPZ(content string) {
	m.managedMu.Lock()
	m.managedRPZ = content
	m.managedMu.Unlock()
}

// RefreshAll compiles every enabled source plus the local zones. The previous
// set is kept if every source fails, so a bad deploy never leaves the
// resolver with no blocklists at all.
func (m *Manager) RefreshAll(ctx context.Context) (State, error) {
	start := time.Now()
	cfg := m.cfg.Get()
	b := NewBuilder()
	var (
		sources []Source
		errs    = map[string]string{}
	)
	m.managedMu.RLock()
	managedRPZ := m.managedRPZ
	m.managedMu.RUnlock()
	if strings.TrimSpace(managedRPZ) != "" {
		rules, st, err := Parse(FormatRPZ, []byte(managedRPZ), ParseOptions{Subdomains: true, DefaultTTL: 60})
		if err != nil {
			errs["managed:address-overrides"] = err.Error()
		} else {
			for _, r := range rules {
				b.Add(r, "managed:address-overrides", nil)
			}
			b.NoteSource(SourceMeta{Name: "managed:address-overrides", Format: string(FormatRPZ), Rules: len(rules), Bytes: st.Bytes, UpdatedAt: start, Duration: config.Duration(time.Since(start)), Enabled: true})
		}
	}

	// Inline block/allow entries first: they are cheap, and putting them
	// before the lists means a later list cannot override an operator's
	// inline decision without a priority bump.
	if len(cfg.Block.BlockEntries) > 0 {
		rules, st, err := parseInline(cfg.Block.BlockEntries, ActionBlock, cfg.Block.ResponseTTL)
		if err == nil {
			for _, r := range rules {
				b.Add(r, "inline:block", nil)
			}
			b.NoteSource(SourceMeta{
				Name: "inline:block", Format: string(FormatPlain), Rules: len(rules),
				Bytes: st.Bytes, UpdatedAt: start, Duration: config.Duration(time.Since(start)), Enabled: true,
			})
		} else {
			errs["inline:block"] = err.Error()
		}
	}
	if len(cfg.Block.AllowEntries) > 0 {
		rules, st, err := parseInline(cfg.Block.AllowEntries, ActionAllow, cfg.Block.ResponseTTL)
		if err == nil {
			for _, r := range rules {
				b.Add(r, "inline:allow", nil)
			}
			b.NoteSource(SourceMeta{
				Name: "inline:allow", Format: string(FormatPlain), Rules: len(rules),
				Bytes: st.Bytes, UpdatedAt: start, Duration: config.Duration(time.Since(start)), Enabled: true,
			})
		} else {
			errs["inline:allow"] = err.Error()
		}
	}

	// Local RPZ-style zone files.
	for _, pattern := range cfg.Block.LocalZones {
		files, err := filepath.Glob(m.cfg.Get().AbsPath(pattern))
		if err != nil || len(files) == 0 {
			// A glob that matches nothing is a configuration smell, not a
			// reason to fail the whole compile.
			errs[pattern] = "no files matched"
			continue
		}
		sort.Strings(files)
		rules, st, err := parseRPZFiles(files)
		if err != nil {
			errs[pattern] = err.Error()
			continue
		}
		name := "local:" + pattern
		for _, r := range rules {
			b.Add(r, name, nil)
		}
		b.NoteSource(SourceMeta{
			Name: name, Format: string(FormatRPZ), Rules: len(rules), Bytes: st.Bytes,
			UpdatedAt: start, Duration: config.Duration(time.Since(start)), Enabled: true, Path: pattern,
		})
	}

	// Remote / on-disk lists.
	for _, l := range cfg.Block.Lists {
		if !l.Enabled {
			b.NoteSource(SourceMeta{
				Name: l.Name, Format: l.Format, Enabled: false, URL: l.URL,
				Path: cfg.ListPath(l),
			})
			continue
		}
		if !m.claim(l.Name) {
			errs[l.Name] = "refresh already in flight"
			continue
		}
		src, err := m.loadOne(ctx, cfg, l)
		m.release(l.Name)
		if err != nil {
			errs[l.Name] = err.Error()
			// Keep the previous content: record what we know about the
			// source so the UI can show it as stale rather than missing.
			if prev, ok := m.previous(l.Name); ok {
				src = rawSource{
					Name: prev.Name, Format: prev.Format, Meta: prev.Meta,
					Stats: prev.Stats, Digest: prev.Digest,
				}
				src.Meta.Err = err.Error()
				src.Meta.Stale = true
			} else {
				continue
			}
		}
		for _, r := range src.rules {
			b.Add(r, l.Name, l.Categories)
		}
		sources = append(sources, src.toSource())
	}

	// Carry forward bookkeeping for sources that contributed nothing this
	// round (disabled, or skipped because a refresh was already in flight), so
	// the UI keeps showing them. Only the *metadata* is carried: their rules
	// are deliberately not re-added, because a disabled source must stop
	// matching, and a source that failed its refresh must fall back to the
	// copy the manager kept for it (above), not to a silent re-add.
	for _, prev := range m.previousAll() {
		if _, ok := b.metaFor(prev.Name); ok {
			continue
		}
		carried := prev
		carried.Meta.Rules = 0
		sources = append(sources, carried)
	}

	set := b.Build()
	state := State{
		Set:      set,
		Sources:  sources,
		Duration: time.Since(start),
		Errors:   errs,
		Digest:   set.Digest(),
	}
	m.hold.Set(set)
	m.stats.cur.Store(&state)
	m.log.Info("blocklist compiled",
		"rules", set.Len(), "sources", len(sources), "errors", len(errs),
		"duration", state.Duration)
	return state, nil
}

func (m *Manager) previous(name string) (Source, bool) {
	st := m.stats.Get()
	if st == nil {
		return Source{}, false
	}
	for _, s := range st.Sources {
		if s.Name == name {
			return s, true
		}
	}
	return Source{}, false
}

func (m *Manager) previousAll() []Source {
	st := m.stats.Get()
	if st == nil {
		return nil
	}
	return st.Sources
}

func (m *Manager) claim(name string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.running[name] {
		return false
	}
	m.running[name] = true
	return true
}

func (m *Manager) release(name string) {
	m.mu.Lock()
	delete(m.running, name)
	m.mu.Unlock()
}

// loadOne fetches and compiles a single list source.
func (m *Manager) loadOne(ctx context.Context, cfg *config.Config, l config.ListSource) (rawSource, error) {
	start := time.Now()
	path := cfg.ListPath(l)

	var (
		data       []byte
		etag       = l.ETag
		lastMod    = l.LastModified
		httpStatus int
	)

	needFetch := true
	if path != "" {
		// Prefer the local copy: it is what a restart without network sees.
		if fi, err := os.Stat(path); err == nil {
			if fi.Size() < l.MaxBytes {
				data, _ = os.ReadFile(path)
				needFetch = false
			}
		}
	}
	if l.URL != "" && needFetch {
		fetched, status, newETag, newLM, err := m.fetch(ctx, l)
		if err != nil {
			if len(data) == 0 {
				return rawSource{}, fmt.Errorf("fetch: %w", err)
			}
			// Keep the local copy and report the failure; a stale list beats
			// no list.
			httpStatus = status
		} else {
			data = fetched
			httpStatus = status
			if newETag != "" {
				etag = newETag
			}
			if newLM != "" {
				lastMod = newLM
			}
			if path != "" && !l.InsecureSkipTLSVerify {
				// Persist for offline restarts. A write failure is not fatal.
				_ = config.WriteFileAtomic(path, data, 0o644)
			}
		}
	} else if l.URL != "" && !needFetch {
		// We have a local copy; try a conditional refresh in the background
		// anyway, but do not fail when it 304s.
		httpStatus = http.StatusOK
	}
	if len(data) == 0 {
		return rawSource{}, errors.New("no data available")
	}

	format := Format(l.Format)
	opt := ParseOptions{
		Subdomains: true,
		DefaultTTL: uint32(cfg.Block.ResponseTTL.D()),
		Categories: l.Categories,
	}
	rules, st, perr := Parse(format, data, opt)
	if perr != nil {
		return rawSource{}, perr
	}
	meta := SourceMeta{
		Name:       l.Name,
		Format:     string(st.Format),
		Rules:      len(rules),
		Bytes:      int64(len(data)),
		ETag:       etag,
		LastMod:    lastMod,
		UpdatedAt:  start,
		Duration:   config.Duration(time.Since(start)),
		Enabled:    true,
		URL:        l.URL,
		Path:       path,
		HTTPStatus: httpStatus,
	}
	return rawSource{
		Name:   l.Name,
		Format: st.Format,
		Meta:   meta,
		Stats:  st,
		Digest: digestOf(data),
		rules:  rules,
	}, nil
}

// rawSource is the internal shape before it becomes a Source.
type rawSource struct {
	Name   string
	Format Format
	Meta   SourceMeta
	Stats  ParseStats
	Digest string
	rules  []Rule
}

func (r rawSource) toSource() Source {
	return Source{Name: r.Name, Format: r.Format, Meta: r.Meta, Stats: r.Stats, Digest: r.Digest}
}

// fetch performs a conditional GET against a list URL, following mirrors.
func (m *Manager) fetch(ctx context.Context, l config.ListSource) (data []byte, status int, etag, lastMod string, err error) {
	urls := append([]string{l.URL}, l.Mirror...)
	var lastErr error
	for _, u := range urls {
		if u == "" {
			continue
		}
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
		if err != nil {
			lastErr = err
			continue
		}
		req.Header.Set("User-Agent", "fancy-dns/1.0 (+blocklist)")
		if l.ETag != "" {
			req.Header.Set("If-None-Match", l.ETag)
		}
		if l.LastModified != "" {
			req.Header.Set("If-Modified-Since", l.LastModified)
		}
		client := m.http
		if l.InsecureSkipTLSVerify {
			// Only for self-hosted mirrors with a private CA; the config
			// field name says what it does.
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.TLSClientConfig = insecureTLS()
			client = &http.Client{Transport: tr, Timeout: client.Timeout}
		}
		resp, err := client.Do(req)
		if err != nil {
			lastErr = err
			continue
		}
		status = resp.StatusCode
		switch resp.StatusCode {
		case http.StatusNotModified:
			resp.Body.Close()
			return nil, status, l.ETag, l.LastModified, nil
		case http.StatusOK:
			etag = resp.Header.Get("ETag")
			lastMod = resp.Header.Get("Last-Modified")
			body, rerr := readLimited(resp.Body, l.MaxBytes)
			resp.Body.Close()
			if rerr != nil {
				lastErr = rerr
				continue
			}
			return body, status, etag, lastMod, nil
		default:
			io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
			resp.Body.Close()
			lastErr = fmt.Errorf("%s: unexpected status %s", u, resp.Status)
		}
	}
	if lastErr == nil {
		lastErr = errors.New("no usable source")
	}
	return nil, 0, "", "", lastErr
}

func readLimited(r io.Reader, max int64) ([]byte, error) {
	if max <= 0 {
		max = 512 << 20
	}
	buf, err := io.ReadAll(io.LimitReader(r, max+1))
	if err != nil {
		return nil, err
	}
	if int64(len(buf)) > max {
		return nil, fmt.Errorf("list exceeds %d bytes", max)
	}
	return buf, nil
}

// Start runs the scheduled refresh loop until ctx is cancelled.
func (m *Manager) Start(ctx context.Context) {
	m.wg.Add(1)
	go func() {
		defer m.wg.Done()
		// Stagger the first pass across sources so a cold cluster does not
		// fetch 20 lists at the same instant.
		for _, l := range m.cfg.Get().Block.Lists {
			if !l.Enabled || l.Refresh <= 0 {
				continue
			}
			jitter := l.Refresh.D() * time.Duration(fraction(l.Name)) / 8
			select {
			case <-ctx.Done():
				return
			case <-time.After(jitter):
			}
			if m.claim(l.Name) {
				if _, err := m.RefreshAll(ctx); err != nil {
					m.log.Warn("initial blocklist refresh failed", "error", err)
				}
				m.release(l.Name)
			}
			// Only stagger the first pass; after that the loop is regular.
			break
		}
		if _, err := m.RefreshAll(ctx); err != nil {
			m.log.Warn("initial blocklist refresh failed", "error", err)
		}
		m.loop(ctx)
	}()
}

func (m *Manager) loop(ctx context.Context) {
	// A 5-minute tick is coarse enough to be free and fine enough that a
	// source's own refresh interval is honoured within 5 minutes.
	t := time.NewTicker(5 * time.Minute)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case name := <-m.trigger:
			// A cluster peer asked for a refresh; it is the leader's job.
			if name != "" && !m.claim(name) {
				continue
			}
			if _, err := m.RefreshAll(ctx); err != nil {
				m.log.Warn("triggered refresh failed", "source", name, "error", err)
			}
			if name != "" {
				m.release(name)
			}
		case <-t.C:
			m.refreshDue(ctx)
		}
	}
}

// refreshDue refreshes the sources whose interval has elapsed.
func (m *Manager) refreshDue(ctx context.Context) {
	cfg := m.cfg.Get()
	now := time.Now()
	var due bool
	for _, l := range cfg.Block.Lists {
		if !l.Enabled || l.Refresh <= 0 {
			continue
		}
		src, ok := m.previous(l.Name)
		if !ok || src.Meta.UpdatedAt.IsZero() {
			due = true
			break
		}
		if now.Sub(src.Meta.UpdatedAt) >= l.Refresh.D() {
			due = true
			break
		}
	}
	if !due {
		return
	}
	// Refreshing every due source in one pass compiles once, which is far
	// cheaper than one compile per source.
	if _, err := m.RefreshAll(ctx); err != nil {
		m.log.Warn("scheduled refresh failed", "error", err)
	}
}

// RefreshOne refreshes a single named source immediately.
func (m *Manager) RefreshOne(ctx context.Context, name string) (State, error) {
	if !m.claim(name) {
		return *m.stats.Get(), fmt.Errorf("refresh of %q already in flight", name)
	}
	defer m.release(name)
	return m.RefreshAll(ctx)
}

// Trigger asks the refresh loop to compile now. Safe to call from an API
// handler; it does not block.
func (m *Manager) Trigger(name string) {
	select {
	case m.trigger <- name:
	default:
	}
}

// Stop shuts the loop down.
func (m *Manager) Stop() {
	select {
	case <-m.stop:
	default:
		close(m.stop)
	}
	m.wg.Wait()
}

// parseInline compiles the config's inline entry lists. The plain-domain
// parser yields block rules, so the requested action is stamped on afterwards:
// an operator's allowlist must not be compiled as a blocklist, which is the
// kind of bug that silently inverts a policy.
func parseInline(entries []string, action Action, ttl config.Duration) ([]Rule, ParseStats, error) {
	var sb strings.Builder
	for _, e := range entries {
		sb.WriteString(e)
		sb.WriteByte('\n')
	}
	opt := ParseOptions{Subdomains: true, DefaultTTL: uint32(ttl.D())}
	rules, st, err := Parse(FormatPlain, []byte(sb.String()), opt)
	if err != nil {
		return nil, st, err
	}
	for i := range rules {
		if rules[i].Action == ActionBlock {
			rules[i].Action = action
		}
	}
	return rules, st, nil
}

func parseRPZFiles(files []string) ([]Rule, ParseStats, error) {
	var sb strings.Builder
	for _, f := range files {
		data, err := os.ReadFile(f)
		if err != nil {
			return nil, ParseStats{}, err
		}
		// Prefix each file with its own $ORIGIN so a multi-file policy zone
		// keeps the per-file origins from its file name.
		fmt.Fprintf(&sb, "$ORIGIN %s\n", zonefile.OriginFromPath(f))
		sb.Write(data)
		sb.WriteByte('\n')
	}
	opt := ParseOptions{Subdomains: true, DefaultTTL: 60}
	return Parse(FormatRPZ, []byte(sb.String()), opt)
}

func digestOf(data []byte) string {
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:8])
}

// fraction maps a name to a stable 0..1 value, for refresh jitter.
func fraction(name string) int {
	sum := sha256.Sum256([]byte(name))
	return int(sum[0]) // 0..255
}

// insecureTLS is the transport tweak for sources that need it. It is only
// reachable when a list explicitly sets insecure_skip_tls_verify.
func insecureTLS() *tls.Config { return &tls.Config{InsecureSkipVerify: true} }