Per-client request rate limits over a minute, an hour and a day (closes #43)

Each client, one IPv4 address or one IPv6 /64, is counted in two buckets per window, the earlier weighted by how much of it the window covers; at most 20,000 clients are kept, least recently seen dropped first. A request over SWWAF_RATE_LIMIT_PER_MINUTE, _HOUR or _DAY (1000, 10000, 50000, or off) gets 429 before reaching the app. Refused requests count, 413s included. A clock set back over a second behind a bucket's start restarts that window. The log line gains limit_hit and the action rate_limited. Deviation from SPEC.md, per the issue: the 20,000 bound and /64 are fixed. Judgement call: golang-lru/v2 holds the table; httprate does not count refused requests. Deviation: go.mod and go.sum hand-written; no make target tidies them. Model: opus-5-5
2026-10-04 04:24:34 +02:00
parent bedd324f3c
commit f51459fbfe
15 changed files with 587 additions and 32 deletions
@@ -0,0 +1,132 @@
+// Package ratelimit counts each client's requests over a minute, an hour
+// and a day, as the "Counting method" section of SPEC.md describes, and
+// tells when a request takes a client over a rate limit. The counts are
+// kept in memory only, for at most 20,000 clients.
+package ratelimit
+
+import (
+	"net/netip"
+	"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, 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. It is safe
+// for concurrent use.
+type Limiter struct {
+	windows [3]window
+
+	mu sync.Mutex
+	// clients holds each client's buckets, one pair for each of windows,
+	// in the same order.
+	clients *simplelru.LRU[netip.Prefix, *[3]buckets]
+}
+
+// New returns a Limiter for limits, with no client counted yet.
+func New(limits Limits) *Limiter {
+	clients, err := simplelru.NewLRU[netip.Prefix, *[3]buckets](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,
+	}
+}
+
+// Count counts a request from client at now, in every window, whether or
+// not it is refused. It returns the window whose limit the request takes
+// the client over, "minute", "hour" or "day", the shortest if it is over
+// several, or "" if it is within every limit.
+func (l *Limiter) Count(client netip.Prefix, now time.Time) string {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	counts, seen := l.clients.Get(client)
+	if !seen {
+		counts = &[3]buckets{}
+		l.clients.Add(client, counts)
+	}
+
+	limitHit := ""
+
+	for i, w := range l.windows {
+		requests := counts[i].add(now, w.length)
+		if limitHit == "" && w.limit > 0 && requests > float64(w.limit) {
+			limitHit = w.name
+		}
+	}
+
+	return limitHit
+}
+
+// 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
+}
+
+// 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
+	current  int64
+	previous 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)
+}
@@ -0,0 +1,200 @@
+package ratelimit_test
+
+import (
+	"net/netip"
+	"testing"
+	"time"
+
+	"sneak.berlin/go/smallwebwaf/internal/ratelimit"
+)
+
+// limit is the limit the tests set.
+const limit = 3
+
+// The windows, as Count names them.
+const (
+	minute = "minute"
+	hour   = "hour"
+)
+
+func TestEachWindowRefusesAtItsLimitAndLetsTheClientBack(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		window string
+		limits ratelimit.Limits
+		length time.Duration
+	}{
+		{minute, ratelimit.Limits{PerMinute: limit}, time.Minute},
+		{hour, ratelimit.Limits{PerHour: limit}, time.Hour},
+		{"day", ratelimit.Limits{PerDay: limit}, 24 * time.Hour},
+	} {
+		t.Run(tc.window, func(t *testing.T) {
+			t.Parallel()
+
+			limiter := ratelimit.New(tc.limits)
+			client := netip.MustParsePrefix("203.0.113.9/32")
+			start := midnight()
+			quarter := tc.length / 4
+
+			for range limit {
+				wantCount(t, limiter, client, start, "")
+			}
+
+			wantCount(t, limiter, client, start, tc.window)
+
+			// A quarter into the next bucket, the window still covers three
+			// quarters of the bucket before, with its four requests: 3 + 1
+			// is over the limit.
+			wantCount(t, limiter, client, start.Add(tc.length+quarter), tc.window)
+
+			// Three quarters into it, a quarter: 1 + 2 is within.
+			wantCount(t, limiter, client, start.Add(tc.length+3*quarter), "")
+		})
+	}
+}
+
+func TestClientBackAfterAWholeBucketIsWithinTheLimitAtOnce(t *testing.T) {
+	t.Parallel()
+
+	limiter := ratelimit.New(ratelimit.Limits{PerHour: limit})
+	client := netip.MustParsePrefix("203.0.113.9/32")
+	start := midnight()
+
+	for range limit {
+		wantCount(t, limiter, client, start, "")
+	}
+
+	wantCount(t, limiter, client, start, hour)
+
+	// No request in the whole next bucket, so a quarter into the one after
+	// it the window covers none of the four requests: 1 is within the
+	// limit. Were they counted as the bucket before, 3 + 1 would be over.
+	wantCount(t, limiter, client, start.Add(2*time.Hour+time.Hour/4), "")
+}
+
+func TestRefusedRequestsCount(t *testing.T) {
+	t.Parallel()
+
+	limiter := ratelimit.New(ratelimit.Limits{PerMinute: limit, PerHour: 2 * limit})
+	refused := netip.MustParsePrefix("203.0.113.9/32")
+	within := netip.MustParsePrefix("203.0.113.10/32")
+	start := midnight()
+
+	for range limit {
+		wantCount(t, limiter, refused, start, "")
+		wantCount(t, limiter, within, start, "")
+	}
+
+	for range limit {
+		wantCount(t, limiter, refused, start, minute)
+	}
+
+	// Half a minute into the next bucket the window covers half of the
+	// bucket before: 3 + 1 is over the minute's limit for the client
+	// whose three refused requests count, and 1.5 + 1 within it for the
+	// other. The first is over the hour's limit too, and the shorter
+	// window is named.
+	halfway := start.Add(time.Minute + time.Minute/2)
+	wantCount(t, limiter, refused, halfway, minute)
+	wantCount(t, limiter, within, halfway, "")
+
+	// The refused requests count in the hour as well: 6 + 1 + 1 is over
+	// its limit, and 3 + 1 + 1 within it.
+	later := start.Add(10 * time.Minute)
+	wantCount(t, limiter, refused, later, hour)
+	wantCount(t, limiter, within, later, "")
+}
+
+func TestRequestCountedLateGoesInTheBucketUnderWay(t *testing.T) {
+	t.Parallel()
+
+	limiter := ratelimit.New(ratelimit.Limits{PerMinute: limit})
+	client := netip.MustParsePrefix("203.0.113.9/32")
+	start := midnight()
+
+	for range limit {
+		wantCount(t, limiter, client, start, "")
+	}
+
+	// A concurrent request dated a moment before the bucket under way, but
+	// counted after it began, is counted in it: 3 + 1 is over the limit.
+	wantCount(t, limiter, client, start.Add(-time.Millisecond), minute)
+}
+
+func TestClockSetBackStartsTheBucketsAfresh(t *testing.T) {
+	t.Parallel()
+
+	limiter := ratelimit.New(ratelimit.Limits{PerHour: limit})
+	client := netip.MustParsePrefix("203.0.113.9/32")
+	start := midnight()
+
+	for range limit {
+		wantCount(t, limiter, client, start, "")
+	}
+
+	// Half an hour into the next bucket: 3 / 2 + 1 is within the limit.
+	wantCount(t, limiter, client, start.Add(time.Hour+time.Hour/2), "")
+
+	// The clock is set back an hour. Counted in the bucket under way, the
+	// next request would find the bucket before it at full weight, 3 + 2,
+	// over the limit until the clock caught up. The buckets start afresh
+	// instead, and the client is refused only past the limit again.
+	setBack := start.Add(time.Hour / 2)
+	for range limit {
+		wantCount(t, limiter, client, setBack, "")
+	}
+
+	wantCount(t, limiter, client, setBack, hour)
+}
+
+func TestKeepsAtMost20000ClientsDroppingTheLeastRecentlySeen(t *testing.T) {
+	t.Parallel()
+
+	const maxClients = 20000
+
+	limiter := ratelimit.New(ratelimit.Limits{PerMinute: 1})
+	now := midnight()
+
+	clients := make([]netip.Prefix, maxClients+1)
+	addr := netip.MustParseAddr("10.0.0.0")
+
+	for i := range clients {
+		clients[i] = netip.PrefixFrom(addr, addr.BitLen())
+		addr = addr.Next()
+	}
+
+	for _, client := range clients[:maxClients] {
+		wantCount(t, limiter, client, now, "")
+	}
+
+	// The first client is seen again: its second request is over the
+	// limit of one, so it is still counted.
+	wantCount(t, limiter, clients[0], now, minute)
+
+	// One client more drops the least recently seen, the second, which
+	// starts afresh, while the first is kept.
+	wantCount(t, limiter, clients[maxClients], now, "")
+	wantCount(t, limiter, clients[1], now, "")
+	wantCount(t, limiter, clients[0], now, minute)
+}
+
+// midnight is the start of a bucket in every window.
+func midnight() time.Time {
+	return time.Date(2026, 10, 4, 0, 0, 0, 0, time.UTC)
+}
+
+// wantCount counts a request from client at now, and checks the window
+// whose limit it goes over, "" for none.
+func wantCount(
+	t *testing.T, limiter *ratelimit.Limiter, client netip.Prefix, now time.Time,
+	want string,
+) {
+	t.Helper()
+
+	got := limiter.Count(client, now)
+	if got != want {
+		t.Errorf("request from %s at %s is over %q, want %q",
+			client, now.Format(time.RFC3339), got, want)
+	}
+}