Initial commit

Assisted-by: GLM 5.3

internal/ratelimit/ratelimit.go

@@ -0,0 +1,230 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
+
+// Package ratelimit implements sliding-window request counters keyed by
+// client IP, with LRU-bounded storage and periodic sweeps.
+package ratelimit
+
+import (
+	"container/list"
+	"sync"
+	"time"
+)
+
+const (
+	apiMaxTracked       = 20_000
+	apiCleanupInterval  = 300 * time.Second
+	loginWindow         = 60 * time.Second
+	loginMaxAttempts    = 10
+	loginMaxTracked     = 10_000
+	loginCleanupEvery   = 300 * time.Second
+	attemptsPerKeyLimit = 64
+)
+
+// window is a sliding-window counter over opaque string keys. Both the
+// public API limiter and the login limiter are one of these with
+// different limits, which is why the storage, the sweep and the
+// eviction live here once.
+type window struct {
+	limit        int
+	duration     time.Duration
+	maxTracked   int
+	cleanupEvery time.Duration
+	now          func() time.Time
+
+	mu          sync.Mutex
+	hits        map[string][]time.Time
+	order       *list.List // front = most recently used key
+	positions   map[string]*list.Element
+	lastCleanup time.Time
+	haveCleaned bool
+}
+
+func newWindow(limit int, duration time.Duration, maxTracked int, cleanupEvery time.Duration, now func() time.Time) *window {
+	if now == nil {
+		now = time.Now
+	}
+	return &window{
+		limit:        limit,
+		duration:     duration,
+		maxTracked:   maxTracked,
+		cleanupEvery: cleanupEvery,
+		now:          now,
+		hits:         map[string][]time.Time{},
+		order:        list.New(),
+		positions:    map[string]*list.Element{},
+	}
+}
+
+// hit records one attempt for key and returns the number of attempts in
+// the window, including this one.
+func (w *window) hit(key string) int {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	now := w.now()
+	w.sweepLocked(now)
+	kept := w.pruneLocked(key, now)
+	kept = append(kept, now)
+	w.hits[key] = kept
+	w.touchLocked(key)
+	w.evictLocked()
+	return len(kept)
+}
+
+// count returns the number of attempts for key inside the window,
+// without recording one.
+func (w *window) count(key string) int {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return len(w.pruneLocked(key, w.now()))
+}
+
+// retryAfter returns how long the caller should wait before the oldest
+// tracked attempt for key leaves the window.
+func (w *window) retryAfter(key string) time.Duration {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	now := w.now()
+	kept := w.pruneLocked(key, now)
+	if len(kept) == 0 {
+		return 0
+	}
+	if wait := kept[0].Add(w.duration).Sub(now); wait > 0 {
+		return wait
+	}
+	return 0
+}
+
+// pruneLocked drops the attempts for key that have left the window and
+// stores what is left. It also caps how many attempts one key may keep:
+// a key that is refused on every request would otherwise grow a slice of
+// timestamps for the whole window, unbounded by anything. The cap is far
+// above the limit, so the count a caller sees is unaffected; the retry
+// delay of a hammering client becomes approximate, which is all it can
+// be once it is already being refused. The caller must hold w.mu.
+func (w *window) pruneLocked(key string, now time.Time) []time.Time {
+	cutoff := now.Add(-w.duration)
+	hits := w.hits[key]
+	kept := hits[:0]
+	for _, t := range hits {
+		if t.After(cutoff) {
+			kept = append(kept, t)
+		}
+	}
+	if len(kept) == 0 {
+		delete(w.hits, key)
+		return kept
+	}
+	if len(kept) > w.limit+attemptsPerKeyLimit {
+		kept = kept[len(kept)-w.limit-attemptsPerKeyLimit:]
+	}
+	w.hits[key] = kept
+	return kept
+}
+
+func (w *window) touchLocked(key string) {
+	if elem, ok := w.positions[key]; ok {
+		w.order.MoveToFront(elem)
+		return
+	}
+	w.positions[key] = w.order.PushFront(key)
+}
+
+// evictLocked drops the least recently used keys once the bound is
+// exceeded. It runs on every insert, so the maps cannot grow between
+// sweeps. The caller must hold w.mu.
+func (w *window) evictLocked() {
+	for w.order.Len() > w.maxTracked {
+		oldest := w.order.Back()
+		if oldest == nil {
+			return
+		}
+		key := oldest.Value.(string)
+		w.order.Remove(oldest)
+		delete(w.positions, key)
+		delete(w.hits, key)
+	}
+}
+
+// sweepLocked drops every attempt that has left its window and trims
+// the key set. It is rate-limited by cleanupEvery so a busy server does
+// not walk the map on every request. The caller must hold w.mu.
+func (w *window) sweepLocked(now time.Time) {
+	if w.haveCleaned && now.Sub(w.lastCleanup) < w.cleanupEvery {
+		return
+	}
+	cutoff := now.Add(-w.duration)
+	for key, hits := range w.hits {
+		kept := hits[:0]
+		for _, t := range hits {
+			if t.After(cutoff) {
+				kept = append(kept, t)
+			}
+		}
+		if len(kept) == 0 {
+			delete(w.hits, key)
+			if elem, ok := w.positions[key]; ok {
+				w.order.Remove(elem)
+				delete(w.positions, key)
+			}
+			continue
+		}
+		w.hits[key] = kept
+	}
+	w.evictLocked()
+	w.lastCleanup = now
+	w.haveCleaned = true
+}
+
+// Limiter is a thread-safe sliding-window request counter.
+type Limiter struct {
+	window *window
+}
+
+// New creates a limiter allowing limit hits per window.
+func New(limit int, window time.Duration) *Limiter {
+	return &Limiter{window: newWindow(limit, window, apiMaxTracked, apiCleanupInterval, time.Now)}
+}
+
+// Limit returns the configured request budget per window.
+func (l *Limiter) Limit() int { return l.window.limit }
+
+// Check records a hit for key and reports whether it is allowed, how
+// many requests remain, and the retry-after delay in seconds.
+func (l *Limiter) Check(key string) (allowed bool, remaining, retryAfter int) {
+	count := l.window.hit(key)
+	allowed = count <= l.window.limit
+	remaining = max(0, l.window.limit-count)
+	if !allowed {
+		retryAfter = int(l.window.retryAfter(key).Seconds()) + 1
+	}
+	return allowed, remaining, retryAfter
+}
+
+// LoginLimiter counts login attempts per IP. Record returns the number
+// of attempts in the current window; Blocked reports whether the
+// threshold was exceeded together with the retry-after delay.
+type LoginLimiter struct {
+	window *window
+}
+
+// NewLoginLimiter creates the shared login-attempt tracker.
+func NewLoginLimiter() *LoginLimiter {
+	return &LoginLimiter{window: newWindow(loginMaxAttempts, loginWindow, loginMaxTracked, loginCleanupEvery, time.Now)}
+}
+
+// Record logs an attempt for ip and returns the count in the window.
+func (l *LoginLimiter) Record(ip string) int { return l.window.hit(ip) }
+
+// Blocked reports whether ip exceeded the login budget and, when it
+// did, the whole-second wait before the oldest attempt expires.
+func (l *LoginLimiter) Blocked(ip string) (bool, int) {
+	if l.window.count(ip) <= l.window.limit {
+		return false, 0
+	}
+	retryAfter := l.window.retryAfter(ip)
+	if retryAfter <= 0 {
+		retryAfter = l.window.duration
+	}
+	return true, int(retryAfter.Seconds()) + 1
+}

internal/ratelimit/ratelimit_test.go

@@ -0,0 +1,167 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
+
+package ratelimit
+
+import (
+	"fmt"
+	"testing"
+	"testing/synctest"
+	"time"
+)
+
+func TestLimiterAllowsWithinBudget(t *testing.T) {
+	l := New(3, time.Minute)
+	for i := range 3 {
+		allowed, remaining, retry := l.Check("ip")
+		if !allowed || remaining != 2-i || retry != 0 {
+			t.Fatalf("hit %d: allowed=%v remaining=%d retry=%d", i, allowed, remaining, retry)
+		}
+	}
+}
+
+func TestLimiterBlocksOverBudget(t *testing.T) {
+	l := New(2, time.Minute)
+	l.Check("ip")
+	l.Check("ip")
+	allowed, remaining, retry := l.Check("ip")
+	if allowed {
+		t.Fatal("third hit allowed")
+	}
+	if remaining != 0 {
+		t.Fatalf("remaining = %d", remaining)
+	}
+	if retry < 1 || retry > 61 {
+		t.Fatalf("retry_after = %d", retry)
+	}
+}
+
+// synctest runs the window in virtual time, so the boundary is exact and
+// the test does not spend a millisecond of the wall clock on it.
+func TestLimiterWindowExpiry(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		l := New(1, 20*time.Millisecond)
+		if allowed, _, _ := l.Check("ip"); !allowed {
+			t.Fatal("first hit blocked")
+		}
+		if allowed, _, _ := l.Check("ip"); allowed {
+			t.Fatal("second hit allowed")
+		}
+		time.Sleep(19 * time.Millisecond)
+		if allowed, _, _ := l.Check("ip"); allowed {
+			t.Fatal("a hit inside the window was allowed")
+		}
+		// Past the window measured from the last recorded hit.
+		time.Sleep(21 * time.Millisecond)
+		if allowed, remaining, _ := l.Check("ip"); !allowed || remaining != 0 {
+			t.Fatalf("hit after window: allowed=%v remaining=%d", allowed, remaining)
+		}
+	})
+}
+
+func TestLimiterTracksKeysSeparately(t *testing.T) {
+	l := New(1, time.Minute)
+	l.Check("a")
+	if allowed, _, _ := l.Check("b"); !allowed {
+		t.Fatal("separate key blocked")
+	}
+}
+
+// A limiter built for one budget does not carry another's hits: the
+// process-level limiters are constructed once, and a fresh one is clean.
+func TestLimiterStartsClean(t *testing.T) {
+	l := New(1, time.Minute)
+	if allowed, _, _ := l.Check("ip"); !allowed {
+		t.Fatal("the first hit was refused")
+	}
+	if allowed, _, _ := l.Check("ip"); allowed {
+		t.Fatal("the second hit was allowed over a budget of one")
+	}
+	fresh := New(10, 20*time.Millisecond)
+	if allowed, _, _ := fresh.Check("ip"); !allowed {
+		t.Fatal("a fresh limiter started with hits recorded")
+	}
+}
+
+func TestLoginLimiterBlocksAfterMaxAttempts(t *testing.T) {
+	l := NewLoginLimiter()
+	for i := 1; i <= loginMaxAttempts; i++ {
+		if count := l.Record("ip"); count != i {
+			t.Fatalf("count = %d, want %d", count, i)
+		}
+		if blocked, _ := l.Blocked("ip"); blocked {
+			t.Fatalf("blocked at attempt %d", i)
+		}
+	}
+	if count := l.Record("ip"); count != loginMaxAttempts+1 {
+		t.Fatalf("count = %d", count)
+	}
+	blocked, retry := l.Blocked("ip")
+	if !blocked {
+		t.Fatal("not blocked after exceeding the budget")
+	}
+	if retry < 1 || retry > 61 {
+		t.Fatalf("retry_after = %d", retry)
+	}
+}
+
+func TestLoginLimiterSeparateIPs(t *testing.T) {
+	l := NewLoginLimiter()
+	for i := 0; i <= loginMaxAttempts; i++ {
+		l.Record("a")
+	}
+	if blocked, _ := l.Blocked("b"); blocked {
+		t.Fatal("unrelated IP blocked")
+	}
+}
+
+// A fresh limiter starts with no attempts, so a restart of the process
+// forgives an address that was being refused.
+func TestLoginLimiterStartsClean(t *testing.T) {
+	l := NewLoginLimiter()
+	for range loginMaxAttempts + 1 {
+		l.Record("ip")
+	}
+	if blocked, _ := l.Blocked("ip"); !blocked {
+		t.Fatal("the limiter did not refuse after the budget was spent")
+	}
+	fresh := NewLoginLimiter()
+	if blocked, _ := fresh.Blocked("ip"); blocked {
+		t.Fatal("a fresh limiter started with attempts recorded")
+	}
+}
+
+// The sweep is what bounds the maps, so it is driven directly with an
+// injected clock rather than by waiting for the real one.
+func TestWindowSweepBoundsTheKeySet(t *testing.T) {
+	now := time.Unix(1_700_000_000, 0)
+	w := newWindow(2, time.Minute, 10, time.Second, func() time.Time { return now })
+	for i := range 50 {
+		w.hit(fmt.Sprintf("key-%d", i))
+		now = now.Add(time.Millisecond)
+	}
+	if got := w.order.Len(); got > 10 {
+		t.Fatalf("tracked %d keys, want at most 10", got)
+	}
+	// Once the window has passed, a sweep clears what is left.
+	now = now.Add(2 * time.Minute)
+	w.hit("fresh")
+	w.sweepLocked(now)
+	if got := w.order.Len(); got != 1 {
+		t.Fatalf("tracked %d keys after the sweep, want 1", got)
+	}
+}
+
+// A key that keeps being refused must not grow its slice without bound.
+func TestWindowCapsAHammeredKey(t *testing.T) {
+	now := time.Unix(1_700_000_000, 0)
+	w := newWindow(2, time.Hour, 100, time.Minute, func() time.Time { return now })
+	for range 5000 {
+		w.hit("hammer")
+		now = now.Add(time.Millisecond)
+	}
+	w.sweepLocked(now)
+	if got := len(w.hits["hammer"]); got > 2+attemptsPerKeyLimit+1 {
+		t.Fatalf("a hammered key kept %d hits", got)
+	}
+}