init

internal/auth/ratelimit_test.go

@@ -0,0 +1,176 @@
+package auth
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+)
+
+// fakeClock drives the limiter's refill without sleeping.
+type fakeClock struct {
+	mu sync.Mutex
+	t  time.Time
+}
+
+func (c *fakeClock) now() time.Time {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.t
+}
+
+func (c *fakeClock) advance(d time.Duration) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.t = c.t.Add(d)
+}
+
+func testLimiter(t *testing.T, burst int, period time.Duration, max int) (*Limiter, *fakeClock) {
+	t.Helper()
+	l := NewLimiter(burst, period, max)
+	clk := &fakeClock{t: time.Unix(1_700_000_000, 0)}
+	l.now = clk.now
+	return l, clk
+}
+
+// Successful requests must cost nothing: a CI runner pushing hundreds of valid
+// deploys a second is not the thing this limiter is defending against.
+func TestAllowDoesNotConsume(t *testing.T) {
+	l, _ := testLimiter(t, 3, time.Minute, 100)
+	for i := 0; i < 1000; i++ {
+		if !l.Allow("10.0.0.1") {
+			t.Fatalf("Allow denied a caller that never failed (i=%d)", i)
+		}
+	}
+}
+
+func TestFailExhaustsBudget(t *testing.T) {
+	l, _ := testLimiter(t, 3, time.Minute, 100)
+	const ip = "10.0.0.1"
+	for i := 0; i < 3; i++ {
+		if !l.Allow(ip) {
+			t.Fatalf("denied before the burst was spent (i=%d)", i)
+		}
+		l.Fail(ip)
+	}
+	if l.Allow(ip) {
+		t.Error("burst exhausted but the caller is still allowed")
+	}
+	// Failing while already throttled must not push the balance negative, or
+	// the client would take proportionally longer to recover the more it tried.
+	for i := 0; i < 100; i++ {
+		l.Fail(ip)
+	}
+	if d := l.RetryAfter(ip); d > 2*time.Minute {
+		t.Errorf("RetryAfter = %v; over-failing drove the bucket negative", d)
+	}
+}
+
+func TestBudgetRefills(t *testing.T) {
+	l, clk := testLimiter(t, 4, time.Minute, 100)
+	const ip = "10.0.0.1"
+	for i := 0; i < 4; i++ {
+		l.Fail(ip)
+	}
+	if l.Allow(ip) {
+		t.Fatal("expected to be throttled")
+	}
+
+	// One quarter of the period restores one of four tokens.
+	clk.advance(15 * time.Second)
+	if !l.Allow(ip) {
+		t.Error("no token after a quarter period")
+	}
+
+	clk.advance(time.Hour)
+	for i := 0; i < 4; i++ {
+		if !l.Allow(ip) {
+			t.Fatalf("bucket did not refill to full (i=%d)", i)
+		}
+		l.Fail(ip)
+	}
+	if l.Allow(ip) {
+		t.Error("bucket refilled past its burst")
+	}
+}
+
+func TestClientsAreIndependent(t *testing.T) {
+	l, _ := testLimiter(t, 2, time.Minute, 100)
+	for i := 0; i < 2; i++ {
+		l.Fail("10.0.0.1")
+	}
+	if l.Allow("10.0.0.1") {
+		t.Error("attacker not throttled")
+	}
+	if !l.Allow("10.0.0.2") {
+		t.Error("one bad client throttled an unrelated one")
+	}
+}
+
+func TestRetryAfter(t *testing.T) {
+	l, clk := testLimiter(t, 2, time.Minute, 100)
+	const ip = "10.0.0.1"
+	if d := l.RetryAfter(ip); d != 0 {
+		t.Errorf("RetryAfter with budget left = %v, want 0", d)
+	}
+	l.Fail(ip)
+	l.Fail(ip)
+
+	d := l.RetryAfter(ip)
+	if d <= 0 {
+		t.Fatal("a throttled client must be told to wait a positive time")
+	}
+	if d > 2*time.Minute {
+		t.Errorf("RetryAfter = %v, unreasonably long for a 1m period", d)
+	}
+	// Waiting the advertised time must actually be enough.
+	clk.advance(d)
+	if !l.Allow(ip) {
+		t.Errorf("still throttled after waiting the advertised %v", d)
+	}
+}
+
+// The bucket map is keyed by attacker-chosen input, so its bound is load
+// bearing: without it the rate limiter becomes the memory exhaustion vector.
+func TestBucketMapIsBounded(t *testing.T) {
+	l, _ := testLimiter(t, 2, time.Minute, 64)
+	for i := 0; i < 10000; i++ {
+		l.Fail(fmt.Sprintf("10.%d.%d.%d", i>>16&0xff, i>>8&0xff, i&0xff))
+	}
+	if got := l.buckets.Len(); got > 64 {
+		t.Errorf("tracking %d clients, cap is 64", got)
+	}
+}
+
+// A nil limiter is the "rate limiting disabled" configuration and must not
+// panic in the request path.
+func TestNilLimiterAllowsEverything(t *testing.T) {
+	var l *Limiter
+	if !l.Allow("10.0.0.1") {
+		t.Error("nil limiter denied a request")
+	}
+	l.Fail("10.0.0.1")
+	if d := l.RetryAfter("10.0.0.1"); d != 0 {
+		t.Errorf("RetryAfter = %v, want 0", d)
+	}
+}
+
+func TestLimiterConcurrentUse(t *testing.T) {
+	l := NewLimiter(50, time.Minute, 256)
+	var wg sync.WaitGroup
+	for g := 0; g < 16; g++ {
+		wg.Add(1)
+		go func(g int) {
+			defer wg.Done()
+			for i := 0; i < 200; i++ {
+				ip := fmt.Sprintf("10.0.0.%d", i%8)
+				l.Allow(ip)
+				if i%3 == 0 {
+					l.Fail(ip)
+				}
+				l.RetryAfter(ip)
+			}
+		}(g)
+	}
+	wg.Wait()
+}