smallwebwaf/internal/bans/bans.go

// Package bans is the ban ledger: the bans smallwebwaf makes on the
// netblocks of clients that break a rate limit, with their notes, as the
// "Bans" section of SPEC.md describes. The bans are kept in memory, and
// written to bans.json and read from it by the state package.
package bans

import (
	"net/netip"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/golang-lru/v2/simplelru"
)

// repeatFactor is how many times as long as the netblock's last ban a ban
// for a limit broken again within the repeat window lasts.
const repeatFactor = 3

// maxTextBytes is how much of each text in a ban's notes is kept.
const maxTextBytes = 256

// Rules are how long a ban for a broken limit lasts, and how many bans
// are held.
type Rules struct {
	// LimitBanDuration is how long a first ban lasts.
	LimitBanDuration time.Duration
	// LimitBanRepeatWindow is how soon after the netblock's last ban
	// ended a broken limit counts as a repeat, which bans for
	// repeatFactor times as long as that ban.
	LimitBanRepeatWindow time.Duration
	// MaxBanDuration is the longest ban; a ban that would be longer is
	// permanent instead.
	MaxBanDuration time.Duration
	// MaxBans is the most bans held, at least one. Past it, the earliest
	// ban of the netblock that has gone longest without a request is
	// dropped.
	MaxBans int
}

// Ban is a ban on a netblock for a broken limit, the only kind of ban
// smallwebwaf makes so far.
type Ban struct {
	Netblock netip.Prefix
	Start    time.Time
	// Expires is when the ban ends, zero for a permanent ban.
	Expires time.Time
	Notes   Notes
}

// Permanent reports whether the ban never runs out.
func (b Ban) Permanent() bool {
	return b.Expires.IsZero()
}

// ActiveAt reports whether the ban refuses requests at now.
func (b Ban) ActiveAt(now time.Time) bool {
	return b.Permanent() || now.Before(b.Expires)
}

// Notes are what an admin needs to decide whether to lift a ban. The
// JSON names are those of bans.json.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type Notes struct {
	// Country is the client's country, when it was looked up.
	Country string `json:"country"`
	// Limit, Window and Count are the limit that was broken, its window,
	// "minute", "hour" or "day", and the count reached: the client's
	// requests in the window, the one that broke the limit included.
	// These are the requests that counted toward the ban, and the window
	// is the time over which they came.
	Limit  int64   `json:"limit"`
	Window string  `json:"window"`
	Count  float64 `json:"count"`
	// Request is the request that broke the limit.
	Request Request `json:"request"`
	// Requests is how many requests the netblock has sent since it was
	// first seen, and Refused how many of them the ban has refused so
	// far. Both go up with each request the ban refuses.
	Requests int64 `json:"requests"`
	Refused  int64 `json:"refused"`
	// EarlierBans is how many bans the netblock had before this one.
	EarlierBans int `json:"earlier_bans"`
}

// Request is a request in a ban's notes. Each text is cut to 256 bytes.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type Request struct {
	Time   time.Time `json:"time"`
	Method string    `json:"method"`
	Host   string    `json:"host"`
	// Path is the path with its query string.
	Path string `json:"path"`
	// Status is what the client was sent, 0 if nothing was.
	Status    int    `json:"status"`
	UserAgent string `json:"user_agent"`
}

// Ledger holds the bans. It is safe for concurrent use.
type Ledger struct {
	rules Rules
	// changed receives a value when a ban is made, unless one is waiting
	// already.
	changed chan struct{}

	mu sync.Mutex
	// netblocks holds each banned netblock's bans, oldest first. Check
	// makes each netblock it finds the most recently seen.
	netblocks *simplelru.LRU[netip.Prefix, *[]Ban]
	// held is how many bans netblocks holds, at most rules.MaxBans.
	held int
	// v4Lengths and v6Lengths are the lengths of the IPv4 and IPv6
	// netblocks that have been banned. Check looks for a ban at each of
	// them, so that a ban read from bans.json refuses every client in its
	// netblock even when it was made with another SWWAF_BAN_SCOPE_V4_PREFIX,
	// or another length of an IPv6 client's netblock.
	v4Lengths, v6Lengths []int
}

// New returns a Ledger with no ban yet.
func New(rules Rules) *Ledger {
	// Every netblock held has a ban, so there are never more netblocks
	// than rules.MaxBans, and the LRU never drops one itself.
	netblocks, err := simplelru.NewLRU[netip.Prefix, *[]Ban](rules.MaxBans, nil)
	if err != nil {
		panic(err) // NewLRU fails only for a size below one
	}

	return &Ledger{
		rules:     rules,
		changed:   make(chan struct{}, 1),
		netblocks: netblocks,
	}
}

// Changed receives a value after a ban is made, so that bans.json can be
// written. Several bans made before it is read leave one value.
func (l *Ledger) Changed() <-chan struct{} {
	return l.changed
}

// Check is called for each request from client, at now. It reports
// whether a ban on a netblock client is in is active, and returns that
// ban, with the request counted among those it refused.
func (l *Ledger) Check(client netip.Addr, now time.Time) (Ban, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	lengths := l.v6Lengths
	if client.Is4() {
		lengths = l.v4Lengths
	}

	for _, length := range lengths {
		bans, found := l.netblocks.Get(netip.PrefixFrom(client, length).Masked())
		if !found {
			continue
		}

		// A ban is made only once the one before has ended, so only the
		// last can be active.
		last := &(*bans)[len(*bans)-1]
		if last.ActiveAt(now) {
			last.Notes.Requests++
			last.Notes.Refused++

			return *last, true
		}
	}

	return Ban{}, false
}

// BanForLimit bans netblock at now for a broken limit, with notes, and
// returns the ban. A first ban lasts LimitBanDuration. A ban made within
// LimitBanRepeatWindow after the netblock's last ban ended lasts
// repeatFactor times as long as that one. A ban that would be longer
// than MaxBanDuration is permanent instead. If a ban on netblock is still
// active, as when two of its requests break a limit at once, that ban is
// returned and no other is made. The ledger fills in the notes' Refused
// and EarlierBans itself.
func (l *Ledger) BanForLimit(netblock netip.Prefix, now time.Time, notes Notes) Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	var last *Ban

	bans, found := l.netblocks.Get(netblock)
	if found {
		last = &(*bans)[len(*bans)-1]
		if last.ActiveAt(now) {
			return *last
		}

		notes.EarlierBans = last.Notes.EarlierBans + 1
	}

	notes.Request = notes.Request.cut()
	ban := Ban{
		Netblock: netblock,
		Start:    now,
		Expires:  l.expiry(last, now),
		Notes:    notes,
	}
	l.add(ban)

	select {
	case l.changed <- struct{}{}:
	default: // a value is waiting already
	}

	return ban
}

// Bans returns the bans held on netblock, oldest first. It is not a
// request from netblock, and leaves when it was last seen unchanged.
func (l *Ledger) Bans(netblock netip.Prefix) []Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	bans, found := l.netblocks.Peek(netblock)
	if !found {
		return nil
	}

	return slices.Clone(*bans)
}

// Snapshot returns every ban held, sorted by netblock, and each
// netblock's bans oldest first, as bans.json lists them.
func (l *Ledger) Snapshot() []Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	held := make([]Ban, 0, l.held)
	for _, bans := range l.netblocks.Values() {
		held = append(held, *bans...)
	}

	slices.SortStableFunc(held, func(a, b Ban) int {
		return a.Netblock.Compare(b.Netblock)
	})

	return held
}

// Load puts bans read from bans.json into the ledger, in place of the
// bans it holds, in the order they started, so that a netblock whose last
// ban started latest counts as the most recently seen. Each netblock is
// masked to its length, so that 203.0.113.9/24 is 203.0.113.0/24, and
// each text in the notes is cut to 256 bytes. Past MaxBans the earliest
// bans are dropped, as when they are made.
func (l *Ledger) Load(bans []Ban) {
	bans = slices.Clone(bans)
	slices.SortStableFunc(bans, func(a, b Ban) int {
		return a.Start.Compare(b.Start)
	})

	l.mu.Lock()
	defer l.mu.Unlock()

	l.netblocks.Purge()
	l.held = 0
	l.v4Lengths, l.v6Lengths = nil, nil

	for _, ban := range bans {
		ban.Netblock = ban.Netblock.Masked()
		ban.Notes.Request = ban.Notes.Request.cut()
		l.add(ban)
	}
}

// add adds ban to its netblock's bans, after the last, and makes its
// netblock the most recently seen. With MaxBans held, it drops one first.
func (l *Ledger) add(ban Ban) {
	if l.held == l.rules.MaxBans {
		l.dropOne()
	}

	// dropOne can have dropped the netblock's last ban, and the netblock
	// with it.
	bans, found := l.netblocks.Get(ban.Netblock)
	if !found {
		bans = &[]Ban{}
		l.netblocks.Add(ban.Netblock, bans)
	}

	*bans = append(*bans, ban)
	l.held++

	lengths := &l.v6Lengths
	if ban.Netblock.Addr().Is4() {
		lengths = &l.v4Lengths
	}

	if !slices.Contains(*lengths, ban.Netblock.Bits()) {
		*lengths = append(*lengths, ban.Netblock.Bits())
	}
}

// expiry returns when a ban for a broken limit made at now ends, or zero
// when it is permanent. last is the netblock's last ban, which has ended,
// or nil when it has none.
func (l *Ledger) expiry(last *Ban, now time.Time) time.Time {
	length := l.rules.LimitBanDuration

	if last != nil && now.Sub(last.Expires) <= l.rules.LimitBanRepeatWindow {
		lastLength := last.Expires.Sub(last.Start)
		// This is repeatFactor * lastLength > MaxBanDuration, written so
		// that it cannot overflow.
		if lastLength > l.rules.MaxBanDuration/repeatFactor {
			return time.Time{}
		}

		length = repeatFactor * lastLength
	}

	if length > l.rules.MaxBanDuration {
		return time.Time{}
	}

	return now.Add(length)
}

// dropOne drops the earliest ban of the netblock that has gone longest
// without a request, and the netblock with it if that was its only ban.
func (l *Ledger) dropOne() {
	netblock, bans, _ := l.netblocks.GetOldest()
	if len(*bans) == 1 {
		l.netblocks.Remove(netblock)
	} else {
		*bans = slices.Delete(*bans, 0, 1)
	}

	l.held--
}

// cut returns r with each text cut to maxTextBytes and copied, so that
// the notes do not keep the rest of the request in memory.
func (r Request) cut() Request {
	r.Method = cutText(r.Method)
	r.Host = cutText(r.Host)
	r.Path = cutText(r.Path)
	r.UserAgent = cutText(r.UserAgent)

	return r
}

// cutText returns a copy of the first maxTextBytes of text.
func cutText(text string) string {
	return strings.Clone(text[:min(len(text), maxTextBytes)])
}