summercms/modules/cooler/ratelimit_test.go

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
}