Initial commit

Assisted-by: GLM 5.3

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