housecarl/internal/server/quota.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"net/netip"
	"sort"
	"strings"
	"time"

	"housecarl/internal/api"
	"housecarl/internal/auth"
	"housecarl/internal/model"
	"housecarl/internal/nft"
	"housecarl/internal/render"
)

// Quotas: nftables counts them, and never resets them. housecarld resets
// each at the start of its period (daily, weekly, monthly), or when asked,
// and saves what they used every few minutes and when it stops, so a
// reboot (whose boot ruleset starts them at zero) doesn't lose it.

const (
	quotaResetsKey = "quota_resets" // rule mark → when its quota was last reset
	QuotaUsageKey  = "quota_usage"  // nft.QuotaUsage, saved
)

// periodStart is when the current period of a reset began (zero for a
// quota reset only by hand), in local time.
func periodStart(now time.Time, reset string, day int) time.Time {
	y, m, d := now.Date()
	midnight := time.Date(y, m, d, 0, 0, 0, 0, now.Location())
	switch reset {
	case "daily":
		return midnight
	case "weekly":
		back := (int(now.Weekday()) + 6) % 7 // Monday
		return midnight.AddDate(0, 0, -back)
	case "monthly":
		if day < 1 {
			day = 1
		}
		start := monthDay(y, m, day, now.Location())
		if start.After(now) {
			py, pm, _ := time.Date(y, m, 1, 0, 0, 0, 0, now.Location()).AddDate(0, -1, 0).Date()
			start = monthDay(py, pm, day, now.Location())
		}
		return start
	}
	return time.Time{}
}

// monthDay is midnight on a day of a month, or on its last day when the
// month is shorter.
func monthDay(y int, m time.Month, day int, loc *time.Location) time.Time {
	last := time.Date(y, m+1, 0, 0, 0, 0, 0, loc).Day()
	if day > last {
		day = last
	}
	return time.Date(y, m, day, 0, 0, 0, 0, loc)
}

// nextReset is when the period after the current one starts.
func nextReset(now time.Time, reset string, day int) time.Time {
	start := periodStart(now, reset, day)
	switch reset {
	case "daily":
		return start.AddDate(0, 0, 1)
	case "weekly":
		return start.AddDate(0, 0, 7)
	case "monthly":
		y, m, _ := start.Date()
		ny, nm, _ := time.Date(y, m, 1, 0, 0, 0, 0, now.Location()).AddDate(0, 1, 0).Date()
		if day < 1 {
			day = 1
		}
		return monthDay(ny, nm, day, now.Location())
	}
	return time.Time{}
}

func (s *Server) quotaResets(ctx context.Context) map[string]time.Time {
	out := map[string]time.Time{}
	if v, err := s.store.GetMeta(ctx, quotaResetsKey); err == nil && v != "" {
		json.Unmarshal([]byte(v), &out)
	}
	return out
}

func (s *Server) saveQuotaResets(ctx context.Context, m map[string]time.Time) {
	b, _ := json.Marshal(m)
	if err := s.store.SetMeta(ctx, quotaResetsKey, string(b)); err != nil {
		s.log.Error("saving the quotas' reset times failed", "err", err)
	}
}

// RunQuotas resets quotas when their period ends and saves their usage,
// until ctx ends (then it saves once more).
func (s *Server) RunQuotas(ctx context.Context) {
	s.checkQuotas(ctx)
	tick := time.NewTicker(time.Minute)
	defer tick.Stop()
	n := 0
	for {
		select {
		case <-ctx.Done():
			s.SaveQuotaUsage(context.Background())
			return
		case <-tick.C:
			s.checkQuotas(ctx)
			if n++; n%5 == 0 {
				s.SaveQuotaUsage(ctx)
			}
		}
	}
}

// checkQuotas resets the quotas whose period has ended since their last
// reset. A quota seen for the first time starts its period now.
func (s *Server) checkQuotas(ctx context.Context) {
	st := s.eng.State()
	resets := s.quotaResets(ctx)
	now := s.nowFn()
	changed := false
	seen := map[string]bool{}
	for _, r := range st.Rules {
		if r.Quota == nil || r.Disabled || r.Raw != "" {
			continue
		}
		key := fmt.Sprintf("r%d", r.ID)
		seen[key] = true
		last, ok := resets[key]
		if !ok {
			resets[key], changed = now, true
			continue
		}
		start := periodStart(now, r.Quota.Reset, r.Quota.Day)
		if start.IsZero() || !last.Before(start) {
			continue
		}
		if err := s.fw.ResetQuota(ctx, r.ID, r.Quota.PerAddress() != ""); err != nil {
			s.log.Error("resetting a quota failed", "rule", r.Label(), "err", err)
			continue
		}
		s.log.Info("quota reset for its new period", "rule", r.Label(), "period", r.Quota.Reset)
		resets[key], changed = now, true
		s.forgetQuotaUsage(ctx, r.ID)
	}
	for key := range resets {
		if !seen[key] {
			delete(resets, key)
			changed = true
		}
	}
	if changed {
		s.saveQuotaResets(ctx, resets)
	}
}

// SaveQuotaUsage saves what the quotas used, for after a reboot.
func (s *Server) SaveQuotaUsage(ctx context.Context) {
	t, err := s.fw.Table(ctx)
	if err != nil {
		return
	}
	b, _ := json.Marshal(t.QuotaUsage)
	if err := s.store.SetMeta(ctx, QuotaUsageKey, string(b)); err != nil {
		s.log.Error("saving the quotas' usage failed", "err", err)
	}
}

// forgetQuotaUsage drops a reset quota from the saved usage, so a restart
// doesn't bring back what it used before.
func (s *Server) forgetQuotaUsage(ctx context.Context, ruleID int64) {
	v, err := s.store.GetMeta(ctx, QuotaUsageKey)
	if err != nil || v == "" {
		return
	}
	var u nft.QuotaUsage
	if json.Unmarshal([]byte(v), &u) != nil {
		return
	}
	name := render.QuotaName(ruleID)
	delete(u.Quotas, name)
	delete(u.QuotaSets, name+"_v4")
	delete(u.QuotaSets, name+"_v6")
	b, _ := json.Marshal(u)
	s.store.SetMeta(ctx, QuotaUsageKey, string(b))
}

// listQuotas answers what each rule's quota used.
func (s *Server) listQuotas(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := require(p, auth.ResLog, false); err != nil {
		return err
	}
	st := s.eng.State()
	var usage nft.QuotaUsage
	if t, err := s.fw.Table(r.Context()); err == nil {
		usage = t.QuotaUsage
	}
	resets := s.quotaResets(r.Context())
	now := s.nowFn()
	out := []api.QuotaInfo{}
	for _, rule := range st.Rules {
		if rule.Quota == nil {
			continue
		}
		q := api.QuotaInfo{RuleID: rule.ID, Rule: rule.Label(), Quota: *rule.Quota}
		name := render.QuotaName(rule.ID)
		if rule.Quota.PerAddress() == "" {
			q.Used = usage.Quotas[name]
		} else {
			v4, v6 := rule.Groups()
			for _, fam := range []string{"_v4", "_v6"} {
				for _, e := range usage.QuotaSets[name+fam] {
					addr := e.Addr
					if fam == "_v4" && v4 < 32 {
						addr = fmt.Sprintf("%s/%d", addr, v4)
					} else if fam == "_v6" && v6 < 128 {
						addr = fmt.Sprintf("%s/%d", addr, v6)
					}
					q.Addresses = append(q.Addresses, api.QuotaAddress{Address: addr, Used: e.Used})
					q.Used += e.Used
				}
			}
			sort.Slice(q.Addresses, func(i, j int) bool { return q.Addresses[i].Used > q.Addresses[j].Used })
		}
		if last, ok := resets[fmt.Sprintf("r%d", rule.ID)]; ok {
			q.LastReset = &last
		}
		if next := nextReset(now, rule.Quota.Reset, rule.Quota.Day); !next.IsZero() {
			q.NextReset = &next
		}
		out = append(out, q)
	}
	writeJSON(w, http.StatusOK, out)
	return nil
}

// resetQuota starts a rule's quota again now: all of it, or one
// address's; or gives some of what it used back, until its next reset.
func (s *Server) resetQuota(w http.ResponseWriter, r *http.Request, p *auth.Principal) error {
	if err := require(p, auth.ResConfig, true); err != nil {
		return err
	}
	o, err := findObject(s.eng.State(), model.KindRule, r.PathValue("ref"))
	if err != nil {
		return err
	}
	rule := o.(*model.Rule)
	if rule.Quota == nil {
		return badRequest("%s has no quota", rule.Label())
	}
	var in api.QuotaReset
	if err := decodeOptional(r, &in); err != nil {
		return err
	}
	if in.Give < 0 {
		return badRequest("give a number of bytes to give back, 1 or more")
	}
	if in.Address == "" && in.Give == 0 {
		if err := s.fw.ResetQuota(r.Context(), rule.ID, rule.Quota.PerAddress() != ""); err != nil {
			return err
		}
		resets := s.quotaResets(r.Context())
		resets[fmt.Sprintf("r%d", rule.ID)] = s.nowFn()
		s.saveQuotaResets(r.Context(), resets)
		s.forgetQuotaUsage(r.Context(), rule.ID)
		s.eng.Audit(r.Context(), p, "quota.reset", rule.Label(), "success", "")
		writeJSON(w, http.StatusOK, map[string]any{"reset": rule.Label()})
		return nil
	}
	name := render.QuotaName(rule.ID)
	adj := nft.QuotaAdjust{Name: name, Give: in.Give}
	if in.Address != "" {
		if rule.Quota.PerAddress() == "" {
			return badRequest("%s has one quota for everything it matches, not one per address: leave out the address", rule.Label())
		}
		key, err := quotaKey(rule, in.Address)
		if err != nil {
			return err
		}
		if !s.quotaHas(r.Context(), name, key) {
			return notFound("%s hasn't used the quota of %s since its last reset", in.Address, rule.Label())
		}
		adj.Addr = key
	}
	if adj.Give == 0 {
		adj.Give = math.MaxInt64 // all of it: the address starts again
	}
	// The next apply writes the quota's new usage into the ruleset.
	s.fw.AdjustQuota(adj)
	if err := s.eng.ApplyCurrent(r.Context()); err != nil {
		return fmt.Errorf("putting the ruleset in force with the quota's new usage: %w", err)
	}
	s.SaveQuotaUsage(r.Context())
	detail, out := "", map[string]any{}
	if in.Give > 0 {
		detail = fmt.Sprintf("gave back %d bytes", in.Give)
		out["given"] = in.Give
	}
	if in.Address != "" {
		detail = strings.TrimSpace(detail + " " + in.Address)
		out["address"] = in.Address
	}
	action := "quota.reset"
	if in.Give > 0 {
		action = "quota.give"
	}
	s.eng.Audit(r.Context(), p, action, rule.Label(), "success", detail)
	out["reset"] = rule.Label()
	writeJSON(w, http.StatusOK, out)
	return nil
}

// quotaKey is the entry of a per-address quota an address counts under:
// the address, or its group's first address (10.0.10.0 for a /24; an
// IPv6 address's /64).
func quotaKey(rule *model.Rule, s string) (string, error) {
	var a netip.Addr
	if p, err := netip.ParsePrefix(s); err == nil {
		a = p.Addr()
	} else if a, err = netip.ParseAddr(s); err != nil {
		return "", badRequest("%q isn't an address", s)
	}
	a = a.Unmap()
	v4, v6 := rule.Groups()
	bits := v4
	if a.Is6() {
		bits = v6
	}
	p, err := a.Prefix(bits)
	if err != nil {
		return "", badRequest("%q: %v", s, err)
	}
	return p.Addr().String(), nil
}

// quotaHas reports whether a per-address quota has an entry for key.
func (s *Server) quotaHas(ctx context.Context, name, key string) bool {
	t, err := s.fw.Table(ctx)
	if err != nil {
		return false
	}
	for _, set := range []string{name + "_v4", name + "_v6"} {
		for _, e := range t.QuotaSets[set] {
			if e.Addr == key {
				return true
			}
		}
	}
	return false
}