guard/internal/state/state.go

// Package state keeps what Guard must remember across a restart: the
// bans it made (for repeat offenders' longer bans, and guard status),
// and the bans Housecarl hasn't taken yet.
package state

import (
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// MaxRecent is how many recent bans are kept for guard status.
const MaxRecent = 200

// Ban is a ban Guard made (or is waiting to make).
type Ban struct {
	Time    time.Time `json:"time"`
	Watch   string    `json:"watch"`
	List    string    `json:"list"`
	Address string    `json:"address"`
	// Group is what is counted for repeats: the address, or its IPv6
	// prefix.
	Group    string `json:"group"`
	Protocol string `json:"protocol,omitempty"`
	Ports    string `json:"ports,omitempty"`
	Seconds  int    `json:"seconds"`
	Reason   string `json:"reason,omitempty"`
	// Result is what Housecarl said: "banned", "renewed", or why not.
	Result string `json:"result,omitempty"`
}

// State is Guard's memory.
type State struct {
	mu   sync.Mutex
	path string
	// Offences are the times each group was banned.
	Offences map[string][]time.Time `json:"offences"`
	Recent   []Ban                  `json:"recent"`
	Pending  []Ban                  `json:"pending"`
}

// Load reads the state (an empty one if the file isn't there).
func Load(path string) (*State, error) {
	s := &State{path: path, Offences: map[string][]time.Time{}}
	data, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		return s, nil
	}
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal(data, s); err != nil {
		return nil, err
	}
	if s.Offences == nil {
		s.Offences = map[string][]time.Time{}
	}
	return s, nil
}

// Save writes the state atomically.
func (s *State) Save() error {
	s.mu.Lock()
	data, err := json.Marshal(s)
	s.mu.Unlock()
	if err != nil || s.path == "" {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(s.path), 0o750); err != nil {
		return err
	}
	tmp := s.path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o640); err != nil {
		return err
	}
	return os.Rename(tmp, s.path)
}

// Repeats counts a group's bans since a time.
func (s *State) Repeats(group string, since time.Time) int {
	s.mu.Lock()
	defer s.mu.Unlock()
	n := 0
	for _, t := range s.Offences[group] {
		if t.After(since) {
			n++
		}
	}
	return n
}

// Banned records a ban Housecarl made.
func (s *State) Banned(b Ban) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Offences[b.Group] = append(s.Offences[b.Group], b.Time)
	s.recent(b)
}

// Failed records a ban Housecarl refused.
func (s *State) Failed(b Ban) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.recent(b)
}

func (s *State) recent(b Ban) {
	s.Recent = append(s.Recent, b)
	if len(s.Recent) > MaxRecent {
		s.Recent = append([]Ban(nil), s.Recent[len(s.Recent)-MaxRecent:]...)
	}
}

// SetPending replaces the bans waiting for Housecarl.
func (s *State) SetPending(p []Ban) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Pending = append([]Ban(nil), p...)
}

// Snapshot returns the recent bans (newest first) and the pending ones.
func (s *State) Snapshot() (recent, pending []Ban) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := len(s.Recent) - 1; i >= 0; i-- {
		recent = append(recent, s.Recent[i])
	}
	return recent, append([]Ban(nil), s.Pending...)
}

// Forget drops offences older than a time.
func (s *State) Forget(before time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for g, ts := range s.Offences {
		var keep []time.Time
		for _, t := range ts {
			if t.After(before) {
				keep = append(keep, t)
			}
		}
		if len(keep) == 0 {
			delete(s.Offences, g)
		} else {
			s.Offences[g] = keep
		}
	}
}