nuntius/internal/handler/state.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

//go:build linux || freebsd

package handler

import (
	"encoding/json"
	"log/slog"
	"os"
	"path/filepath"
	"time"
)

// rateLimitSnapshot is the on-disk shape of persisted limiter buckets.
// Schema lets future formats be detected instead of misparsed.
type rateLimitSnapshot struct {
	Schema int                                   `json:"schema"`
	Saved  time.Time                             `json:"saved"`
	Forms  map[string]map[string]persistedBucket `json:"forms"`
}

// persistedBucket is the exported wire form of an internal bucket.
type persistedBucket struct {
	Tokens float64   `json:"tokens"`
	Last   time.Time `json:"last"`
}

const snapshotSchema = 1

// statePath returns the snapshot file location under the data directory.
func (h *ContactHandler) statePath() string {
	return filepath.Join(h.dataDir, "ratelimit-snapshot.json")
}

// PersistState writes every rate-limit bucket to a snapshot file next to
// the newsletter logs. The write is atomic (temp file plus rename) and
// best-effort: a failure is logged, never fatal, because losing buckets
// only resets limits to their startup defaults. Call it during shutdown,
// before Close.
func (h *ContactHandler) PersistState() {
	if h.dataDir == "" {
		return
	}
	snap := rateLimitSnapshot{
		Schema: snapshotSchema,
		Saved:  time.Now().UTC(),
		Forms:  make(map[string]map[string]persistedBucket, len(h.rateLimits)),
	}
	for path, lim := range h.rateLimits {
		live := lim.snapshot()
		wire := make(map[string]persistedBucket, len(live))
		for ip, b := range live {
			wire[ip] = persistedBucket{Tokens: b.tokens, Last: b.last}
		}
		snap.Forms[path] = wire
	}

	line, err := json.Marshal(snap)
	if err != nil {
		slog.Warn("rate limit snapshot marshal failed", "err", err)
		return
	}
	dir := filepath.Dir(h.statePath())
	if err := os.MkdirAll(dir, 0o755); err != nil {
		slog.Warn("rate limit snapshot mkdir failed", "dir", dir, "err", err)
		return
	}
	tmp := h.statePath() + ".tmp"
	if err := os.WriteFile(tmp, line, 0o600); err != nil {
		slog.Warn("rate limit snapshot write failed", "err", err)
		return
	}
	if err := os.Rename(tmp, h.statePath()); err != nil {
		slog.Warn("rate limit snapshot rename failed", "err", err)
	}
}

// restoreState loads the previous snapshot, if any, back into the fresh
// rate limiters. Each limiter drops entries older than its configured
// bucket age; corrupt files are ignored with a warning, and nothing here
// is fatal: a missing or broken snapshot behaves like an empty one.
func (h *ContactHandler) restoreState() {
	if h.dataDir == "" {
		return
	}
	raw, err := os.ReadFile(h.statePath())
	if err != nil {
		return // no snapshot yet: the common first-start path
	}
	var snap rateLimitSnapshot
	if err := json.Unmarshal(raw, &snap); err != nil {
		slog.Warn("ignoring corrupt rate limit snapshot", "path", h.statePath(), "err", err)
		return
	}
	if snap.Schema != snapshotSchema {
		slog.Warn("ignoring rate limit snapshot with unknown schema",
			"path", h.statePath(), "schema", snap.Schema)
		return
	}
	now := time.Now()
	restoredForms := 0
	for path, wire := range snap.Forms {
		lim, ok := h.rateLimits[path]
		if !ok {
			continue // form removed from config since the snapshot
		}
		entries := make(map[string]bucket, len(wire))
		for ip, pb := range wire {
			entries[ip] = bucket{tokens: pb.Tokens, last: pb.Last}
		}
		lim.restore(entries, now)
		restoredForms++
	}
	slog.Info("restored rate limit buckets", "file", h.statePath(), "forms", restoredForms)
}