jarl/internal/stats/collector.go

package stats

import (
	"cmp"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/netip"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"jarl/internal/blocking"
)

// Kinds of top lists.
const (
	KindName     = "name"     // most queried names
	KindClient   = "client"   // clients sending the most queries
	KindNXDomain = "nxdomain" // names answered with NXDOMAIN
	KindServFail = "servfail" // names answered with SERVFAIL
	KindRefused  = "refused"  // clients that were refused
	KindBlocked  = "blocked"  // names blocked (docs/blocking.md)
	// KindBlockedClient: clients with the most blocked queries.
	KindBlockedClient = "blocked_client"
)

// Kinds lists the top-list kinds.
var Kinds = []string{KindName, KindClient, KindNXDomain, KindServFail, KindRefused, KindBlocked, KindBlockedClient}

// OtherZone is the zone key for names outside every local zone (answered
// by recursion, or refused).
const OtherZone = "-"

// Options configure a Collector.
type Options struct {
	QueryLog          bool
	LogFile           string
	StatisticsChannel string // ADDRESS:PORT
	TopN              int
	RetainMinute      time.Duration
	RetainHour        time.Duration
	RetainDay         time.Duration
	Anonymize         bool
	Exclude           []netip.Prefix
	Location          *time.Location // for days and log timestamps (default: local time)
	// EnsureLogging turns named's response logging on (rndc responselog
	// on). It is called at start and when named answers queries but no
	// log lines arrive.
	EnsureLogging func(context.Context) (bool, error)
}

// ZoneIndex returns, per view name ("" when no views are used), the names
// of the zones this server is authoritative for (lowercase, without the
// final dot).
type ZoneIndex func() map[string][]string

// Status describes the collector's health.
type Status struct {
	Enabled           bool       `json:"enabled"`
	QueryLog          bool       `json:"query_log"`
	LogFile           string     `json:"log_file,omitempty"`
	StatisticsChannel string     `json:"statistics_channel"`
	LastPoll          *time.Time `json:"last_poll,omitempty"`
	PollError         string     `json:"poll_error,omitempty"`
	LastLine          *time.Time `json:"last_log_line,omitempty"`
	LinesRead         int64      `json:"log_lines_read"`
	LogError          string     `json:"log_error,omitempty"`
	Warnings          []string   `json:"warnings"`
	FirstData         *time.Time `json:"first_data,omitempty"`
	Retention         struct {
		Minute string `json:"minute"`
		Hour   string `json:"hour"`
		Day    string `json:"day"`
	} `json:"retention"`
}

// Collector gathers statistics into a Store.
type Collector struct {
	opts   Options
	store  *Store
	zones  ZoneIndex
	log    *slog.Logger
	poller *Poller
	tail   *Tailer
	now    func() time.Time

	mu      sync.Mutex
	minutes map[int64]map[string]int64            // minute -> counter -> value
	hours   map[int64]map[string]map[string]int64 // live hours: hour -> kind -> key -> count
	// late holds counts for hours that were already written and added to
	// their day (lines read after a restart); they are added, not
	// replaced.
	late     map[int64]map[string]map[string]int64
	liveHour int64
	prev     *Snapshot
	zoneIdx  map[string][]string
	zoneAt   time.Time
	status   Status
	recentQ  []int64 // queries per minute from the counters, last 10 minutes
	recentL  []int64 // response log lines per minute, last 10 minutes
	linesMin int64
	ensureAt time.Time
	// rewritten holds queries named blocked (client#port|name) until their
	// response line arrives.
	rewritten map[string]bool
}

// maxKeys bounds the per-hour maps; rarely seen keys are dropped first.
const maxKeys = 200_000

// NewCollector creates a collector. Run starts it.
func NewCollector(store *Store, opts Options, zones ZoneIndex, log *slog.Logger) *Collector {
	if opts.Location == nil {
		opts.Location = time.Local
	}
	c := &Collector{opts: opts, store: store, zones: zones, log: log, poller: NewPoller(opts.StatisticsChannel), now: time.Now,
		minutes: map[int64]map[string]int64{}, hours: map[int64]map[string]map[string]int64{}, late: map[int64]map[string]map[string]int64{}}
	if opts.QueryLog {
		c.tail = &Tailer{Path: opts.LogFile}
	}
	c.status = Status{Enabled: true, QueryLog: opts.QueryLog, StatisticsChannel: opts.StatisticsChannel, Warnings: []string{}}
	if opts.QueryLog {
		c.status.LogFile = opts.LogFile
	}
	c.status.Retention.Minute = humanDuration(opts.RetainMinute)
	c.status.Retention.Hour = humanDuration(opts.RetainHour)
	c.status.Retention.Day = humanDuration(opts.RetainDay)
	return c
}

// Store returns the collector's store.
func (c *Collector) Store() *Store { return c.store }

// Options returns the collector's options.
func (c *Collector) Options() Options { return c.opts }

// Status returns the current status.
func (c *Collector) Status() Status {
	c.mu.Lock()
	defer c.mu.Unlock()
	s := c.status
	s.Warnings = append([]string{}, c.status.Warnings...)
	return s
}

func (c *Collector) minuteOf(t time.Time) int64 { return t.Unix() - t.Unix()%60 }
func (c *Collector) hourOf(t time.Time) int64   { return t.Unix() - t.Unix()%3600 }

// dayOf returns the start of the local day containing t.
func (c *Collector) dayOf(t time.Time) int64 {
	t = t.In(c.opts.Location)
	return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, c.opts.Location).Unix()
}

// Run collects until ctx is done.
func (c *Collector) Run(ctx context.Context) {
	c.start(ctx)
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	next := c.nextFlush()
	for {
		select {
		case <-ctx.Done():
			// Write what was collected; a restart continues the minute.
			c.readLog()
			c.flush(context.Background(), true)
			if c.tail != nil {
				c.tail.Close()
			}
			return
		case <-tick.C:
			c.readLog()
			if !c.now().Before(next) {
				c.flush(ctx, false)
				next = c.nextFlush()
			}
		}
	}
}

// nextFlush is two seconds into the next minute, so the minute that just
// ended is complete.
func (c *Collector) nextFlush() time.Time {
	return time.Unix(c.minuteOf(c.now())+62, 0)
}

func (c *Collector) start(ctx context.Context) {
	if c.tail != nil {
		if v, _ := c.store.Meta(ctx, "tail"); v != "" {
			var p Position
			if json.Unmarshal([]byte(v), &p) == nil {
				c.tail.Resume(p)
			}
		}
	}
	if v, _ := c.store.Meta(ctx, "first_data"); v != "" {
		if n, err := strconv.ParseInt(v, 10, 64); err == nil {
			t := time.Unix(n, 0)
			c.status.FirstData = &t
		}
	}
	// Continue the current hour's top lists, and fold hours that ended
	// while jarld was not running into their days.
	cur := c.hourOf(c.now())
	c.liveHour = cur
	b := &Batch{}
	merged, _ := c.store.Meta(ctx, "merged_through")
	if last, err := strconv.ParseInt(merged, 10, 64); err == nil {
		from := max(last+3600, cur-48*3600)
		for h := from; h < cur; h += 3600 {
			for _, k := range Kinds {
				if es, err := c.store.TopsAt(ctx, ResHour, h, k); err == nil && len(es) > 0 {
					b.AddTops(ResDay, c.dayOf(time.Unix(h, 0)), k, es, 3*c.opts.TopN)
				}
			}
		}
	}
	b.SetMeta("merged_through", strconv.FormatInt(cur-3600, 10))
	for _, k := range Kinds {
		es, err := c.store.TopsAt(ctx, ResHour, cur, k)
		if err != nil || len(es) == 0 {
			continue
		}
		m := c.topsMap(c.hours, cur, k)
		for _, e := range es {
			m[e.Key] = e.Count
		}
	}
	if err := c.store.Commit(ctx, b); err != nil {
		c.log.Warn("statistics: catching up failed", "err", err)
	}
	c.ensureLogging(ctx)
	// The first poll is the baseline for the counters.
	if snap, err := c.poller.Poll(ctx); err == nil {
		c.prev = snap
		t := c.now()
		c.status.LastPoll = &t
	} else {
		c.status.PollError = err.Error()
	}
	c.prune(ctx)
	c.log.Info("statistics collector started", "query_log", c.opts.QueryLog, "statistics_channel", c.opts.StatisticsChannel)
}

func (c *Collector) topsMap(hours map[int64]map[string]map[string]int64, hour int64, kind string) map[string]int64 {
	h := hours[hour]
	if h == nil {
		h = map[string]map[string]int64{}
		hours[hour] = h
	}
	m := h[kind]
	if m == nil {
		m = map[string]int64{}
		h[kind] = m
	}
	return m
}

// hourMaps returns the per-kind maps of an hour, creating them.
func (c *Collector) hourMaps(hours map[int64]map[string]map[string]int64, hour int64) map[string]map[string]int64 {
	for _, k := range Kinds {
		c.topsMap(hours, hour, k)
	}
	return hours[hour]
}

func (c *Collector) count(minute int64, key string, n int64) {
	m := c.minutes[minute]
	if m == nil {
		m = map[string]int64{}
		c.minutes[minute] = m
	}
	m[key] += n
}

// readLog reads new response log lines.
func (c *Collector) readLog() {
	if c.tail == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.refreshZones()
	_, err := c.tail.Read(c.handleLine)
	if err != nil {
		c.status.LogError = err.Error()
	} else {
		c.status.LogError = ""
	}
}

// HandleLine processes one log line (exported for tests).
func (c *Collector) HandleLine(line string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.refreshZones()
	c.handleLine(line)
}

func (c *Collector) handleLine(line string) {
	ev, ok := ParseResponse(line, c.opts.Location)
	if !ok {
		if rw, ok := ParseRewrite(line, c.opts.Location); ok {
			c.handleRewrite(rw)
		}
		return
	}
	blocked := c.rewritten[ev.Client.String()+"#"+ev.Port+"|"+ev.Name]
	if blocked {
		delete(c.rewritten, ev.Client.String()+"#"+ev.Port+"|"+ev.Name)
	}
	c.status.LinesRead++
	c.linesMin++
	t := ev.Time
	c.status.LastLine = &t
	for _, p := range c.opts.Exclude {
		if p.Contains(ev.Client) {
			return
		}
	}
	minute, hour := c.minuteOf(ev.Time), c.hourOf(ev.Time)
	view := ev.View
	if view == "" {
		view = "_default"
	}
	c.count(minute, "responses", 1)
	c.count(minute, "view:"+view, 1)
	c.count(minute, "zone:"+c.zoneOf(ev.Name, ev.View), 1)
	client := clientKey(ev.Client, c.opts.Anonymize)
	hours := c.hours
	if hour < c.liveHour {
		hours = c.late
	}
	hours[hour] = c.hourMaps(hours, hour)
	tops := hours[hour]
	tops[KindName][ev.Name]++
	tops[KindClient][client]++
	switch {
	case blocked:
		// Counted by handleRewrite; not a name that doesn't exist.
	case ev.RCode == "NXDOMAIN":
		tops[KindNXDomain][ev.Name]++
	case ev.RCode == "SERVFAIL":
		tops[KindServFail][ev.Name]++
	case ev.RCode == "REFUSED":
		tops[KindRefused][client]++
	}
	for _, m := range tops {
		if len(m) > maxKeys {
			shrink(m)
		}
	}
}

// handleRewrite counts a query that blocking answered: "blocked", and
// "blocklist:NAME" for the list (or "rules") that matched.
func (c *Collector) handleRewrite(rw Rewrite) {
	if rw.Action == "PASSTHRU" {
		return // an allow rule: answered normally
	}
	for _, p := range c.opts.Exclude {
		if p.Contains(rw.Client) {
			return
		}
	}
	if c.rewritten == nil || len(c.rewritten) > 10_000 {
		// Rewrites whose response line never came (named only logs
		// responses while response logging is on) don't pile up.
		c.rewritten = map[string]bool{}
	}
	c.rewritten[rw.Client.String()+"#"+rw.Port+"|"+rw.Name] = true
	minute, hour := c.minuteOf(rw.Time), c.hourOf(rw.Time)
	c.count(minute, "blocked", 1)
	c.count(minute, "blocklist:"+policySource(rw.Via), 1)
	hours := c.hours
	if hour < c.liveHour {
		hours = c.late
	}
	hours[hour] = c.hourMaps(hours, hour)
	hours[hour][KindBlocked][rw.Name]++
	hours[hour][KindBlockedClient][clientKey(rw.Client, c.opts.Anonymize)]++
}

// policySource names what a rewrite came from: a blocklist's name, "rules"
// for the rules entered in Jarl, or the policy zone for others.
func policySource(via string) string {
	via = strings.TrimSuffix(via, ".")
	if strings.HasSuffix(via, "."+blocking.RulesZone) {
		return "rules"
	}
	if i := strings.LastIndex(via, ".blocklist.jarl.invalid"); i > 0 && i+len(".blocklist.jarl.invalid") == len(via) {
		rest := via[:i]
		return rest[strings.LastIndex(rest, ".")+1:]
	}
	return via
}

// shrink drops the rarest keys until the map is well under maxKeys.
func shrink(m map[string]int64) {
	for limit := int64(1); len(m) > maxKeys*8/10; limit++ {
		for k, v := range m {
			if v <= limit {
				delete(m, k)
			}
		}
	}
}

// clientKey is how a client is recorded: its address, or its network
// when anonymizing.
func clientKey(a netip.Addr, anonymize bool) string {
	if !anonymize {
		return a.String()
	}
	bits := 24
	if a.Is6() {
		bits = 48
	}
	p, _ := a.Prefix(bits)
	return p.String()
}

func (c *Collector) refreshZones() {
	if c.zones == nil || c.zoneIdx != nil && c.now().Sub(c.zoneAt) < 5*time.Minute {
		return
	}
	idx := c.zones()
	for v, zs := range idx {
		// Longest first, so the most specific zone matches.
		slices.SortFunc(zs, func(a, b string) int { return cmp.Compare(len(b), len(a)) })
		idx[v] = zs
	}
	c.zoneIdx, c.zoneAt = idx, c.now()
}

// zoneOf returns the local zone a name belongs to in a view, or OtherZone.
func (c *Collector) zoneOf(name, view string) string {
	for _, z := range c.zoneIdx[view] {
		if name == z || z == "." || strings.HasSuffix(name, "."+z) {
			return z
		}
	}
	return OtherZone
}

// flush polls named's counters and writes everything collected for
// finished minutes (all minutes when final).
func (c *Collector) flush(ctx context.Context, final bool) {
	now := c.now()
	cur := c.minuteOf(now)
	snap, perr := c.poller.Poll(ctx)

	c.mu.Lock()
	if perr != nil {
		c.status.PollError = perr.Error()
	} else {
		c.status.PollError = ""
		c.status.LastPoll = &now
		if c.prev != nil {
			d := Delta(c.prev, snap)
			// The change happened during the minute that just ended.
			for k, v := range d {
				c.count(cur-60, k, v)
			}
			c.noteRecent(d["queries"])
		}
		c.prev = snap
	}
	b := &Batch{}
	first := c.status.FirstData == nil
	for m, vals := range c.minutes {
		if m >= cur && !final {
			continue
		}
		t := time.Unix(m, 0)
		b.AddCounters(ResMinute, m, vals)
		b.AddCounters(ResHour, c.hourOf(t), vals)
		b.AddCounters(ResDay, c.dayOf(t), vals)
		delete(c.minutes, m)
		if first && len(vals) > 0 {
			ft := t
			c.status.FirstData = &ft
			first = false
			b.SetMeta("first_data", strconv.FormatInt(m, 10))
		}
	}
	curHour := c.hourOf(now)
	for h, kinds := range c.late {
		for k, m := range kinds {
			if len(m) > 0 {
				b.AddTops(ResHour, h, k, top(m, c.opts.TopN), c.opts.TopN)
				b.AddTops(ResDay, c.dayOf(time.Unix(h, 0)), k, top(m, c.opts.TopN), 3*c.opts.TopN)
			}
		}
		delete(c.late, h)
	}
	for h, kinds := range c.hours {
		for k, m := range kinds {
			b.SetTops(ResHour, h, k, top(m, c.opts.TopN))
		}
		if h < curHour {
			// The hour is over: add it to its day.
			for k, m := range kinds {
				b.AddTops(ResDay, c.dayOf(time.Unix(h, 0)), k, top(m, c.opts.TopN), 3*c.opts.TopN)
			}
			delete(c.hours, h)
			b.SetMeta("merged_through", strconv.FormatInt(h, 10))
		}
	}
	c.liveHour = curHour
	if c.tail != nil {
		pos, _ := json.Marshal(c.tail.Position())
		b.SetMeta("tail", string(pos))
	}
	c.updateWarnings()
	c.mu.Unlock()

	if err := c.store.Commit(ctx, b); err != nil {
		c.log.Error("statistics: writing failed", "err", err)
	}
	// Queries but no log lines for a while: named's response logging may
	// have been turned off (or named restarted without it).
	c.mu.Lock()
	quiet := len(c.recentL) >= 2 && c.recentL[len(c.recentL)-1] == 0 && c.recentL[len(c.recentL)-2] == 0 &&
		c.recentQ[len(c.recentQ)-1] > 0 && c.recentQ[len(c.recentQ)-2] > 0
	c.mu.Unlock()
	if quiet && !final {
		c.ensureLogging(ctx)
	}
	// Prune once an hour.
	if now.Minute() == 0 || final {
		c.prune(ctx)
	}
}

// ensureLogging makes sure named logs responses (at most every 5 minutes).
func (c *Collector) ensureLogging(ctx context.Context) {
	if !c.opts.QueryLog || c.opts.EnsureLogging == nil || c.now().Sub(c.ensureAt) < 5*time.Minute {
		return
	}
	c.ensureAt = c.now()
	changed, err := c.opts.EnsureLogging(ctx)
	switch {
	case err != nil:
		c.log.Warn("statistics: could not turn on named's response logging (rndc responselog on)", "err", err)
	case changed:
		c.log.Info("statistics: turned on named's response logging (rndc responselog on)")
	}
}

// noteRecent records the last ten minutes of query counts and log lines
// for the health warnings.
func (c *Collector) noteRecent(queries int64) {
	c.recentQ = append(c.recentQ, queries)
	c.recentL = append(c.recentL, c.linesMin)
	c.linesMin = 0
	if len(c.recentQ) > 10 {
		c.recentQ, c.recentL = c.recentQ[1:], c.recentL[1:]
	}
}

func (c *Collector) updateWarnings() {
	w := []string{}
	if c.status.PollError != "" {
		w = append(w, c.status.PollError)
	}
	if c.opts.QueryLog {
		var q, l int64
		for i := range c.recentQ {
			q += c.recentQ[i]
			l += c.recentL[i]
		}
		if len(c.recentQ) >= 5 && q > 0 && l == 0 {
			w = append(w, fmt.Sprintf("named answered %d queries in the last %d minutes, but Jarl read no response log lines from %s: check that /var/log/named exists and is writable by named, and that responselog is not turned off", q, len(c.recentQ), c.opts.LogFile))
		}
		if c.status.LogError != "" {
			w = append(w, "reading the response log failed: "+c.status.LogError)
		}
	}
	c.status.Warnings = w
}

// top returns the n largest entries of m.
func top(m map[string]int64, n int) []Entry {
	out := make([]Entry, 0, len(m))
	for k, v := range m {
		out = append(out, Entry{Key: k, Count: v})
	}
	slices.SortFunc(out, func(a, b Entry) int {
		if a.Count != b.Count {
			return cmp.Compare(b.Count, a.Count)
		}
		return cmp.Compare(a.Key, b.Key)
	})
	if len(out) > n {
		out = out[:n]
	}
	return out
}

func (c *Collector) prune(ctx context.Context) {
	now := c.now().Unix()
	for res, keep := range map[string]time.Duration{ResMinute: c.opts.RetainMinute, ResHour: c.opts.RetainHour, ResDay: c.opts.RetainDay} {
		if keep <= 0 {
			continue
		}
		if n, err := c.store.Prune(ctx, res, now-int64(keep.Seconds())); err != nil {
			c.log.Warn("statistics: pruning failed", "err", err)
		} else if n > 0 {
			c.log.Debug("statistics: pruned", "resolution", res, "rows", n)
		}
	}
}

func humanDuration(d time.Duration) string {
	if d%(24*time.Hour) == 0 && d > 0 {
		return strconv.Itoa(int(d/(24*time.Hour))) + "d"
	}
	return d.String()
}