fix(security): prune and cap the rate limiter state
REVIEW-001 (M0) M3s: the fixed-window limiter's state map was never pruned, so a flood of distinct keys grew it without bound (memory DoS). Sweep elapsed windows on insert — gated on a size threshold and at most once per window, so it stays amortised — and hard-cap the map (65536 entries) as the backstop. Expose Len() for tests/observability and cover both the sweep and the cap.
This commit is contained in:
parent
6f2ea67503
commit
f7d232950f
2 changed files with 98 additions and 5 deletions
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
|
|
@ -77,6 +78,51 @@ func TestRateLimiterWindowResets(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// TestRateLimiterPrunesElapsedWindows covers M3s: the state map must not grow
|
||||
// without bound. Once enough distinct keys have been seen, an elapsed window
|
||||
// makes the next insert sweep them away.
|
||||
func TestRateLimiterPrunesElapsedWindows(t *testing.T) {
|
||||
l := middleware.NewRateLimiter(1, time.Minute)
|
||||
now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
l.SetClock(func() time.Time { return now })
|
||||
|
||||
// Seed enough distinct keys to make a sweep worthwhile.
|
||||
const seeded = 5000
|
||||
for i := 0; i < seeded; i++ {
|
||||
l.Allow("k" + strconv.Itoa(i))
|
||||
}
|
||||
if got := l.Len(); got != seeded {
|
||||
t.Fatalf("Len = %d, want %d live windows", got, seeded)
|
||||
}
|
||||
|
||||
// Two windows later every seeded window has elapsed; the next new key
|
||||
// triggers the sweep.
|
||||
now = now.Add(2 * time.Minute)
|
||||
l.Allow("fresh")
|
||||
|
||||
if got := l.Len(); got > 10 {
|
||||
t.Errorf("Len after the sweeps = %d, want the %d elapsed windows pruned (M3s)", got, seeded)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiterMapIsCapped covers M3s's hard cap: a flood of distinct keys
|
||||
// must not grow the map without bound.
|
||||
func TestRateLimiterMapIsCapped(t *testing.T) {
|
||||
l := middleware.NewRateLimiter(1, time.Minute)
|
||||
// The clock is frozen, so the time-gated sweep cannot fire — only the hard
|
||||
// cap bounds the map.
|
||||
now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
l.SetClock(func() time.Time { return now })
|
||||
|
||||
const inserted = 70_000
|
||||
for i := 0; i < inserted; i++ {
|
||||
l.Allow("k" + strconv.Itoa(i))
|
||||
}
|
||||
if got := l.Len(); got >= inserted {
|
||||
t.Errorf("the limiter held %d keys after %d inserts — the map is unbounded (M3s)", got, inserted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiterMiddleware429(t *testing.T) {
|
||||
l := middleware.NewRateLimiter(1, time.Minute)
|
||||
h := l.Middleware(func(*http.Request) string { return "k" },
|
||||
|
|
|
|||
|
|
@ -13,12 +13,17 @@ import (
|
|||
// guards the auth-sensitive routes: /{GUID}, /access, /setup. State is
|
||||
// in-memory and single-process, which matches the single-instance model
|
||||
// (D-012).
|
||||
//
|
||||
// The state map is pruned: elapsed windows are swept as new keys arrive and the
|
||||
// map has a hard cap, so a flood of distinct keys cannot grow it without bound
|
||||
// (M3s — a memory-DoS defence).
|
||||
type RateLimiter struct {
|
||||
mu sync.Mutex
|
||||
limit int
|
||||
window time.Duration
|
||||
now func() time.Time
|
||||
state map[string]*window
|
||||
mu sync.Mutex
|
||||
limit int
|
||||
window time.Duration
|
||||
now func() time.Time
|
||||
state map[string]*window
|
||||
sweptAt time.Time // when the last sweep ran
|
||||
}
|
||||
|
||||
type window struct {
|
||||
|
|
@ -33,6 +38,15 @@ const (
|
|||
DefaultRateWindow = time.Minute
|
||||
)
|
||||
|
||||
// The map's bounds (M3s). sweepThreshold is the size above which an elapsed
|
||||
// window triggers a sweep — below it the map is too small to be worth it — and
|
||||
// sweeps run at most once per window. maxStateEntries is the hard cap, so the
|
||||
// map can never grow without bound whatever the arrival pattern.
|
||||
const (
|
||||
sweepThreshold = 1 << 12
|
||||
maxStateEntries = 1 << 16
|
||||
)
|
||||
|
||||
// NewRateLimiter builds a limiter allowing limit requests per window per key.
|
||||
// Non-positive values fall back to the defaults.
|
||||
func NewRateLimiter(limit int, win time.Duration) *RateLimiter {
|
||||
|
|
@ -58,6 +72,9 @@ func (l *RateLimiter) Allow(key string) bool {
|
|||
now := l.now()
|
||||
w, ok := l.state[key]
|
||||
if !ok || !now.Before(w.reset) {
|
||||
if !ok {
|
||||
l.pruneLocked(now)
|
||||
}
|
||||
l.state[key] = &window{count: 1, reset: now.Add(l.window)}
|
||||
return true
|
||||
}
|
||||
|
|
@ -65,6 +82,36 @@ func (l *RateLimiter) Allow(key string) bool {
|
|||
return w.count <= l.limit
|
||||
}
|
||||
|
||||
// pruneLocked drops elapsed windows (M3s). It is called on insert of a NEW
|
||||
// key, and is cheap: a sweep needs the map to be worth sweeping
|
||||
// (sweepThreshold) and runs at most once per window. At the hard cap the map is
|
||||
// reset outright — a bounded limiter that occasionally forgets beats an
|
||||
// unbounded one an attacker can hold open.
|
||||
func (l *RateLimiter) pruneLocked(now time.Time) {
|
||||
if len(l.state) >= maxStateEntries {
|
||||
l.state = make(map[string]*window)
|
||||
l.sweptAt = now
|
||||
return
|
||||
}
|
||||
if len(l.state) < sweepThreshold || now.Sub(l.sweptAt) < l.window {
|
||||
return
|
||||
}
|
||||
for k, w := range l.state {
|
||||
if !now.Before(w.reset) {
|
||||
delete(l.state, k)
|
||||
}
|
||||
}
|
||||
l.sweptAt = now
|
||||
}
|
||||
|
||||
// Len reports how many distinct keys the limiter is tracking. It is for tests
|
||||
// and observability; the map is bounded by maxStateEntries (M3s).
|
||||
func (l *RateLimiter) Len() int {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return len(l.state)
|
||||
}
|
||||
|
||||
// SetClock injects a clock (tests).
|
||||
func (l *RateLimiter) SetClock(fn func() time.Time) {
|
||||
if fn == nil {
|
||||
|
|
|
|||
Loading…
Reference in a new issue