smallwebwaf/internal/ratelimit/ratelimit.go

// Package ratelimit keeps the table of clients: each client's requests
// and bytes 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 or a byte 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

// The kinds of limits, as the metrics name them.
const (
	// KindRequests is a rate limit, on a client's requests.
	KindRequests = "requests"
	// KindBytes is a byte limit, on a client's bytes.
	KindBytes = "bytes"
)

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

// Limiter counts each client's requests and bytes 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 and Client.byteBuckets.
	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
// of requests and of bytes in each window, and its history.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type Client struct {
	Client      netip.Prefix `json:"client"`
	Minute      Buckets      `json:"minute"`
	Hour        Buckets      `json:"hour"`
	Day         Buckets      `json:"day"`
	MinuteBytes Buckets      `json:"minute_bytes"`
	HourBytes   Buckets      `json:"hour_bytes"`
	DayBytes    Buckets      `json:"day_bytes"`
	History     History      `json:"history"`
}

// Buckets are a client's two buckets in one window: the requests, or the
// bytes, 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"`
	// ASN, ASName and Country are the client's AS number, AS name and
	// country as last looked up, each empty when the lookup could not
	// find it, and LookedUp is when the lookup gave that answer; all are
	// empty while the client never was looked up.
	ASN      string    `json:"asn,omitempty"`
	ASName   string    `json:"as_name,omitempty"`
	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, a 401 at
	// smallwebwaf's own endpoints included, and neither the others
	// smallwebwaf answered there.
	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 or a byte limit.
	Limit int64 `json:"limit"`
}

// Request is what a client's history keeps of one of its requests.
type Request struct {
	// Forwarded is true for a request passed to the app, Refused for one
	// refused before anything reached it, a 401 at smallwebwaf's own
	// endpoints included. Both are false for any other request smallwebwaf
	// answered there.
	Forwarded bool
	Refused   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 or a byte
	// 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, byteLimit: limits.BytesPerMinute,
			},
			{
				name: "hour", length: time.Hour,
				limit: limits.PerHour, byteLimit: limits.BytesPerHour,
			},
			{
				name: "day", length: day,
				limit: limits.PerDay, byteLimit: limits.BytesPerDay,
			},
		},
		clients: clients,
	}
}

// Hit is a request that takes a client over a rate limit, or whose bytes
// take it over a byte limit.
type Hit struct {
	// Kind is KindRequests for a rate limit, KindBytes for a byte limit.
	Kind string
	// Window is "minute", "hour" or "day".
	Window string
	// Limit is the window's limit.
	Limit int64
	// Count is the client's requests, or bytes, counted in the window,
	// this request's included.
	Count float64
}

// Counts are a client's requests and bytes in the minute, the hour and
// the day that end at a request, that request's included.
//
//nolint:tagliatelle // SPEC.md's request log names its fields in snake_case
type Counts struct {
	Minute      float64 `json:"minute"`
	Hour        float64 `json:"hour"`
	Day         float64 `json:"day"`
	MinuteBytes float64 `json:"minute_bytes"`
	HourBytes   float64 `json:"hour_bytes"`
	DayBytes    float64 `json:"day_bytes"`
}

// Count counts a request from client at now, in every window, whether or
// not it is refused, and returns the client's counts in each window. It
// reports whether the request takes the client over a rate limit, and the
// hit: the window whose limit it goes over, the shortest if it is over
// several.
func (l *Limiter) Count(client netip.Prefix, now time.Time) (Counts, Hit, bool) {
	return l.count(client, now, 1, 0)
}

// CountBytes counts bytes, those of a request from client that has ended,
// at now, in every window, and returns the client's counts in each window.
// It reports whether the bytes take the client over a byte limit, and the
// hit, as Count does.
func (l *Limiter) CountBytes(
	client netip.Prefix, now time.Time, bytes int64,
) (Counts, Hit, bool) {
	return l.count(client, now, 0, bytes)
}

// Reset sets client's counts of requests and of bytes 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{}
		c.MinuteBytes, c.HourBytes, c.DayBytes = 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

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

	if r.Refused {
		h.Refused++
	}

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

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

// AddLookup gives client's history its AS number, AS name and country, as
// the lookup gave them at lookedUp, if the table of clients holds the
// client.
// It does not make the client the most recently seen.
func (l *Limiter) AddLookup(
	client netip.Prefix, lookedUp time.Time, asn, asName, country string,
) {
	l.mu.Lock()
	defer l.mu.Unlock()

	c, held := l.clients.Peek(client)
	if !held {
		return
	}

	h := &c.History
	h.ASN, h.ASName, h.Country, h.LookedUp = asn, asName, country, lookedUp
}

// 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
}

// Client returns client as the table holds it, and whether it does. It is
// not a request from client, and leaves when it was last seen unchanged.
func (l *Limiter) Client(client netip.Prefix) (Client, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	c, seen := l.clients.Peek(client)
	if !seen {
		return Client{}, false
	}

	return *c, true
}

// Len returns how many clients are in the table.
func (l *Limiter) Len() int {
	l.mu.Lock()
	defer l.mu.Unlock()

	return l.clients.Len()
}

// 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 the table, in place of
// the clients it holds, 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) {
	clients = slices.Clone(clients)
	slices.SortStableFunc(clients, func(a, b Client) int {
		return a.History.LastSeen.Compare(b.History.LastSeen)
	})

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

	l.clients.Purge()

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

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

// count adds requests and bytes from client at now to its buckets in
// every window, and returns its counts. A limit is broken only by what is
// added to it, so that a request whose bytes are counted after another of
// the client's requests broke a rate limit does not break it too.
func (l *Limiter) count(
	client netip.Prefix, now time.Time, requests, bytes int64,
) (Counts, Hit, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	c := l.get(client)
	requestBuckets, byteBuckets := c.buckets(), c.byteBuckets()

	var (
		requestCounts, byteCounts [3]float64
		hit                       Hit
	)

	for i, w := range l.windows {
		requestCounts[i] = requestBuckets[i].add(now, w.length, requests)
		byteCounts[i] = byteBuckets[i].add(now, w.length, bytes)

		switch {
		case hit.Window != "":
		case requests > 0 && w.limit > 0 && requestCounts[i] > float64(w.limit):
			hit = Hit{
				Kind: KindRequests, Window: w.name, Limit: w.limit, Count: requestCounts[i],
			}
		case bytes > 0 && w.byteLimit > 0 && byteCounts[i] > float64(w.byteLimit):
			hit = Hit{
				Kind: KindBytes, Window: w.name, Limit: w.byteLimit, Count: byteCounts[i],
			}
		}
	}

	counts := Counts{
		Minute: requestCounts[0], Hour: requestCounts[1], Day: requestCounts[2],
		MinuteBytes: byteCounts[0], HourBytes: byteCounts[1], DayBytes: byteCounts[2],
	}

	return counts, hit, hit.Window != ""
}

// 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 of requests in the minute, the hour and the
// day.
func (c *Client) buckets() [3]*Buckets {
	return [3]*Buckets{&c.Minute, &c.Hour, &c.Day}
}

// byteBuckets returns c's buckets of bytes in the minute, the hour and the
// day.
func (c *Client) byteBuckets() [3]*Buckets {
	return [3]*Buckets{&c.MinuteBytes, &c.HourBytes, &c.DayBytes}
}

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

// add counts n requests, or n bytes, at now in a window of length, and
// returns the client's count in the window that ends at now: what is in
// the bucket under way, and what is in the bucket before it weighted by
// how much of that bucket the window still covers. With n zero it counts
// nothing, and returns the count.
//
// 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, n int64) 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 += n

	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++
	}
}