smallwebwaf/internal/ratelimit/ratelimit.go

// Package ratelimit keeps the table of clients: each client's requests
// counted over a minute, an hour and a day, as the "Counting method"
// section of SPEC.md describes, which tell when a request takes the client
// over a rate limit, and each client's history since it was first seen.
// At most 20,000 clients are kept, in memory, and written to clients.json
// and read from it by the state package.
package ratelimit

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

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

// maxClients is how many clients are kept. Past it, the least recently
// seen client is dropped, with its history, and starts afresh if it comes
// back.
const maxClients = 20000

const day = 24 * time.Hour

// Limits are the most requests a client may make in a minute, an hour and
// a day. Zero is no limit.
type Limits struct {
	PerMinute int64
	PerHour   int64
	PerDay    int64
}

// Limiter counts each client's requests against the limits, and keeps
// its history. It is safe for concurrent use.
type Limiter struct {
	// windows are the minute, the hour and the day, in the order of
	// Client.buckets.
	windows [3]window

	mu      sync.Mutex
	clients *simplelru.LRU[netip.Prefix, *Client]
}

// Client is a client in the table, as clients.json holds it: its buckets
// in each window, and its history.
type Client struct {
	Client  netip.Prefix `json:"client"`
	Minute  Buckets      `json:"minute"`
	Hour    Buckets      `json:"hour"`
	Day     Buckets      `json:"day"`
	History History      `json:"history"`
}

// Buckets are a client's two buckets in one window: the requests in the
// bucket under way, which began at Start, and in the bucket before it.
type Buckets struct {
	Start    time.Time `json:"start"`
	Current  int64     `json:"current"`
	Previous int64     `json:"previous"`
}

// History is what is known of a client since it was first seen.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type History struct {
	FirstSeen time.Time `json:"first_seen"`
	LastSeen  time.Time `json:"last_seen"`
	// Country is the client's country as it was last looked up, and
	// LookedUp when that was; both are empty while it never was.
	Country  string    `json:"country,omitempty"`
	LookedUp time.Time `json:"looked_up,omitzero"`
	// Requests are all the client's requests: Forwarded those passed to
	// the app, Refused those refused before anything reached it.
	Requests  int64 `json:"requests"`
	Forwarded int64 `json:"forwarded"`
	Refused   int64 `json:"refused"`
	// RequestBytes and ResponseBytes are the body bytes of its requests
	// and of the responses it was sent.
	RequestBytes  int64     `json:"request_bytes"`
	ResponseBytes int64     `json:"response_bytes"`
	Responses     Responses `json:"responses,omitzero"`
	Offences      Offences  `json:"offences,omitzero"`
}

// Responses are the responses a client was sent, by status class;
// Status5xx counts every status from 500 up.
type Responses struct {
	Status1xx int64 `json:"1xx,omitempty"`
	Status2xx int64 `json:"2xx,omitempty"`
	Status3xx int64 `json:"3xx,omitempty"`
	Status4xx int64 `json:"4xx,omitempty"`
	Status5xx int64 `json:"5xx,omitempty"`
}

// Offences are a client's offences, by kind.
type Offences struct {
	// Limit is its requests that broke a rate limit.
	Limit int64 `json:"limit"`
}

// Request is what a client's history keeps of one of its requests.
type Request struct {
	// Country is the client's country, when the request looked it up.
	Country string
	// Forwarded is true for a request passed to the app, false for one
	// refused before anything reached it.
	Forwarded bool
	// Status is what the client was sent, 0 if nothing was.
	Status int
	// RequestBytes and ResponseBytes are the body bytes of the request
	// and of its response.
	RequestBytes  int64
	ResponseBytes int64
	// BrokeLimit is true for a request that broke a rate limit.
	BrokeLimit bool
}

// New returns a Limiter for limits, with no client counted yet.
func New(limits Limits) *Limiter {
	clients, err := simplelru.NewLRU[netip.Prefix, *Client](maxClients, nil)
	if err != nil {
		panic(err) // NewLRU fails only for a size below one
	}

	return &Limiter{
		windows: [3]window{
			{name: "minute", length: time.Minute, limit: limits.PerMinute},
			{name: "hour", length: time.Hour, limit: limits.PerHour},
			{name: "day", length: day, limit: limits.PerDay},
		},
		clients: clients,
	}
}

// Hit is a request that takes a client over a rate limit.
type Hit struct {
	// Window is "minute", "hour" or "day".
	Window string
	// Limit is the window's limit.
	Limit int64
	// Requests is the client's requests counted in the window, this one
	// included.
	Requests float64
}

// Count counts a request from client at now, in every window, whether or
// not it is refused. It reports whether the request takes the client over
// a limit, and the window whose limit it goes over, the shortest if it is
// over several.
func (l *Limiter) Count(client netip.Prefix, now time.Time) (Hit, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	var hit Hit

	for i, b := range l.get(client).buckets() {
		w := l.windows[i]

		requests := b.add(now, w.length)
		if hit.Window == "" && w.limit > 0 && requests > float64(w.limit) {
			hit = Hit{Window: w.name, Limit: w.limit, Requests: requests}
		}
	}

	return hit, hit.Window != ""
}

// Reset sets client's counts in every window back to zero. Its history
// keeps its totals.
func (l *Limiter) Reset(client netip.Prefix) {
	l.mu.Lock()
	defer l.mu.Unlock()

	c, seen := l.clients.Peek(client)
	if seen {
		c.Minute, c.Hour, c.Day = Buckets{}, Buckets{}, Buckets{}
	}
}

// AddToHistory adds r, a request from client at now, to the client's
// history.
func (l *Limiter) AddToHistory(client netip.Prefix, now time.Time, r Request) {
	l.mu.Lock()
	defer l.mu.Unlock()

	h := &l.get(client).History
	if h.FirstSeen.IsZero() {
		h.FirstSeen = now
	}

	h.LastSeen = now

	if r.Country != "" {
		h.Country = r.Country
		h.LookedUp = now
	}

	h.Requests++
	if r.Forwarded {
		h.Forwarded++
	} else {
		h.Refused++
	}

	h.RequestBytes += r.RequestBytes
	h.ResponseBytes += r.ResponseBytes
	h.Responses.add(r.Status)

	if r.BrokeLimit {
		h.Offences.Limit++
	}
}

// Requests returns how many requests the clients inside netblock have
// sent, as their histories count them.
func (l *Limiter) Requests(netblock netip.Prefix) int64 {
	l.mu.Lock()
	defer l.mu.Unlock()

	// Most often the netblock is one client.
	c, seen := l.clients.Peek(netblock)
	if seen {
		return c.History.Requests
	}

	var requests int64

	for _, c := range l.clients.Values() {
		if netblock.Overlaps(c.Client) {
			requests += c.History.Requests
		}
	}

	return requests
}

// Snapshot returns every client in the table, sorted by address, as
// clients.json lists them.
func (l *Limiter) Snapshot() []Client {
	l.mu.Lock()

	clients := make([]Client, 0, l.clients.Len())
	for _, c := range l.clients.Values() {
		clients = append(clients, *c)
	}

	l.mu.Unlock()

	slices.SortFunc(clients, func(a, b Client) int {
		return a.Client.Compare(b.Client)
	})

	return clients
}

// Load puts clients read from clients.json into a table that holds none
// yet, in the order they were last seen, so that the least recently seen
// is dropped first. Buckets whose time has passed at now are emptied.
func (l *Limiter) Load(clients []Client, now time.Time) {
	l.mu.Lock()
	defer l.mu.Unlock()

	clients = slices.Clone(clients)
	slices.SortStableFunc(clients, func(a, b Client) int {
		return a.History.LastSeen.Compare(b.History.LastSeen)
	})

	for _, c := range clients {
		for i, b := range c.buckets() {
			// The window that ends at now covers neither bucket once it
			// begins after the bucket under way has ended.
			length := l.windows[i].length
			if !now.Add(-length).Before(b.Start.Add(length)) {
				*b = Buckets{}
			}
		}

		l.clients.Add(c.Client, &c)
	}
}

// get returns client's entry in the table, a new one if it has none, and
// makes it the most recently seen.
func (l *Limiter) get(client netip.Prefix) *Client {
	c, seen := l.clients.Get(client)
	if !seen {
		c = &Client{Client: client}
		l.clients.Add(client, c)
	}

	return c
}

// buckets returns c's buckets in the minute, the hour and the day.
func (c *Client) buckets() [3]*Buckets {
	return [3]*Buckets{&c.Minute, &c.Hour, &c.Day}
}

// window is a length of time over which requests are counted, and the
// most requests a client may make in it.
type window struct {
	name   string
	length time.Duration
	limit  int64
}

// add counts a request at now in a window of length, and returns the
// client's requests in the window that ends at now: those in the bucket
// under way, and those in the bucket before it weighted by how much of
// that bucket the window still covers.
//
// Concurrent requests can be counted out of order, so now can be a moment
// before the bucket under way began; such a request is counted in that
// bucket. A request dated more than a second before it means the clock
// was set back, and the buckets start afresh: otherwise the bucket before
// would keep its full weight until the clock caught up.
func (b *Buckets) add(now time.Time, length time.Duration) float64 {
	if now.Before(b.Start.Add(-time.Second)) {
		*b = Buckets{}
	}

	start := now.Truncate(length)
	if start.After(b.Start) {
		if start.Equal(b.Start.Add(length)) {
			b.Previous = b.Current
		} else {
			b.Previous = 0
		}

		b.Start = start
		b.Current = 0
	}

	b.Current++

	elapsed := max(now.Sub(b.Start), 0)
	covered := 1 - float64(elapsed)/float64(length)

	return float64(b.Previous)*covered + float64(b.Current)
}

// add counts a response with status in its class. A status of 0, for
// nothing sent, is not a response.
func (r *Responses) add(status int) {
	switch {
	case status >= http.StatusInternalServerError:
		r.Status5xx++
	case status >= http.StatusBadRequest:
		r.Status4xx++
	case status >= http.StatusMultipleChoices:
		r.Status3xx++
	case status >= http.StatusOK:
		r.Status2xx++
	case status >= http.StatusContinue:
		r.Status1xx++
	}
}