xtablo-source/backend/internal/auth/ratelimit_test.go

package auth

import (
	"sync"
	"testing"
	"time"
)

// fakeNow returns a function that returns a fixed time, advancing by delta on
// each call. For tests that need a static clock, pass delta=0.
func fakeNow(t0 time.Time, delta time.Duration) func() time.Time {
	mu := sync.Mutex{}
	current := t0
	return func() time.Time {
		mu.Lock()
		defer mu.Unlock()
		t := current
		current = current.Add(delta)
		return t
	}
}

// staticNow returns a function that always returns the same time.
func staticNow(t0 time.Time) func() time.Time {
	return func() time.Time { return t0 }
}

// TestRateLimit_BurstAllowsFiveThenDenies verifies that with rate=5/min and
// burst=5, exactly five consecutive Allow("k") calls return true at the same
// fake timestamp, and the sixth returns false (Pattern 8).
func TestRateLimit_BurstAllowsFiveThenDenies(t *testing.T) {
	t0 := time.Now()
	s := newLimiterStoreWithClock(staticNow(t0))

	const key = "user@example.com:127.0.0.1"
	for i := 1; i <= 5; i++ {
		if !s.Allow(key) {
			t.Fatalf("call %d: expected true (within burst), got false", i)
		}
	}
	if s.Allow(key) {
		t.Fatal("call 6: expected false (burst exhausted), got true")
	}
}

// TestRateLimit_RefillsAfter12s checks that after the burst is exhausted, a
// new token refills after ~12 seconds (rate.Every(12s) = 5/min).
func TestRateLimit_RefillsAfter12s(t *testing.T) {
	t0 := time.Now()
	s := newLimiterStoreWithClock(staticNow(t0))

	const key = "user@example.com:127.0.0.1"
	// Exhaust the burst.
	for i := 0; i < 5; i++ {
		s.Allow(key)
	}
	// Sixth attempt at t0 must fail.
	if s.Allow(key) {
		t.Fatal("expected false immediately after burst exhaustion, got true")
	}

	// Advance clock by 12 seconds — one token should have refilled.
	t1 := t0.Add(12 * time.Second)
	s.now = staticNow(t1)

	if !s.Allow(key) {
		t.Fatal("expected true after 12s refill window, got false")
	}
}

// TestRateLimit_PerKeyIsolation ensures that exhausting one key does not affect
// another key sharing the same IP.
func TestRateLimit_PerKeyIsolation(t *testing.T) {
	t0 := time.Now()
	s := newLimiterStoreWithClock(staticNow(t0))

	keyA := "alice@example.com:10.0.0.1"
	keyB := "bob@example.com:10.0.0.1"

	// Exhaust key A.
	for i := 0; i < 5; i++ {
		s.Allow(keyA)
	}
	if s.Allow(keyA) {
		t.Fatal("keyA: expected false after burst exhaustion, got true")
	}

	// Key B must still have its full burst available.
	if !s.Allow(keyB) {
		t.Fatal("keyB: expected true (isolated limiter), got false")
	}
}

// TestRateLimit_JanitorEvictsIdle verifies that cleanupNow removes entries
// whose lastSeen is older than idleTTL.
func TestRateLimit_JanitorEvictsIdle(t *testing.T) {
	t0 := time.Now()
	s := newLimiterStoreWithClock(staticNow(t0))

	// Insert two entries by calling Allow.
	s.Allow("keyA:1.2.3.4")
	s.Allow("keyB:1.2.3.4")

	if sz := s.size(); sz != 2 {
		t.Fatalf("before cleanup: size = %d; want 2", sz)
	}

	// Advance clock past idleTTL (10 min default).
	t1 := t0.Add(11 * time.Minute)
	s.now = staticNow(t1)

	s.cleanupNow()

	if sz := s.size(); sz != 0 {
		t.Fatalf("after cleanup: size = %d; want 0", sz)
	}
}

// TestRateLimit_ConcurrentAllowDoesNotPanic runs 100 goroutines calling Allow
// on overlapping keys concurrently. Run with -race to detect data races.
func TestRateLimit_ConcurrentAllowDoesNotPanic(t *testing.T) {
	t0 := time.Now()
	s := newLimiterStoreWithClock(staticNow(t0))

	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func(n int) {
			defer wg.Done()
			// Three distinct keys with overlapping access.
			keys := []string{"k1:ip1", "k2:ip1", "k1:ip2"}
			s.Allow(keys[n%3])
		}(i)
	}
	wg.Wait()
}