package cooler_test import ( "bytes" "context" "fmt" "log/slog" "net" "strings" "sync" "sync/atomic" "testing" "time" "git.golem15.com/golem15/summercms/modules/backpack" "git.golem15.com/golem15/summercms/modules/compass" "git.golem15.com/golem15/summercms/modules/cooler" "git.golem15.com/golem15/summercms/modules/surf" "github.com/alicebob/miniredis/v2" ) // Both stores must satisfy surf.Store; the compiler checks the Redis one here. var _ surf.Store = (*cooler.RateLimitStore)(nil) // parityStore is one rate-limit store under test with the way its clock moves. type parityStore struct { name string store surf.Store advance func(time.Duration) } // parityStep is one Attempt, after the store's clock moved by advance. type parityStep struct { advance time.Duration key string max int allowed bool attempts int // newWindow: retryAfter is close to the full decay (a window opened). // Otherwise retryAfter never grows within the window and stays at or // below decay minus the time advanced since the window opened. newWindow bool } // parityDecay is long enough for real-time stores to cross it in a test and // short enough to keep the table fast. const parityDecay = 300 * time.Millisecond // parityTable is the shared fixed-window behaviour of surf.MemoryStore and // the Redis store: the first attempt opens the window, a denial does not // increment, the window never extends, retryAfter never grows and expiry // opens a new window. var parityTable = []parityStep{ {key: "login|192.0.2.1", max: 3, allowed: true, attempts: 1, newWindow: true}, {key: "login|192.0.2.1", max: 3, allowed: true, attempts: 2}, {key: "login|192.0.2.1", max: 3, allowed: true, attempts: 3}, {key: "login|192.0.2.1", max: 3, allowed: false, attempts: 3}, {key: "login|192.0.2.1", max: 3, allowed: false, attempts: 3}, // Another key has its own window. {key: "login|192.0.2.2", max: 3, allowed: true, attempts: 1, newWindow: true}, // A max of 0 denies without counting. {key: "closed", max: 0, allowed: false, attempts: 0, newWindow: true}, // Inside the window: still denied, the window did not move. {advance: 120 * time.Millisecond, key: "login|192.0.2.1", max: 3, allowed: false, attempts: 3}, // Past the window: a new one opens at 1. {advance: 250 * time.Millisecond, key: "login|192.0.2.1", max: 3, allowed: true, attempts: 1, newWindow: true}, {key: "login|192.0.2.1", max: 3, allowed: true, attempts: 2}, } // runParity drives the table through s and checks every answer. func runParity(t *testing.T, s parityStore) { t.Helper() opened := map[string]time.Duration{} // advanced since the key's window opened last := map[string]time.Duration{} // last retryAfter per key for i, step := range parityTable { if step.advance > 0 { s.advance(step.advance) for k := range opened { opened[k] += step.advance } } allowed, attempts, retry := s.store.Attempt(step.key, step.max, parityDecay) if allowed != step.allowed || attempts != step.attempts { t.Fatalf("%s step %d (%s max %d): allowed=%v attempts=%d, want %v %d", s.name, i, step.key, step.max, allowed, attempts, step.allowed, step.attempts) } if retry <= 0 || retry > parityDecay { t.Fatalf("%s step %d: retryAfter %s outside (0, %s]", s.name, i, retry, parityDecay) } if step.newWindow { if retry < parityDecay-100*time.Millisecond { t.Fatalf("%s step %d: a new window answered retryAfter %s, want about %s", s.name, i, retry, parityDecay) } opened[step.key] = 0 } else { if prev, ok := last[step.key]; ok && retry > prev { t.Fatalf("%s step %d: retryAfter grew from %s to %s", s.name, i, prev, retry) } if limit := parityDecay - opened[step.key]; retry > limit { t.Fatalf("%s step %d: retryAfter %s, want at most %s (the window extended)", s.name, i, retry, limit) } } last[step.key] = retry } } func sleepFor(d time.Duration) { time.Sleep(d) } // miniStore returns a Redis rate-limit store on a fresh miniredis. func miniStore(t *testing.T, prefix string, fallback cooler.Attempter) (*miniredis.Miniredis, *cooler.RateLimitStore) { t.Helper() mr := miniredis.RunT(t) client, err := cooler.OpenRedis(t.Context(), "redis://"+mr.Addr()+"/0") if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = client.Close() }) return mr, cooler.NewRateLimitStore(client, prefix, fallback) } func TestRateLimitParity(t *testing.T) { t.Run("surf.MemoryStore", func(t *testing.T) { mem := surf.NewMemoryStore(0) runParity(t, parityStore{name: "memory", store: mem, advance: sleepFor}) }) t.Run("Redis store on miniredis", func(t *testing.T) { mr, store := miniStore(t, "acme:", nil) runParity(t, parityStore{name: "miniredis", store: store, advance: mr.FastForward}) }) } // countAdmitted runs workers concurrent attempts on one key and returns how // many were admitted. func countAdmitted(s surf.Store, workers, max int) int { var admitted atomic.Int64 var wg sync.WaitGroup start := make(chan struct{}) for range workers { wg.Add(1) go func() { defer wg.Done() <-start if ok, _, _ := s.Attempt("burst", max, time.Minute); ok { admitted.Add(1) } }() } close(start) wg.Wait() return int(admitted.Load()) } func TestRateLimitConcurrentNeverExceedsMax(t *testing.T) { const workers, max = 50, 10 if got := countAdmitted(surf.NewMemoryStore(0), workers, max); got != max { t.Fatalf("MemoryStore admitted %d of %d, want %d", got, workers, max) } mr, store := miniStore(t, "acme:", nil) if got := countAdmitted(store, workers, max); got != max { t.Fatalf("Redis store admitted %d of %d, want %d", got, workers, max) } if got, err := mr.Get("acme:rl:burst"); err != nil || got != fmt.Sprint(max) { t.Fatalf("Redis counter = %q %v, want %d (denials must not increment)", got, err, max) } } // captureDefault routes slog.Default to a buffer for the test. Stores built // after the call log there. func captureDefault(t *testing.T) *lockedBuffer { t.Helper() buf := &lockedBuffer{} prev := slog.Default() slog.SetDefault(slog.New(slog.NewTextHandler(buf, nil))) t.Cleanup(func() { slog.SetDefault(prev) }) return buf } type lockedBuffer struct { mu sync.Mutex buf bytes.Buffer } func (b *lockedBuffer) Write(p []byte) (int, error) { b.mu.Lock() defer b.mu.Unlock() return b.buf.Write(p) } func (b *lockedBuffer) lines() []string { b.mu.Lock() defer b.mu.Unlock() var out []string for _, l := range strings.Split(strings.TrimSpace(b.buf.String()), "\n") { if l != "" { out = append(out, l) } } return out } func TestRateLimitFallbackOnError(t *testing.T) { t.Run("a stopped Redis counts in the fallback and warns once", func(t *testing.T) { logs := captureDefault(t) fallback := surf.NewMemoryStore(0) mr, store := miniStore(t, "acme:", fallback) if ok, n, _ := store.Attempt("login", 2, time.Minute); !ok || n != 1 { t.Fatalf("with Redis up: %v %d", ok, n) } addr := mr.Addr() mr.Close() want := []struct { allowed bool attempts int }{{true, 1}, {true, 2}, {false, 2}} for i, w := range want { start := time.Now() ok, n, retry := store.Attempt("login", 2, time.Minute) if elapsed := time.Since(start); elapsed > cooler.AttemptTimeout+300*time.Millisecond { t.Fatalf("attempt %d took %s, want at most about %s", i, elapsed, cooler.AttemptTimeout) } if ok != w.allowed || n != w.attempts || retry <= 0 { t.Fatalf("fallback attempt %d = %v %d %s, want %v %d", i, ok, n, retry, w.allowed, w.attempts) } } // The counts live in the fallback store itself. if ok, n, _ := fallback.Attempt("login", 2, time.Minute); ok || n != 2 { t.Fatalf("fallback store = %v %d, want denied at 2", ok, n) } lines := logs.lines() if len(lines) != 1 || !strings.Contains(lines[0], "level=WARN") || !strings.Contains(lines[0], "rate-limit store failed") || !strings.Contains(lines[0], "endpoint="+addr) { t.Fatalf("log = %q, want one warning naming the endpoint", lines) } }) t.Run("without a fallback an outage admits", func(t *testing.T) { captureDefault(t) mr, store := miniStore(t, "acme:", nil) mr.Close() for range 3 { if ok, n, retry := store.Attempt("login", 1, time.Minute); !ok || n != 0 || retry != 0 { t.Fatalf("Attempt = %v %d %s, want true 0 0", ok, n, retry) } } }) t.Run("a store without a client uses the fallback", func(t *testing.T) { captureDefault(t) store := cooler.NewRateLimitStore(nil, "acme:", surf.NewMemoryStore(0)) if ok, n, _ := store.Attempt("k", 1, time.Minute); !ok || n != 1 { t.Fatalf("first = %v %d", ok, n) } if ok, n, _ := store.Attempt("k", 1, time.Minute); ok || n != 1 { t.Fatalf("second = %v %d", ok, n) } }) t.Run("a script error answers from the fallback", func(t *testing.T) { captureDefault(t) mr, store := miniStore(t, "acme:", surf.NewMemoryStore(0)) mr.SetError("LOADING Redis is loading the dataset in memory") if ok, n, _ := store.Attempt("k", 1, time.Minute); !ok || n != 1 { t.Fatalf("Attempt = %v %d", ok, n) } mr.SetError("") if ok, n, _ := store.Attempt("k", 1, time.Minute); !ok || n != 1 { t.Fatalf("after recovery Redis counts again from 1: %v %d", ok, n) } }) } func TestRateLimitKeyPrefix(t *testing.T) { mr, store := miniStore(t, "acme:", nil) store.Attempt("login|192.0.2.1", 5, 90*time.Second) if got, err := mr.Get("acme:rl:login|192.0.2.1"); err != nil || got != "1" { t.Fatalf("acme:rl:login|192.0.2.1 = %q %v", got, err) } if ttl := mr.TTL("acme:rl:login|192.0.2.1"); ttl != 90*time.Second { t.Fatalf("window TTL = %s, want 90s", ttl) } // A sub-millisecond decay is clamped to 1 ms rather than failing. if ok, _, _ := store.Attempt("tiny", 1, time.Microsecond); !ok { t.Fatal("a sub-ms decay must still admit") } if ttl := mr.TTL("acme:rl:tiny"); ttl != time.Millisecond { t.Fatalf("sub-ms window TTL = %s, want 1ms", ttl) } if keys := mr.Keys(); len(keys) != 2 { t.Fatalf("keys = %q, want only the two rl keys", keys) } // Through Install: the prefix derives from app.name and the store from // http.rate_limit.store. cfg := emptyConfig(t) mustSet(t, cfg, "app.name", "Acme Shop") mustSet(t, cfg, cooler.RateLimitStoreKey, "redis") mustSet(t, cfg, cooler.RedisURLKey, "redis://"+mr.Addr()+"/0") app := backpack.New(cfg) h, err := cooler.Install(context.Background(), app) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = h.Close() }) if !h.RateLimitRedis() { t.Fatal("RateLimitRedis = false with http.rate_limit.store redis") } fallback := surf.NewMemoryStore(0) s := h.RateLimitStore(fallback) if s == cooler.Attempter(fallback) { t.Fatal("RateLimitStore returned the fallback with the Redis store selected") } s.Attempt("login", 1, time.Minute) if !mr.Exists("acme_shop:rl:login") { t.Fatalf("keys = %q, want acme_shop:rl:login", mr.Keys()) } // cache.prefix wins over app.name. mustSet(t, cfg, cooler.CachePrefixKey, "shop1:") h2, err := cooler.Install(context.Background(), backpack.New(cfg)) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = h2.Close() }) h2.RateLimitStore(fallback).Attempt("login", 1, time.Minute) if !mr.Exists("shop1:rl:login") { t.Fatalf("keys = %q, want shop1:rl:login", mr.Keys()) } } // emptyConfig returns a config that ignores the process environment. func emptyConfig(t *testing.T) *compass.Config { t.Helper() cfg, err := compass.Open(compass.Options{Dir: t.TempDir(), Environ: []string{}}) if err != nil { t.Fatal(err) } return cfg } func mustSet(t *testing.T, cfg *compass.Config, key string, value any) { t.Helper() if err := cfg.Set(key, value); err != nil { t.Fatal(err) } } // deadAddr returns a localhost address nothing listens on. func deadAddr(t *testing.T) string { t.Helper() ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } addr := ln.Addr().String() _ = ln.Close() return addr }