fancy-dns/internal/engine/engine.go

// Package engine is the query pipeline: the single place where a DNS question
// is turned into an answer.
//
// The order of operations matters and is the whole design:
//
//	rate limit -> ACL policy -> local zone -> blocklist -> cache -> upstream
//
// Everything in front of the cache is cheap and lock-light; the cache is a
// single shard-local lookup; only a miss reaches the network. The pipeline is
// a plain method call chain rather than a plugin stack, because the hot path
// runs millions of times a second and a plugin registry is pure overhead.
package engine

import (
	"context"
	"encoding/hex"
	"net/netip"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/acl"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/blocklist"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/cache"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/localzone"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/metrics"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/qlog"
	"git.k8s-git-k6n.net/mats/fancy-dns/internal/resolver"
)

// upstreamNSID is the nameserver identifier we advertise to upstreams when a
// client asks for one. It is hex-encoded on the way into the EDNS0 option:
// dns.EDNS0_NSID.Nsid is a hex string in this miekg/dns version, not the
// opaque bytes RFC 5001 describes.
const upstreamNSID = "fancy-dns"

// Request describes one inbound query.
type Request struct {
	// Msg is the parsed query. The engine may modify it (to attach an OPT or
	// strip EDNS) so the transport can write it back out.
	Msg *dns.Msg
	// Remote is the client address.
	Remote netip.Addr
	// Transport is "udp", "tcp", "dot" or "doh", used for metrics labels and
	// the log.
	Transport string
	// TCP reports whether the response may exceed the UDP size limit.
	TCP bool
	// Proto is the negotiated protocol string, kept for the log.
	Proto string
	// Start is when the request arrived.
	Start time.Time
	// ClientNames are reverse-DNS names for the client, if known. The server
	// fills these in from its reverse cache so the ACL can match on names
	// without a synchronous PTR lookup on the query path.
	ClientNames []string
	// ClientTags are derived labels for tag-based ACL matching.
	ClientTags []string
	// DoHPath records which DoH endpoint served the request, for the log.
	DoHPath string
}

// Response is the outcome plus the metadata the caller logs.
type Response struct {
	Msg *dns.Msg
	// Action is the pipeline classification.
	Action qlog.Action
	// Group is the matched ACL group.
	Group string
	// BlockSource is the blocklist/rule that blocked the query, for the log.
	BlockSource string
	// CacheHit reports whether the answer came from cache.
	CacheHit bool
	// CacheKind is "hit", "stale" or "miss".
	CacheKind string
	// Upstream is the address that answered.
	Upstream string
	// Local reports that the answer came from the local zone rather than from
	// the cache or an upstream. It is what makes metrics.OutcomeLocal
	// reachable: an authoritative local answer is a real, distinct outcome
	// and the dashboard has a column for it, but nothing set it, so every
	// locally answered query was counted as an upstream one and the
	// outcome="local" series was never published.
	Local bool
	// Duration is the total handling time.
	Duration time.Duration
	// Client is the resolved client string.
	Client string
	// Do reports whether the response carried the DNSSEC OK bit.
	Do bool
	// Truncated is set when the answer was shortened to fit the client buffer.
	Truncated bool
	// Drop is true when no response should be sent at all.
	Drop bool
	// GroupPolicy is the effective policy, for the UI.
	GroupPolicy string
}

// Engine holds the pipeline's swappable components.
type Engine struct {
	cfgStore *config.Store

	acl     atomic.Pointer[acl.Engine]
	blocks  *blocklist.Holder
	cache   *cache.Manager
	aclH    *acl.Holder
	local   *localzone.Holder
	pool    atomic.Pointer[resolver.Pool]
	logger  *qlog.Logger
	reverse *ReverseCache

	// concurrency limiter: a buffered channel as a semaphore, which is
	// cheaper than a sync.WaitGroup and gives us "try or fail" semantics.
	// Held in an atomic pointer because a config swap replaces it while
	// queries are in flight.
	sem     atomic.Pointer[chan struct{}]
	semSize atomic.Int64

	// perClient semaphores for clients with an explicit cap.
	clientMu   sync.Mutex
	clientSems map[string]*clientLimiter
	// rate limiters for clients with an explicit cap.
	rateMu    sync.Mutex
	rateState map[string]*rateState

	// inflight is the current in-flight count for the gauge.
	inflight atomic.Int64
	// dropped counts queries rejected by the concurrency cap.
	dropped atomic.Uint64

	// upstreamTimeout bounds a single upstream exchange.
	upstreamTimeout atomic.Int64

	// groupStats tracks per-group counters for the UI.
	groupStats *groupStatTable

	// flusher is the app's flush callback, used by the admin API.
	flusher atomic.Pointer[func()]
}

type clientLimiter struct {
	sem chan struct{}
}

type rateState struct {
	windowStart time.Time
	count       int
}

// Options configures a new Engine.
type Options struct {
	// Config is the hot-reloadable config store.
	Config *config.Store
	// Logger receives query entries; may be nil.
	Logger *qlog.Logger
	// Blocklists is the compiled blocklist holder.
	Blocklists *blocklist.Holder
	// Cache is the answer cache.
	Cache *cache.Manager
	// ACL is the compiled ACL engine.
	ACL *acl.Holder
	// Local is the local zone store.
	Local *localzone.Holder
	// Pool is the upstream pool.
	Pool *resolver.Pool
	// Reverse is the reverse-DNS cache used for name-based ACL matching.
	Reverse *ReverseCache
}

// New builds an Engine and wires the config-store hooks so a reload updates
// every swappable piece.
func New(opt Options) *Engine {
	e := &Engine{
		cfgStore:   opt.Config,
		logger:     opt.Logger,
		blocks:     opt.Blocklists,
		cache:      opt.Cache,
		aclH:       opt.ACL,
		local:      opt.Local,
		reverse:    opt.Reverse,
		groupStats: newGroupStatTable(),
		clientSems: make(map[string]*clientLimiter),
		rateState:  make(map[string]*rateState),
	}
	e.pool.Store(opt.Pool)
	if e.cfgStore != nil {
		cfg := e.cfgStore.Get()
		e.applyConfig(cfg)
		e.cfgStore.OnSwap(e.applyConfig)
	} else if opt.Pool != nil {
		e.upstreamTimeout.Store(int64(2 * time.Second))
	}
	// Prefetching a stale entry re-runs the pipeline for just that key.
	if e.cache != nil {
		e.cache.SetPrefetcher(func(k cache.Key) { e.prefetch(k) })
	}
	return e
}

// applyConfig hot-swaps the components affected by a config change.
func (e *Engine) applyConfig(cfg *config.Config) {
	semSize := cfg.Server.MaxConcurrent
	if semSize < 1 {
		semSize = 1024
	}
	sem := make(chan struct{}, semSize)
	e.sem.Store(&sem)
	e.semSize.Store(int64(semSize))
	e.upstreamTimeout.Store(int64(cfg.Server.Timeouts.Upstre.D()))

	if e.aclH != nil {
		e.aclH.Set(acl.Compile(cfg.ACL))
		metrics.SetACLGroups(len(cfg.ACL.Groups))
	}
	if e.local != nil && len(cfg.Block.LocalZones) > 0 {
		// Local zones are loaded asynchronously by the server; here we only
		// record that a reload happened.
		metrics.SetLocalRecords(e.local.Get().Count())
	}
	if e.cache != nil {
		e.cache.Resize(cfg.Cache.Size, cfg.Cache.Shards)
	}
	metrics.SetConcurrent(0)
}

// Logger exposes the query logger (may be nil).
func (e *Engine) Logger() *qlog.Logger { return e.logger }

// Cache exposes the answer cache (may be nil).
func (e *Engine) Cache() *cache.Manager { return e.cache }

// Pool returns the current upstream pool.
func (e *Engine) Pool() *resolver.Pool { return e.pool.Load() }

// Inflight is the number of queries currently being handled.
func (e *Engine) Inflight() int64 { return e.inflight.Load() }

// Dropped is the number of queries rejected by the concurrency cap.
func (e *Engine) Dropped() uint64 { return e.dropped.Load() }

// GroupStats returns per-group counters.
func (e *Engine) GroupStats() []acl.GroupSummary {
	base := e.aclH.Get().Summaries()
	return e.groupStats.merge(base)
}

// Handle runs the full pipeline for one request and returns the response to
// send. It never returns nil: a synthetic SERVFAIL is the worst case.
func (e *Engine) Handle(ctx context.Context, req Request) Response {
	start := req.Start
	if start.IsZero() {
		start = time.Now()
	}
	done := e.inflight.Add(1)
	metrics.SetConcurrent(int(e.inflight.Load()))
	defer func() {
		e.inflight.Add(-1)
		metrics.SetConcurrent(int(e.inflight.Load()))
		_ = done
	}()

	client := req.Remote.String()
	decision := e.decide(req)
	res := Response{Client: client, Group: decision.Group, GroupPolicy: string(decision.Policy)}
	q := req.Msg

	// A malformed or empty query is answered with FORMERR before anything
	// else looks at it.
	if err := validate(q); err != nil {
		res.Msg = formErr(q)
		res.Action = qlog.ActionError
		e.finish(req, res, start, decision)
		return res
	}

	qname := q.Question[0].Name
	qtype := q.Question[0].Qtype

	// ACL deny wins over everything, including the local zone: an operator who
	// denied a client means it.
	if !decision.Allows() {
		res.Msg = refuse(q, decision)
		res.Action = qlog.ActionRefused
		e.finish(req, res, start, decision)
		return res
	}

	// Concurrency cap. A per-client cap is checked first so one noisy client
	// cannot exhaust the global budget.
	if cl := decision.MaxConcurrent; cl > 0 {
		if !e.acquireClient(client, cl) {
			res.Msg = servFail(q, "client concurrency limit")
			res.Action = qlog.ActionError
			e.finish(req, res, start, decision)
			return res
		}
		defer e.releaseClient(client)
	}
	// A nil semaphore means "no limit configured yet", which only happens
	// before the first applyConfig; treat it as unbounded rather than dropping
	// every query.
	sem := e.sem.Load()
	if sem != nil {
		select {
		case *sem <- struct{}{}:
			defer func() { <-*sem }()
		default:
			e.dropped.Add(1)
			res.Msg = servFail(q, "server overloaded")
			res.Action = qlog.ActionError
			e.finish(req, res, start, decision)
			return res
		}
	}

	// Rate limiting.
	if decision.RateLimit > 0 && !e.allowRate(client, decision.RateLimit) {
		res.Msg = refuse(q, decision)
		res.Action = qlog.ActionRefused
		e.finish(req, res, start, decision)
		return res
	}

	// Group-level domain policy.
	if e.groupBlocks(decision, qname) {
		res.Msg = e.finalize(blockedResponse(q, e.blockCfg(), nil), q, req)
		res.Action = qlog.ActionBlock
		res.BlockSource = "acl:" + decision.Group
		e.finish(req, res, start, decision)
		return res
	}

	// Local zone data.
	if store := e.local.Get(); store != nil && store.Count() > 0 {
		if local, ok := e.answerLocal(q, store); ok {
			res.Msg = e.finalize(local, q, req)
			res.Action = qlog.ActionAllow
			// Authoritative local data is ActionAllow for the query log (the
			// query was permitted and answered), but it is its own metrics
			// outcome: counting it as upstream inflated the upstream column
			// with queries that never left the process.
			res.Local = true
			e.finish(req, res, start, decision)
			return res
		}
	}

	// Blocklist.
	var blocked *blocklist.Entry
	if !decision.SkipBlocklist && e.blocks != nil {
		if set := e.blocks.Get(); set != nil && set.Len() > 0 {
			blocked = set.Match(qname)
		}
	}
	if blocked != nil && !blocked.IsAllow() {
		switch blocked.Action {
		case blocklist.ActionDrop:
			res.Action = qlog.ActionDrop
			res.BlockSource = blocked.Source
			res.Drop = true
			res.Duration = time.Since(start)
			e.finish(req, res, start, decision)
			return res
		default:
			res.Msg = e.finalize(e.blockedMsg(q, blocked), q, req)
			res.Action = qlog.ActionBlock
			res.BlockSource = blocked.Source
			res.CacheHit = false
			e.finish(req, res, start, decision)
			return res
		}
	}

	// Cache. A group that skips the cache still gets served from it on a
	// prefetched entry only if the entry is fresh; we deliberately skip the
	// lookup entirely to keep the guarantee honest.
	if e.cache != nil && !decision.SkipCache {
		key := cache.QueryKey(qname, qtype, q.Question[0].Qclass, cache.Do(q), 0)
		hit := e.cache.Get(key)
		switch hit.Kind {
		case cache.Fresh, cache.Stale:
			res.Msg = hit.Msg
			res.CacheHit = true
			res.CacheKind = hit.Kind.String()
			res.Upstream = hit.Upstream
			if hit.Kind == cache.Stale {
				metrics.RecordStaleHit()
			}
			res.Msg = e.finalize(res.Msg, q, req)
			res.Action = qlog.ActionCached
			if hit.Kind == cache.Stale {
				// The cached answer is expired, so its TTL is 0; hand back a
				// short one so clients do not cache the stale copy forever.
				res.Msg.Answer = ageTo(res.Msg.Answer, 1)
				res.Msg.Ns = ageTo(res.Msg.Ns, 1)
			}
			e.finish(req, res, start, decision)
			return res
		default:
			metrics.RecordMiss()
		}
	}

	// Upstream.
	pool := e.pool.Load()
	if pool == nil {
		res.Msg = servFail(q, "no upstream configured")
		res.Action = qlog.ActionError
		e.finish(req, res, start, decision)
		return res
	}
	fwd := e.prepareUpstream(q, req)
	upStart := time.Now()
	up, addr, err := pool.Exchange(ctx, fwd, decision.Upstreams, time.Duration(e.upstreamTimeout.Load()))
	res.Duration = time.Since(start)
	if err != nil || up == nil {
		metrics.RecordUpstreamError()
		res.Msg = servFail(q, "upstream failure: "+shortErr(err))
		res.Action = qlog.ActionError
		e.finish(req, res, start, decision)
		return res
	}
	metrics.RecordUpstreamLatency(addr, time.Since(upStart))
	// finalize is what makes the answer safe to put on the wire: it echoes the
	// client's question and id verbatim and truncates to the client's
	// advertised buffer. Without it an uncached multi-kilobyte answer left
	// here oversized for a 512-byte UDP client, which no responder can deliver
	// and the client silently drops. (The cache-hit path above has always
	// done this; the uncached path did not.)
	reply := e.finalize(e.fromUpstream(q, up, req), q, req)
	res.Msg = reply
	res.Upstream = addr
	res.Action = qlog.ActionAllow
	// A non-success RCODE is reported as an error through finish's metrics
	// call, so the counters stay in one place.
	if reply.Rcode != dns.RcodeServerFailure && reply.Rcode != dns.RcodeSuccess {
		res.Action = qlog.ActionError
	} else if reply.Rcode == dns.RcodeServerFailure {
		res.Action = qlog.ActionError
	}
	if e.cache != nil && !decision.SkipCache {
		key := cache.QueryKey(qname, qtype, q.Question[0].Qclass, cache.Do(q), 0)
		// Cache the untruncated upstream answer so a small client later does
		// not poison the entry with a truncated copy.
		if !up.Truncated {
			e.cache.PutResponse(key, up, addr)
		}
	}
	e.finish(req, res, start, decision)
	return res
}

// finish attaches duration, records metrics and writes the query log.
func (e *Engine) finish(req Request, res Response, start time.Time, d acl.Decision) {
	if res.Duration == 0 {
		res.Duration = time.Since(start)
	}
	metrics.Record(req.Transport, outcomeFor(res), res.Duration)
	if res.Action == qlog.ActionAllow && res.CacheHit {
		metrics.Record(req.Transport, metrics.OutcomeCache, res.Duration)
	}
	e.groupStats.observe(d.Group, res)
	if e.logger == nil {
		return
	}
	if !e.logger.ShouldLog(res.Action) && !d.LogQueries {
		return
	}
	entry := qlog.Entry{
		TS:       start,
		Client:   res.Client,
		Group:    res.Group,
		Duration: res.Duration.Milliseconds(),
		Action:   res.Action,
		Upstream: res.Upstream,
		List:     res.BlockSource,
		Protocol: req.Transport,
	}
	if res.Msg != nil {
		entry.FromMsg(res.Msg)
	}
	e.logger.Log(entry)
}

// outcomeFor maps a finished query onto its metrics outcome. The action alone
// is not enough: ActionAllow covers both an upstream answer and a local-zone
// answer, and the two are different things for an operator reading the
// dashboard, so Response.Local is consulted first.
func outcomeFor(res Response) string {
	if res.Local {
		return metrics.OutcomeLocal
	}
	switch res.Action {
	case qlog.ActionCached:
		return metrics.OutcomeCache
	case qlog.ActionBlock:
		return metrics.OutcomeBlocked
	case qlog.ActionError, qlog.ActionDrop:
		return metrics.OutcomeError
	case qlog.ActionRefused:
		return metrics.OutcomeRefused
	default:
		return metrics.OutcomeUpstream
	}
}

// decide resolves the ACL decision for a request.
func (e *Engine) decide(req Request) acl.Decision {
	eng := e.aclH.Get()
	if eng == nil || eng.Count() == 0 {
		return acl.Decision{Policy: acl.PolicyAllow, Default: true}
	}
	return eng.Match(acl.Client{
		IP:    req.Remote,
		Names: req.ClientNames,
		Tags:  req.ClientTags,
	})
}

// groupBlocks applies a group's allow/deny domain lists.
func (e *Engine) groupBlocks(d acl.Decision, qname string) bool {
	if len(d.DenyDomains) > 0 {
		if acl.NewDomainSet(d.DenyDomains).Match(qname) {
			return true
		}
	}
	if len(d.AllowDomains) > 0 {
		return !acl.NewDomainSet(d.AllowDomains).Match(qname)
	}
	return false
}

// acquireClient takes a per-client slot, creating the limiter on demand.
func (e *Engine) acquireClient(client string, limit int) bool {
	e.clientMu.Lock()
	cl, ok := e.clientSems[client]
	if !ok || cap(cl.sem) != limit {
		cl = &clientLimiter{sem: make(chan struct{}, limit)}
		e.clientSems[client] = cl
	}
	e.clientMu.Unlock()
	select {
	case cl.sem <- struct{}{}:
		return true
	default:
		return false
	}
}

func (e *Engine) releaseClient(client string) {
	e.clientMu.Lock()
	cl, ok := e.clientSems[client]
	e.clientMu.Unlock()
	if !ok {
		return
	}
	select {
	case <-cl.sem:
	default:
	}
}

// allowRate applies a fixed-window rate limit. A fixed window is fine here:
// the ACL cap exists to stop runaway clients, and it costs one map lookup.
func (e *Engine) allowRate(client string, perSec int) bool {
	now := time.Now()
	e.rateMu.Lock()
	defer e.rateMu.Unlock()
	st, ok := e.rateState[client]
	if !ok {
		st = &rateState{windowStart: now}
		e.rateState[client] = st
	}
	if now.Sub(st.windowStart) >= time.Second {
		st.windowStart = now
		st.count = 0
	}
	if st.count >= perSec {
		return false
	}
	st.count++
	return true
}

// answerLocal serves a question from the local zone store, converting a
// negative result into NXDOMAIN or NODATA with the zone's SOA.
func (e *Engine) answerLocal(q *dns.Msg, store *localzone.Store) (*dns.Msg, bool) {
	if len(q.Question) == 0 {
		return nil, false
	}
	question := q.Question[0]
	res := store.Lookup(question)
	if res.Delegated {
		return nil, false
	}
	m := new(dns.Msg)
	m.SetReply(q)
	m.Authoritative = true
	m.RecursionAvailable = true
	if !res.Found {
		// Only claim NXDOMAIN when we are actually authoritative for the name.
		if res.Zone == "" {
			return nil, false
		}
		if res.Authority != nil {
			m.Rcode = dns.RcodeNameError
			m.Ns = []dns.RR{res.Authority}
		} else {
			m.Rcode = dns.RcodeSuccess // NODATA without an SOA
		}
		return m, true
	}
	if len(res.Records) == 0 {
		if res.Authority != nil {
			m.Ns = []dns.RR{res.Authority}
		}
		return m, true
	}
	// CNAME answers for other types need the target resolved too; we do not
	// chase here (the caller re-queries upstream if the target is not local),
	// but a CNAME-only answer for a CNAME question is complete.
	if len(m.Question) > 0 && m.Question[0].Qtype != res.Type {
		if cname, ok := firstOfType(res.Records, dns.TypeCNAME); ok {
			m.Answer = []dns.RR{cname}
			return m, true
		}
	}
	m.Answer = res.Records
	return m, true
}

func firstOfType(rrs []dns.RR, t uint16) (dns.RR, bool) {
	for _, rr := range rrs {
		if rr.Header().Rrtype == t {
			return rr, true
		}
	}
	return nil, false
}

// blockedMsg synthesises the blocklist answer.
func (e *Engine) blockedMsg(q *dns.Msg, entry *blocklist.Entry) *dns.Msg {
	cfg := e.blockCfg()
	if entry != nil && entry.TTL > 0 {
		c := cfg
		c.ResponseTTL = config.Duration(entry.TTL) * config.Duration(time.Second)
		cfg = c
	}
	if entry != nil {
		switch entry.Action {
		case blocklist.ActionNXDomain:
			return nxDomain(q, uint32(cfg.ResponseTTL))
		case blocklist.ActionZero:
			c := cfg
			c.DefaultAction = "zero"
			if entry.Response != "" {
				c.ResponseIP = entry.Response
			}
			return sinkhole(q, c)
		case blocklist.ActionRedirect:
			return redirect(q, entry.Response, uint32(cfg.ResponseTTL))
		case blocklist.ActionCNAME:
			return redirect(q, entry.Response, uint32(cfg.ResponseTTL))
		}
	}
	return blockedResponse(q, cfg, entry)
}

func (e *Engine) blockCfg() config.Block {
	if e.cfgStore == nil {
		return config.Default().Block
	}
	return e.cfgStore.Get().Block
}

// prepareUpstream builds the message to forward: the question plus the
// client's EDNS0 preferences, and nothing else (we never leak our own OPT to
// upstream, and we never forward the client's AD bit).
func (e *Engine) prepareUpstream(q *dns.Msg, req Request) *dns.Msg {
	cfg := e.cfgStore.Get()
	m := new(dns.Msg)
	m.SetQuestion(q.Question[0].Name, q.Question[0].Qtype)
	m.Question[0].Qclass = q.Question[0].Qclass
	m.RecursionDesired = true
	m.CheckingDisabled = q.CheckingDisabled

	var maxSize uint16 = cfg.DNS.MaxUDPSize
	if o := q.IsEdns0(); o != nil {
		if cfg.DNS.EDNS.PreserveEDNS0 {
			for _, opt := range o.Option {
				switch opt.Option() {
				case dns.EDNS0SUBNET:
					// Forward the client's ECS verbatim: the upstream uses it
					// to pick a nearby anycast location, which is the whole
					// point of the feature.
					m.SetEdns0(maxSize, o.Do())
					if opt4, ok := opt.(*dns.EDNS0_SUBNET); ok {
						appendECS(m, opt4)
					}
					continue
				case dns.EDNS0PADDING, dns.EDNS0COOKIE:
					// Cookie and padding are hop-by-hop; forwarding padding
					// upstream is harmless, cookies are not.
					if opt.Option() == dns.EDNS0PADDING {
						m.SetEdns0(maxSize, o.Do())
						po := new(dns.EDNS0_PADDING)
						po.Padding = opt4Padding(opt)
						m.IsEdns0().Option = append(m.IsEdns0().Option, po)
					}
					continue
				}
			}
		}
		if o.Do() {
			m.SetEdns0(maxSize, true)
		}
		if opt := o.Option; opt != nil {
			for _, oo := range opt {
				if oo.Option() == dns.EDNS0NSID {
					// Our own NSID, not the client's — but only when the
					// operator has not asked us to hide. dns.EDNS0_NSID.Nsid
					// is hex on the wire (the library hex-decodes it in pack
					// and hex-encodes it in unpack), so a plain string here
					// does not pack and every NSID query gets SERVFAIL.
					if e.cfgStore == nil || !e.cfgStore.Get().DNS.HideIdentity {
						nsid := new(dns.EDNS0_NSID)
						nsid.Nsid = hex.EncodeToString([]byte(upstreamNSID))
						if m.IsEdns0() == nil {
							m.SetEdns0(maxSize, o.Do())
						}
						m.IsEdns0().Option = append(m.IsEdns0().Option, nsid)
					}
				}
			}
		}
	}
	if m.IsEdns0() == nil {
		m.SetEdns0(maxSize, false)
	}
	return m
}

func appendECS(m *dns.Msg, src *dns.EDNS0_SUBNET) {
	sub := new(dns.EDNS0_SUBNET)
	sub.Code = src.Code
	sub.Family = src.Family
	sub.SourceNetmask = src.SourceNetmask
	sub.SourceScope = src.SourceScope
	sub.Address = src.Address
	if m.IsEdns0() == nil {
		m.SetEdns0(1232, false)
	}
	m.IsEdns0().Option = append(m.IsEdns0().Option, sub)
}

func opt4Padding(opt dns.EDNS0) []byte {
	if p, ok := opt.(*dns.EDNS0_PADDING); ok {
		return p.Padding
	}
	return nil
}

// fromUpstream converts an upstream answer into a client reply: the question
// and ID come from the client, the payload from upstream.
func (e *Engine) fromUpstream(q, up *dns.Msg, req Request) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(q)
	m.Rcode = up.Rcode
	if up.CheckingDisabled {
		m.CheckingDisabled = true
	}
	if up.AuthenticatedData {
		m.AuthenticatedData = true
	}
	if up.RecursionAvailable {
		m.RecursionAvailable = true
	}
	m.Answer = up.Answer
	m.Ns = up.Ns
	// Keep the client's EDNS0 shape in the reply, but carry upstream's DO/extended
	// rcode when present.
	if o := up.IsEdns0(); o != nil {
		m.SetEdns0(0, o.Do())
		if o.ExtendedRcode() != 0 || o.Version() != 0 {
			ue := m.IsEdns0()
			ue.SetExtendedRcode(uint16(o.ExtendedRcode()))
			ue.SetVersion(o.Version())
			ue.SetDo(o.Do())
		}
	} else if o := q.IsEdns0(); o != nil {
		m.SetEdns0(0, o.Do())
	} else if e.cfgStore != nil && e.cfgStore.Get().DNS.MaxUDPSize > 512 {
		m.SetEdns0(0, false)
	}
	if len(up.Question) > 0 {
		m.Question = up.Question
	}
	return m
}

// finalize prepares a response for the wire: attach EDNS0, apply the client's
// UDP size limit, and set the AD/CD bits as needed.
func (e *Engine) finalize(msg, q *dns.Msg, req Request) *dns.Msg {
	if msg == nil {
		return servFail(q, "nil response")
	}
	// Preserve the client's question verbatim (including the 0x20 case we may
	// have canonicalised away).
	if len(q.Question) > 0 {
		msg.Question = q.Question
	}
	msg.Id = q.Id
	msg.Response = true
	if q.Opcode == dns.OpcodeQuery {
		msg.RecursionDesired = q.RecursionDesired
	}
	if !req.TCP {
		limit := clientUDPSize(q)
		if limit > 0 && wouldOverflow(msg, limit) {
			// RFC 2181 §9: TC means "ask again over TCP". Truncating to a
			// question-only response is the correct, interoperable choice.
			trunc := new(dns.Msg)
			trunc.SetReply(q)
			trunc.Truncated = true
			if o := q.IsEdns0(); o != nil {
				trunc.SetEdns0(0, o.Do())
			}
			return trunc
		}
	}
	return msg
}

func clientUDPSize(q *dns.Msg) uint16 {
	if o := q.IsEdns0(); o != nil {
		if s := o.UDPSize(); s >= 512 {
			return s
		}
	}
	return 512
}

func wouldOverflow(m *dns.Msg, limit uint16) bool {
	b, err := m.Pack()
	if err != nil {
		return true
	}
	return len(b) > int(limit)
}

func ageTo(rrs []dns.RR, ttl uint32) []dns.RR {
	for _, rr := range rrs {
		rr.Header().Ttl = ttl
	}
	return rrs
}

// prefetch re-resolves a stale key in the background.
func (e *Engine) prefetch(k cache.Key) {
	if e.cache == nil {
		return
	}
	pool := e.pool.Load()
	if pool == nil {
		return
	}
	q := new(dns.Msg)
	q.SetQuestion(dns.Fqdn(k.Name), k.Type)
	q.RecursionDesired = true
	if k.DO {
		q.SetEdns0(1232, true)
	}
	timeout := time.Duration(e.upstreamTimeout.Load())
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	up, addr, err := pool.Exchange(ctx, q, nil, timeout)
	if err != nil || up == nil || up.Truncated {
		return
	}
	e.cache.PutResponse(k, up, addr)
}

// validate rejects malformed queries early.
func validate(q *dns.Msg) error {
	if q == nil {
		return errMalformed
	}
	if q.Response {
		return errMalformed
	}
	if len(q.Question) != 1 {
		return errMalformed
	}
	if q.Opcode != dns.OpcodeQuery && q.Opcode != dns.OpcodeUpdate {
		return errMalformed
	}
	name := q.Question[0].Name
	if name == "" || len(name) > 255*4 {
		return errMalformed
	}
	return nil
}

var errMalformed = errString("malformed query")

type errString string

func (e errString) Error() string { return string(e) }

func formErr(q *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetRcode(q, dns.RcodeFormatError)
	return m
}

func servFail(q *dns.Msg, reason string) *dns.Msg {
	m := new(dns.Msg)
	m.SetRcode(q, dns.RcodeServerFailure)
	return m
}

func nxDomain(q *dns.Msg, ttl uint32) *dns.Msg {
	m := new(dns.Msg)
	m.SetRcode(q, dns.RcodeNameError)
	m.Ns = []dns.RR{syntheticSOA(q, ttl)}
	return m
}

// syntheticSOA is the authority record attached to synthesised negative
// answers. RFC 2308 requires a SOA in NXDOMAIN's authority section for the TTL
// to be honoured, and clients use it to cache the negative result.
func syntheticSOA(q *dns.Msg, ttl uint32) dns.RR {
	name := "."
	if len(q.Question) > 0 {
		labels := strings.Split(strings.TrimSuffix(q.Question[0].Name, "."), ".")
		if len(labels) >= 2 {
			name = strings.Join(labels[len(labels)-2:], ".") + "."
		}
	}
	if ttl == 0 {
		ttl = 60
	}
	soa := &dns.SOA{
		Hdr:     dns.RR_Header{Name: name, Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: ttl},
		Ns:      "ns.fancy-dns.local.",
		Mbox:    "hostmaster.fancy-dns.local.",
		Serial:  uint32(time.Now().Unix()),
		Refresh: 3600,
		Retry:   600,
		Expire:  604800,
		Minttl:  ttl,
	}
	return soa
}

// blockedResponse builds the configured blocklist answer.
func blockedResponse(q *dns.Msg, cfg config.Block, entry *blocklist.Entry) *dns.Msg {
	ttl := uint32(cfg.ResponseTTL.D())
	if ttl == 0 {
		ttl = 60
	}
	switch cfg.DefaultAction {
	case "drop":
		// A drop is handled by the caller, which knows not to answer; reaching
		// here means the caller wants a body, so return a timeout-style empty
		// NOERROR.
		m := new(dns.Msg)
		m.SetReply(q)
		return m
	case "refused":
		return refuseConfig(q)
	case "zero":
		return sinkhole(q, cfg)
	case "redirect":
		return redirect(q, cfg.RedirectName, ttl)
	default:
		return nxDomain(q, ttl)
	}
}

func refuseConfig(q *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetRcode(q, dns.RcodeRefused)
	return m
}

func refuse(q *dns.Msg, _ acl.Decision) *dns.Msg { return refuseConfig(q) }

// sinkhole answers with the configured zero address, choosing the family from
// the question so a v4 client never gets a v6 answer.
func sinkhole(q *dns.Msg, cfg config.Block) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(q)
	m.Authoritative = true
	ttl := uint32(cfg.ResponseTTL.D())
	if ttl == 0 {
		ttl = 60
	}
	if len(q.Question) == 0 {
		return m
	}
	switch q.Question[0].Qtype {
	case dns.TypeA:
		ip := cfg.ResponseIP
		if ip == "" {
			ip = "0.0.0.0"
		}
		rr, _ := dns.NewRR(q.Question[0].Name + " 0 IN A " + ip)
		if rr != nil {
			rr.Header().Ttl = ttl
			m.Answer = []dns.RR{rr}
		}
	case dns.TypeAAAA:
		ip := cfg.ResponseIPv6
		if ip == "" {
			ip = "::"
		}
		rr, _ := dns.NewRR(q.Question[0].Name + " 0 IN AAAA " + ip)
		if rr != nil {
			rr.Header().Ttl = ttl
			m.Answer = []dns.RR{rr}
		}
	default:
		m.Ns = []dns.RR{syntheticSOA(q, ttl)}
	}
	return m
}

// redirect answers with a CNAME to a configured name, plus an A/AAAA for the
// target when we know it (a CNAME to a name that does not resolve would be a
// worse answer than no CNAME at all).
func redirect(q *dns.Msg, target string, ttl uint32) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(q)
	m.Authoritative = true
	if target == "" {
		return m
	}
	if len(q.Question) == 0 {
		return m
	}
	cname, err := dns.NewRR(q.Question[0].Name + " 0 IN CNAME " + dns.Fqdn(target))
	if err != nil {
		return m
	}
	cname.Header().Ttl = ttl
	m.Answer = []dns.RR{cname}
	switch q.Question[0].Qtype {
	case dns.TypeA:
		rr, _ := dns.NewRR(dns.Fqdn(target) + " 0 IN A 127.0.0.1")
		if rr != nil {
			rr.Header().Ttl = ttl
			m.Answer = append(m.Answer, rr)
		}
	case dns.TypeAAAA:
		rr, _ := dns.NewRR(dns.Fqdn(target) + " 0 IN AAAA ::1")
		if rr != nil {
			rr.Header().Ttl = ttl
			m.Answer = append(m.Answer, rr)
		}
	}
	return m
}

func shortErr(err error) string {
	if err == nil {
		return ""
	}
	s := err.Error()
	if len(s) > 120 {
		s = s[:120]
	}
	return s
}

// groupStatTable tracks per-group query and block counts for the UI.
type groupStatTable struct {
	mu    sync.RWMutex
	stats map[string]*groupStat
}

type groupStat struct {
	queries uint64
	blocked uint64
	errors  uint64
	cache   uint64
}

func newGroupStatTable() *groupStatTable {
	return &groupStatTable{stats: make(map[string]*groupStat)}
}

func (t *groupStatTable) observe(group string, res Response) {
	if group == "" {
		group = "(default)"
	}
	t.mu.Lock()
	s, ok := t.stats[group]
	if !ok {
		s = &groupStat{}
		t.stats[group] = s
	}
	s.queries++
	switch res.Action {
	case qlog.ActionBlock:
		s.blocked++
	case qlog.ActionError, qlog.ActionDrop:
		s.errors++
	case qlog.ActionCached:
		s.cache++
	}
	t.mu.Unlock()
}

func (t *groupStatTable) merge(base []acl.GroupSummary) []acl.GroupSummary {
	t.mu.RLock()
	defer t.mu.RUnlock()
	byName := make(map[string]acl.GroupSummary, len(base))
	for _, g := range base {
		byName[g.Name] = g
	}
	// Include groups that were removed from the config but still have traffic.
	for name := range t.stats {
		if _, ok := byName[name]; ok {
			continue
		}
		byName[name] = acl.GroupSummary{Name: name, Policy: string(acl.PolicyAllow)}
	}
	out := make([]acl.GroupSummary, 0, len(byName))
	for name, g := range byName {
		if s, ok := t.stats[name]; ok {
			g.QueryCount = s.queries
			g.BlockedCount = s.blocked
		}
		out = append(out, g)
	}
	return out
}

var _ = strconv.Itoa

// Flush persists buffered state: the cache snapshot and the query log. It is
// wired to the server's flusher interface so the admin API can force a flush.
func (e *Engine) Flush() {
	if p := e.flusher.Load(); p != nil {
		(*p)()
	}
}

// SetFlusher installs the flush callback (set once during app wiring).
func (e *Engine) SetFlusher(fn func()) {
	if fn == nil {
		e.flusher.Store(nil)
		return
	}
	e.flusher.Store(&fn)
}