smallwebwaf/internal/state/state.go

// Package state keeps smallwebwaf's state in JSON files in
// SWWAF_STATE_DIR, as the "Persistent state" section of SPEC.md describes:
// bans.json holds the bans, clients.json each client's counters and
// history, and lookups.json GeoJS's answers. Load reads them at start,
// Watch takes in an admin's edit of one while smallwebwaf runs, and Run
// and WriteAll write them. The disk is read and written outside the
// parts' locks, which are held only to take a snapshot or to put in what
// a file holds, so that no request waits on the disk.
package state

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"
	"net/netip"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"sneak.berlin/go/smallwebwaf/internal/bans"
	"sneak.berlin/go/smallwebwaf/internal/lookup"
	"sneak.berlin/go/smallwebwaf/internal/metrics"
	"sneak.berlin/go/smallwebwaf/internal/ratelimit"
)

// version is the version of the files' format, the only one read.
const version = 1

// fileMode lets the smallwebwaf user alone read and write the files, which
// hold visitors' addresses.
const fileMode = 0o600

// The state files' names.
const (
	bansJSON    = "bans.json"
	clientsJSON = "clients.json"
	lookupsJSON = "lookups.json"
)

var (
	errVersion = errors.New("unknown version")
	// errMissing is for an entry without a field it needs.
	errMissing = errors.New("has no")
	errCause   = errors.New("is not limit, attack or admin")
)

// Params are what Load needs.
type Params struct {
	// Dir is the directory of the state files (SWWAF_STATE_DIR).
	Dir string
	// WriteDelay is how long after a ban is made bans.json is written
	// (SWWAF_STATE_WRITE_DELAY), and CounterInterval how often every file
	// is (SWWAF_STATE_COUNTER_INTERVAL).
	WriteDelay      time.Duration
	CounterInterval time.Duration
	// Ledger, Limiter and GeoJS hold the state.
	Ledger  *bans.Ledger
	Limiter *ratelimit.Limiter
	GeoJS   *lookup.GeoJS
	// Now tells the time by which the counters' buckets run out, normally
	// time.Now in UTC.
	Now func() time.Time
	// ProcessLog receives what was read and taken in, the edits set aside,
	// and the writes that fail.
	ProcessLog *slog.Logger
	// Metrics count each file's writes, and the edits taken in and set
	// aside.
	Metrics *metrics.Metrics
}

// Files are the state files of a running smallwebwaf.
type Files struct {
	params Params

	// mu is held while a file is read for an edit, and while it is
	// written, so that Watch and the writes take turns. No request takes
	// it.
	mu sync.Mutex
	// sums are the SHA-256 sums of what each file held, by name, when
	// smallwebwaf last read or wrote it. A file that holds anything else
	// has been edited since.
	sums map[string][sha256.Size]byte
}

// bansFile is bans.json, indented for an admin to read and edit.
type bansFile struct {
	Version int        `json:"version"`
	Bans    []banEntry `json:"bans"`
}

// banEntry is a ban as bans.json holds it: a permanent ban's expires is
// null, a ban an admin added may have no cause, which makes it an
// admin's, and lifted is left out until an admin lifts the ban.
type banEntry struct {
	Netblock netip.Prefix `json:"netblock"`
	Start    time.Time    `json:"start"`
	Expires  *time.Time   `json:"expires"`
	Cause    string       `json:"cause"`
	Reason   string       `json:"reason,omitempty"`
	Lifted   *time.Time   `json:"lifted,omitempty"`
	Notes    bans.Notes   `json:"notes"`
}

// clientsFile is clients.json, with each client on a line of its own.
type clientsFile struct {
	Version int                `json:"version"`
	Clients []ratelimit.Client `json:"clients"`
}

// lookupsFile is lookups.json, with each answer on a line of its own.
type lookupsFile struct {
	Version int             `json:"version"`
	Lookups []lookup.Answer `json:"lookups"`
}

// stateFile is the struct of a state file. Once the file is decoded, its
// check refuses the first entry without a field it needs, which would
// otherwise be read as something the entry does not say. data is the
// file, for a field that may be null or "" but not left out, which the
// struct cannot tell apart.
type stateFile interface {
	check(data []byte) error
}

// Load checks that files can be written in Dir, and reads the state files
// in it into the ledger, the limiter and GeoJS. A missing file is empty
// state, as on a first start. A file that does not parse, has an unknown
// version, or has an entry without a field it needs, is an error that
// names the file and, where the JSON decoder tells it, the line and
// column, or else the entry.
func Load(params Params) (*Files, error) {
	err := checkWritable(params.Dir)
	if err != nil {
		return nil, fmt.Errorf("SWWAF_STATE_DIR cannot be written: %w", err)
	}

	f := &Files{params: params, sums: map[string][sha256.Size]byte{}}

	bansRead, bansErr := f.read(bansJSON)
	clientsRead, clientsErr := f.read(clientsJSON)
	lookupsRead, lookupsErr := f.read(lookupsJSON)

	err = errors.Join(bansErr, clientsErr, lookupsErr)
	if err != nil {
		return nil, err
	}

	params.ProcessLog.Info("read the state files", "directory", params.Dir,
		"bans", bansRead, "clients", clientsRead, "lookups", lookupsRead)

	return f, nil
}

// Run writes bans.json WriteDelay after a ban is made, with every ban
// made in between, and every file every CounterInterval, until ctx is
// done. A write that fails is logged, and the file is written again at
// its next write. Each write takes in an admin's edit of its file first,
// as writeFile describes.
func (f *Files) Run(ctx context.Context) {
	interval := time.NewTicker(f.params.CounterInterval)
	defer interval.Stop()

	var bansDue <-chan time.Time // nil while no ban waits to be written

	for {
		select {
		case <-ctx.Done():
			return
		case <-f.params.Ledger.Changed():
			if bansDue == nil {
				bansDue = time.After(f.params.WriteDelay)
			}
		case <-bansDue:
			bansDue = nil

			f.logFailure(f.writeFile(bansJSON))
		case <-interval.C:
			f.logFailure(f.WriteAll())
		}
	}
}

// WriteAll writes every state file, as smallwebwaf stops. A file that
// fails does not keep the others from being written.
func (f *Files) WriteAll() error {
	return errors.Join(f.writeFile(bansJSON), f.writeFile(clientsJSON),
		f.writeFile(lookupsJSON))
}

// Watch watches Dir until ctx is done, and takes in an admin's edit of a
// state file as soon as it is saved: what the file holds replaces what
// smallwebwaf held for it. An edit that does not parse is left for the
// file's next write, which sets it aside, since a file can be read while
// an editor is still writing it. If Dir cannot be watched, that is
// logged, and an edit is taken in only before its file is written.
func (f *Files) Watch(ctx context.Context) {
	watcher, err := fsnotify.NewWatcher()
	if err == nil {
		defer func() {
			_ = watcher.Close()
		}()

		err = watcher.Add(f.params.Dir)
	}

	if err != nil {
		f.params.ProcessLog.Error("cannot watch the state files for edits",
			"error", err.Error())

		return
	}

	f.params.ProcessLog.Info("watching the state files for edits",
		"directory", f.params.Dir)

	for {
		select {
		case <-ctx.Done():
			return
		case event := <-watcher.Events:
			switch name := filepath.Base(event.Name); name {
			case bansJSON, clientsJSON, lookupsJSON:
				f.fileChanged(name)
			}
		case err = <-watcher.Errors:
			f.params.ProcessLog.Warn("watching the state files failed",
				"error", err.Error())
		}
	}
}

// logFailure logs a write that failed.
func (f *Files) logFailure(err error) {
	if err != nil {
		f.params.ProcessLog.Error("writing the state files failed",
			"error", err.Error())
	}
}

// fileChanged takes in what the state file name holds, as Watch sees it
// change, if that is an edit made since smallwebwaf last read or wrote
// the file. A file that cannot be read or does not parse is left for its
// next write.
func (f *Files) fileChanged(name string) {
	f.mu.Lock()
	defer f.mu.Unlock()

	data, changed, err := f.readChanged(name)
	if err != nil || !changed {
		return
	}

	_ = f.takeInEdit(name, data)
}

// takeInEdit takes in data, an edit of the state file name, as takeIn
// does, and counts and logs it. Every edit taken in while smallwebwaf
// runs, by Watch or by a write, is taken in here. An edit that does not
// parse is neither counted nor logged, and takeIn's error returned.
func (f *Files) takeInEdit(name string, data []byte) error {
	_, err := f.takeIn(name, data, true)
	if err != nil {
		return err
	}

	// Counted before it is logged, so that the count is there once the
	// log line is.
	f.params.Metrics.StateFileEditTakenIn(name)
	f.params.ProcessLog.Info("took in an edit of a state file",
		"file", filepath.Join(f.params.Dir, name))

	return nil
}

// read takes in the state file name at start, and returns how many
// entries it holds. A missing file holds none.
func (f *Files) read(name string) (int, error) {
	data, changed, err := f.readChanged(name)
	if err != nil || !changed {
		return 0, err
	}

	return f.takeIn(name, data, false)
}

// readChanged returns what the state file name holds, and whether that
// has changed since smallwebwaf last read or wrote the file, as it has
// for a file smallwebwaf never read or wrote. A missing file has not
// changed: it is written again at its next write.
func (f *Files) readChanged(name string) ([]byte, bool, error) {
	path := filepath.Join(f.params.Dir, name)

	data, err := os.ReadFile(path) //nolint:gosec // a state file, in SWWAF_STATE_DIR
	if errors.Is(err, fs.ErrNotExist) {
		return nil, false, nil
	}

	if err != nil {
		return nil, false, err
	}

	return data, sha256.Sum256(data) != f.sums[name], nil
}

// takeIn parses data, what the state file name holds, puts it into the
// part that keeps that state, in place of what the part held, and returns
// how many entries the file holds. edit is whether data is an admin's
// edit taken in while smallwebwaf runs, rather than the file read at the
// start. An error names the file and, where the JSON decoder tells it,
// the line and column, or else the entry.
func (f *Files) takeIn(name string, data []byte, edit bool) (int, error) {
	path := filepath.Join(f.params.Dir, name)

	var entries int

	switch name {
	case bansJSON:
		var file bansFile

		err := parse(path, data, &file)
		if err != nil {
			return 0, err
		}

		held := make([]bans.Ban, 0, len(file.Bans))
		for _, entry := range file.Bans {
			held = append(held, entry.ban())
		}

		if edit {
			f.params.Ledger.LoadEdit(held)
		} else {
			f.params.Ledger.Load(held)
		}

		entries = len(held)
	case clientsJSON:
		var file clientsFile

		err := parse(path, data, &file)
		if err != nil {
			return 0, err
		}

		f.params.Limiter.Load(file.Clients, f.params.Now())
		entries = len(file.Clients)
	case lookupsJSON:
		var file lookupsFile

		err := parse(path, data, &file)
		if err != nil {
			return 0, err
		}

		f.params.GeoJS.Load(file.Lookups)
		entries = len(file.Lookups)
	}

	f.sums[name] = sha256.Sum256(data)

	return entries, nil
}

// writeFile writes the state file name from what smallwebwaf holds. An
// edit made since smallwebwaf last read or wrote the file is taken in
// first, so that it is not overwritten, or set aside if it does not
// parse. A file that cannot be read, or an edit that cannot be set
// aside, is left as it is, and the write given up. Every write is counted
// in the metrics, and one that fails or is given up as a failure.
func (f *Files) writeFile(name string) error {
	f.mu.Lock()
	defer f.mu.Unlock()

	data, changed, err := f.readChanged(name)
	if err == nil && changed {
		err = f.takeInEdit(name, data)
		if err != nil {
			err = f.setAside(name, err)
		}
	}

	if err == nil {
		data, err = f.encode(name)
		if err != nil {
			err = fmt.Errorf("encode %s: %w", name, err)
		}
	}

	if err == nil {
		err = write(f.params.Dir, name, data)
	}

	if err == nil {
		// The file holds data from here on, even if the directory sync
		// fails, so that its next read does not take it for an admin's
		// edit.
		f.sums[name] = sha256.Sum256(data)
		err = syncDirectory(f.params.Dir)
	}

	f.params.Metrics.StateFileWritten(name, len(data), err)

	return err
}

// setAside renames the state file name, an edit that does not parse with
// parseErr, to name.bad, for the admin to mend, and logs it with where in
// the file the error is. If the rename fails, the edit is left as it is,
// and the error returned is parseErr joined with the rename's.
func (f *Files) setAside(name string, parseErr error) error {
	path := filepath.Join(f.params.Dir, name)

	err := os.Rename(path, path+".bad")
	if err != nil {
		return errors.Join(parseErr, err)
	}

	f.params.ProcessLog.Error("set aside an edit of a state file that does not parse",
		"file", path+".bad", "error", parseErr.Error())
	f.params.Metrics.StateFileEditSetAside(name)

	return nil
}

// encode returns the state file name as smallwebwaf writes it, from a
// snapshot of the part that keeps that state.
func (f *Files) encode(name string) ([]byte, error) {
	switch name {
	case bansJSON:
		held := f.params.Ledger.Snapshot()

		file := bansFile{Version: version, Bans: make([]banEntry, 0, len(held))}
		for _, ban := range held {
			file.Bans = append(file.Bans, newBanEntry(ban))
		}

		data, err := json.MarshalIndent(file, "", "    ")
		if err != nil {
			return nil, err
		}

		return append(data, '\n'), nil
	case clientsJSON:
		return encodeOnePerLine("clients", f.params.Limiter.Snapshot())
	default: // lookups.json
		return encodeOnePerLine("lookups", f.params.GeoJS.Snapshot())
	}
}

// newBanEntry returns ban as bans.json holds it.
func newBanEntry(ban bans.Ban) banEntry {
	entry := banEntry{
		Netblock: ban.Netblock, Start: ban.Start, Cause: ban.Cause, Reason: ban.Reason,
		Notes: ban.Notes,
	}
	if !ban.Permanent() {
		entry.Expires = &ban.Expires
	}

	if !ban.Lifted.IsZero() {
		entry.Lifted = &ban.Lifted
	}

	return entry
}

// ban returns the ban an entry of bans.json holds.
func (e banEntry) ban() bans.Ban {
	ban := bans.Ban{
		Netblock: e.Netblock, Start: e.Start, Cause: e.Cause, Reason: e.Reason,
		Notes: e.Notes,
	}
	if e.Expires != nil {
		ban.Expires = *e.Expires
	}

	if e.Lifted != nil {
		ban.Lifted = *e.Lifted
	}

	return ban
}

// check refuses a ban without a netblock, which would refuse every IPv6
// client, a start, from which the length of the netblock's next ban is
// worked out, or an expires, which would make it permanent. A permanent
// ban's expires is null, which Bans cannot tell from a missing one, so
// each expires is read again as written. A cause other than limit,
// attack or admin, most likely misspelt, is refused too.
func (f *bansFile) check(data []byte) error {
	var written struct {
		Bans []struct {
			Expires json.RawMessage `json:"expires"`
		} `json:"bans"`
	}

	err := json.Unmarshal(data, &written)
	if err != nil {
		return err
	}

	for i, entry := range f.Bans {
		switch {
		case !entry.Netblock.IsValid():
			return missing(i, "netblock")
		case entry.Start.IsZero():
			return missing(i, "start")
		case written.Bans[i].Expires == nil:
			return missing(i, "expires")
		case entry.Cause != "" && entry.Cause != bans.CauseLimit &&
			entry.Cause != bans.CauseAttack && entry.Cause != bans.CauseAdmin:
			return fmt.Errorf("entry %d's cause %q %w", i+1, entry.Cause, errCause)
		}
	}

	return nil
}

// check refuses a client without its address, which would count nobody's
// requests, or with requests in a window but no start, which would drop
// them and give the client a fresh allowance.
func (f *clientsFile) check([]byte) error {
	for i, client := range f.Clients {
		switch {
		case !client.Client.IsValid():
			return missing(i, "client")
		case countsWithoutStart(client.Minute):
			return missing(i, "minute.start")
		case countsWithoutStart(client.Hour):
			return missing(i, "hour.start")
		case countsWithoutStart(client.Day):
			return missing(i, "day.start")
		}
	}

	return nil
}

// check refuses an answer without a client, which would answer for
// nobody, a country, which would place the client nowhere, or the time
// GeoJS gave it, which would drop it. "" is the country of a client
// GeoJS cannot place, which Lookups cannot tell from a missing one, so
// each country is read again as written.
func (f *lookupsFile) check(data []byte) error {
	var written struct {
		Lookups []struct {
			Country *string `json:"country"`
		} `json:"lookups"`
	}

	err := json.Unmarshal(data, &written)
	if err != nil {
		return err
	}

	for i, answer := range f.Lookups {
		switch {
		case !answer.Client.IsValid():
			return missing(i, "client")
		case written.Lookups[i].Country == nil:
			return missing(i, "country")
		case answer.Answered.IsZero():
			return missing(i, "answered")
		}
	}

	return nil
}

// countsWithoutStart reports whether b holds requests but no start, which
// places them in time.
func countsWithoutStart(b ratelimit.Buckets) bool {
	return b.Start.IsZero() && (b.Current != 0 || b.Previous != 0)
}

// missing returns the error for entry i, counted from 0, of a state file,
// which has no field.
func missing(i int, field string) error {
	return fmt.Errorf("entry %d %w %q", i+1, errMissing, field)
}

// encodeOnePerLine encodes a state file whose entries, under key, are one
// to a line, so that grep shows everything about one client.
func encodeOnePerLine[E any](key string, entries []E) ([]byte, error) {
	var b bytes.Buffer

	fmt.Fprintf(&b, "{\n    \"version\": %d,\n    %q: [", version, key)

	for i, entry := range entries {
		line, err := json.Marshal(entry)
		if err != nil {
			return nil, err
		}

		if i > 0 {
			b.WriteString(",")
		}

		b.WriteString("\n        ")
		b.Write(line)
	}

	b.WriteString("\n    ]\n}\n")

	return b.Bytes(), nil
}

// checkWritable makes a file in dir and removes it again.
func checkWritable(dir string) error {
	file, err := os.CreateTemp(dir, "write-check-*")
	if err != nil {
		return err
	}

	return errors.Join(file.Close(), os.Remove(file.Name()))
}

// parse reads data, what the state file at path holds, into file, a
// pointer to that file's struct, and checks its entries.
func parse(path string, data []byte, file stateFile) error {
	// The version is read first, so that a file of another version is
	// refused for that, and not for an entry this version cannot read.
	var header struct {
		Version int `json:"version"`
	}

	err := json.Unmarshal(data, &header)
	if err == nil && header.Version != version {
		err = fmt.Errorf("%w %d, where this smallwebwaf reads version %d",
			errVersion, header.Version, version)
	}

	if err == nil {
		decoder := json.NewDecoder(bytes.NewReader(data))
		// A field this version does not know is most likely misspelt, and
		// its value would be lost without a word.
		decoder.DisallowUnknownFields()
		err = decoder.Decode(file)
	}

	if err == nil {
		err = file.check(data)
	}

	if err != nil {
		return fmt.Errorf("%s%s: %w", path, position(data, err), err)
	}

	return nil
}

// position returns where in data err was found, as ", line L, column C"
// of the last byte the JSON decoder read, or "" when err does not tell.
func position(data []byte, err error) string {
	var (
		syntaxErr *json.SyntaxError
		typeErr   *json.UnmarshalTypeError
		read      int64
	)

	switch {
	case errors.As(err, &syntaxErr):
		read = syntaxErr.Offset
	case errors.As(err, &typeErr):
		read = typeErr.Offset
	default:
		return ""
	}

	before := data[:max(min(read, int64(len(data)))-1, 0)]
	line := bytes.Count(before, []byte("\n")) + 1
	column := len(before) - bytes.LastIndexByte(before, '\n')

	return fmt.Sprintf(", line %d, column %d", line, column)
}

// write writes data to the file name in dir so that a crash at any
// moment leaves either the old file or the new one, whole: data goes to a
// temporary file in the same directory, which is synced and renamed over
// name. syncDirectory must follow, so that the rename lasts.
func write(dir, name string, data []byte) error {
	path := filepath.Join(dir, name)
	temporary := path + ".tmp"

	err := writeSynced(temporary, data)
	if err == nil {
		err = os.Rename(temporary, path)
	}

	if err != nil {
		_ = os.Remove(temporary)
	}

	return err
}

// syncDirectory syncs dir to the disk, so that a rename in it lasts.
func syncDirectory(dir string) error {
	directory, err := os.Open(dir) //nolint:gosec // SWWAF_STATE_DIR itself
	if err != nil {
		return err
	}

	return errors.Join(directory.Sync(), directory.Close())
}

// writeSynced writes data to the file at path, and syncs it to the disk.
func writeSynced(path string, data []byte) error {
	//nolint:gosec // a state file's temporary file, in SWWAF_STATE_DIR
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, fileMode)
	if err != nil {
		return err
	}

	_, err = file.Write(data)
	if err == nil {
		err = file.Sync()
	}

	return errors.Join(err, file.Close())
}