pages/internal/auth/ratelimit_test.go

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