starter code

internal/ratelimit/memory.go

@@ -0,0 +1,80 @@
+package ratelimit
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+// Memory is a fixed-window limiter held in process. Useful for tests and for a
+// single-instance self-hoster who would rather not run Redis.
+//
+// Not suitable for more than one backend process: each would keep its own
+// counters, so N processes allow N times the limit.
+type Memory struct {
+	mu      sync.Mutex
+	windows map[string]*window
+	Now     func() time.Time // injectable so tests need not sleep
+}
+
+type window struct {
+	count   int
+	expires time.Time
+}
+
+func NewMemory() *Memory {
+	return &Memory{
+		windows: make(map[string]*window),
+		Now:     time.Now,
+	}
+}
+
+func (m *Memory) Allow(_ context.Context, key string, limit int, dur time.Duration) (Result, error) {
+	if limit <= 0 {
+		return Result{Allowed: true, Limit: 0, Remaining: 0}, nil
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	now := m.Now()
+
+	w, ok := m.windows[key]
+	if !ok || now.After(w.expires) {
+		w = &window{expires: now.Add(dur)}
+		m.windows[key] = w
+	}
+
+	w.count++
+
+	remaining := limit - w.count
+	if remaining < 0 {
+		remaining = 0
+	}
+
+	return Result{
+		Allowed:    w.count <= limit,
+		Limit:      limit,
+		Remaining:  remaining,
+		RetryAfter: w.expires.Sub(now),
+	}, nil
+}
+
+// Sweep drops expired windows. Without it the map grows for every key ever
+// seen, which for per-user keys is unbounded. Call it periodically.
+func (m *Memory) Sweep() int {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	now := m.Now()
+	removed := 0
+
+	for k, w := range m.windows {
+		if now.After(w.expires) {
+			delete(m.windows, k)
+			removed++
+		}
+	}
+
+	return removed
+}

internal/ratelimit/ratelimit.go

@@ -0,0 +1,62 @@
+// Package ratelimit enforces per-key and per-end-user request ceilings.
+//
+// Redis is authoritative for the live window; api.usage_daily holds completed
+// days. A Redis flush therefore loses at most the window in progress, which for
+// abuse detection is an acceptable trade.
+package ratelimit
+
+import (
+	"context"
+	"time"
+)
+
+// Result is what a limiter says about one request.
+type Result struct {
+	Allowed    bool
+	Limit      int
+	Remaining  int
+	RetryAfter time.Duration
+}
+
+// Limiter is deliberately narrow so the middleware can be tested without Redis,
+// and so a self-hoster running a single instance could drop in an in-process
+// implementation instead of standing up Redis at all.
+type Limiter interface {
+	// Allow records one request against key and reports whether it may
+	// proceed. Implementations must be atomic: a check-then-increment race
+	// under load is exactly when the limit matters most.
+	Allow(ctx context.Context, key string, limit int, window time.Duration) (Result, error)
+}
+
+// Quota is the ceiling for one tier, from api.keys.quota_tier.
+type Quota struct {
+	// Requests allowed per window for the key as a whole.
+	KeyLimit int
+
+	// Requests allowed per window for a single end user within that key. One
+	// Discord user should not be able to consume a whole self-hosted
+	// instance's allowance, and a key whose traffic is 90% one user is the
+	// enumeration signal worth alerting on.
+	UserLimit int
+
+	Window time.Duration
+}
+
+// Tiers maps api.keys.quota_tier to its ceiling. Values are a starting point,
+// not a considered policy: watch real usage before hardening them.
+var Tiers = map[string]Quota{
+	"default":  {KeyLimit: 600, UserLimit: 120, Window: time.Minute},
+	"verified": {KeyLimit: 3000, UserLimit: 600, Window: time.Minute},
+	"trusted":  {KeyLimit: 12000, UserLimit: 2400, Window: time.Minute},
+	"internal": {KeyLimit: 0, UserLimit: 0, Window: time.Minute}, // 0 = unlimited
+}
+
+// QuotaFor falls back to the tightest tier for an unknown name, so a typo in
+// the database cannot accidentally grant unlimited access.
+func QuotaFor(tier string) Quota {
+	if q, ok := Tiers[tier]; ok {
+		return q
+	}
+
+	return Tiers["default"]
+}

internal/ratelimit/ratelimit_test.go

@@ -0,0 +1,208 @@
+package ratelimit
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestMemoryWindow(t *testing.T) {
+	now := time.Date(2026, 9, 4, 12, 0, 0, 0, time.UTC)
+	m := NewMemory()
+	m.Now = func() time.Time { return now }
+
+	ctx := context.Background()
+
+	for i := 1; i <= 3; i++ {
+		res, err := m.Allow(ctx, "k", 3, time.Minute)
+		if err != nil || !res.Allowed {
+			t.Fatalf("request %d: allowed=%v err=%v", i, res.Allowed, err)
+		}
+		if res.Remaining != 3-i {
+			t.Errorf("request %d: remaining %d, want %d", i, res.Remaining, 3-i)
+		}
+	}
+
+	res, _ := m.Allow(ctx, "k", 3, time.Minute)
+	if res.Allowed {
+		t.Error("4th request allowed past a limit of 3")
+	}
+	if res.RetryAfter <= 0 {
+		t.Error("RetryAfter not set on a rejection")
+	}
+
+	// Window rolls over.
+	now = now.Add(time.Minute + time.Second)
+	if res, _ := m.Allow(ctx, "k", 3, time.Minute); !res.Allowed {
+		t.Error("request rejected after the window expired")
+	}
+}
+
+func TestMemoryKeysAreIndependent(t *testing.T) {
+	m := NewMemory()
+	ctx := context.Background()
+
+	for i := 0; i < 5; i++ {
+		m.Allow(ctx, "a", 5, time.Minute)
+	}
+
+	if res, _ := m.Allow(ctx, "a", 5, time.Minute); res.Allowed {
+		t.Error("key a not limited")
+	}
+	if res, _ := m.Allow(ctx, "b", 5, time.Minute); !res.Allowed {
+		t.Error("key b affected by key a")
+	}
+}
+
+func TestMemoryZeroLimitIsUnlimited(t *testing.T) {
+	m := NewMemory()
+	for i := 0; i < 1000; i++ {
+		if res, _ := m.Allow(context.Background(), "k", 0, time.Minute); !res.Allowed {
+			t.Fatalf("zero limit rejected at %d", i)
+		}
+	}
+}
+
+func TestMemorySweep(t *testing.T) {
+	now := time.Now()
+	m := NewMemory()
+	m.Now = func() time.Time { return now }
+
+	m.Allow(context.Background(), "a", 5, time.Minute)
+	m.Allow(context.Background(), "b", 5, time.Hour)
+
+	if n := m.Sweep(); n != 0 {
+		t.Errorf("swept %d live windows", n)
+	}
+
+	now = now.Add(2 * time.Minute)
+
+	if n := m.Sweep(); n != 1 {
+		t.Errorf("swept %d, want 1 (a expired, b did not)", n)
+	}
+}
+
+func TestMemoryIsConcurrencySafe(t *testing.T) {
+	m := NewMemory()
+	ctx := context.Background()
+
+	const goroutines, each = 20, 50
+
+	done := make(chan int, goroutines)
+
+	for g := 0; g < goroutines; g++ {
+		go func() {
+			allowed := 0
+			for i := 0; i < each; i++ {
+				if res, _ := m.Allow(ctx, "shared", 100, time.Minute); res.Allowed {
+					allowed++
+				}
+			}
+			done <- allowed
+		}()
+	}
+
+	total := 0
+	for g := 0; g < goroutines; g++ {
+		total += <-done
+	}
+
+	// Exactly 100 of the 1000 attempts may pass. A check-then-increment race
+	// would let more through.
+	if total != 100 {
+		t.Errorf("%d requests allowed under concurrency, want exactly 100", total)
+	}
+}
+
+// fakeRunner stands in for Redis so the script result handling can be tested.
+type fakeRunner struct {
+	ret  any
+	err  error
+	keys []string
+	args []any
+}
+
+func (f *fakeRunner) Eval(_ context.Context, _ string, keys []string, args ...any) (any, error) {
+	f.keys, f.args = keys, args
+	return f.ret, f.err
+}
+
+func TestRedisResultParsing(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("under limit", func(t *testing.T) {
+		f := &fakeRunner{ret: []any{int64(3), int64(45000)}}
+		r := NewRedis(f)
+
+		res, err := r.Allow(ctx, "key:1", 10, time.Minute)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !res.Allowed || res.Remaining != 7 {
+			t.Errorf("allowed=%v remaining=%d, want true/7", res.Allowed, res.Remaining)
+		}
+		if res.RetryAfter != 45*time.Second {
+			t.Errorf("RetryAfter %v, want 45s", res.RetryAfter)
+		}
+		if len(f.keys) != 1 || f.keys[0] != "rl:key:1" {
+			t.Errorf("keys %v, want [rl:key:1]", f.keys)
+		}
+	})
+
+	t.Run("over limit", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{int64(11), int64(1000)}})
+		res, _ := r.Allow(ctx, "k", 10, time.Minute)
+		if res.Allowed || res.Remaining != 0 {
+			t.Errorf("allowed=%v remaining=%d, want false/0", res.Allowed, res.Remaining)
+		}
+	})
+
+	t.Run("exactly at limit is allowed", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{int64(10), int64(1000)}})
+		if res, _ := r.Allow(ctx, "k", 10, time.Minute); !res.Allowed {
+			t.Error("request at exactly the limit was rejected")
+		}
+	})
+
+	t.Run("plain int accepted", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{1, 1000}})
+		if _, err := r.Allow(ctx, "k", 10, time.Minute); err != nil {
+			t.Errorf("int result rejected: %v", err)
+		}
+	})
+
+	t.Run("missing ttl falls back to the window", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{int64(11), int64(-1)}})
+		res, _ := r.Allow(ctx, "k", 10, time.Minute)
+		if res.RetryAfter != time.Minute {
+			t.Errorf("RetryAfter %v, want 1m", res.RetryAfter)
+		}
+	})
+
+	t.Run("eval error propagates", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{err: errors.New("nope")})
+		if _, err := r.Allow(ctx, "k", 10, time.Minute); err == nil {
+			t.Error("expected an error so the caller can fail open")
+		}
+	})
+
+	t.Run("garbage result is an error not a panic", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: "not a list"})
+		if _, err := r.Allow(ctx, "k", 10, time.Minute); err == nil {
+			t.Error("expected an error")
+		}
+	})
+
+	t.Run("zero limit skips redis entirely", func(t *testing.T) {
+		f := &fakeRunner{ret: []any{int64(1), int64(1)}}
+		r := NewRedis(f)
+		res, _ := r.Allow(ctx, "k", 0, time.Minute)
+		if !res.Allowed {
+			t.Error("zero limit rejected")
+		}
+		if f.keys != nil {
+			t.Error("redis was called for an unlimited tier")
+		}
+	})
+}

internal/ratelimit/redis.go

@@ -0,0 +1,98 @@
+package ratelimit
+
+import (
+	"context"
+	"fmt"
+	"time"
+)
+
+// ScriptRunner is the slice of a Redis client this package needs. Narrow on
+// purpose: it keeps go-redis out of this file's imports, which means the window
+// logic below is testable without a Redis server. See redis_adapter.go for the
+// ten-line bridge to *redis.Client.
+type ScriptRunner interface {
+	Eval(ctx context.Context, script string, keys []string, args ...any) (any, error)
+}
+
+// INCR and EXPIRE must be one atomic step. Done as two commands, a crash
+// between them leaves a counter with no TTL, which silently becomes a permanent
+// ban for that key.
+//
+// Returns {count, ttl_ms}.
+const allowScript = `
+local current = redis.call('INCR', KEYS[1])
+if current == 1 then
+  redis.call('PEXPIRE', KEYS[1], ARGV[1])
+end
+return {current, redis.call('PTTL', KEYS[1])}
+`
+
+// Redis is a fixed-window limiter.
+//
+// Fixed window has a known flaw: a client can send `limit` requests at the end
+// of one window and `limit` more at the start of the next, briefly achieving
+// twice the rate. A sliding window or token bucket avoids it at the cost of
+// more state. For protecting a backend from a runaway bot instance the simple
+// version is enough, and the burst is bounded at 2x rather than unbounded.
+type Redis struct {
+	Client ScriptRunner
+	Prefix string // namespace, so SCAN can find these keys when debugging
+}
+
+func NewRedis(c ScriptRunner) *Redis {
+	return &Redis{Client: c, Prefix: "rl:"}
+}
+
+func (r *Redis) Allow(ctx context.Context, key string, limit int, dur time.Duration) (Result, error) {
+	if limit <= 0 {
+		return Result{Allowed: true}, nil
+	}
+
+	raw, err := r.Client.Eval(ctx, allowScript,
+		[]string{r.Prefix + key}, dur.Milliseconds())
+	if err != nil {
+		return Result{}, fmt.Errorf("ratelimit: eval: %w", err)
+	}
+
+	vals, ok := raw.([]any)
+	if !ok || len(vals) != 2 {
+		return Result{}, fmt.Errorf("ratelimit: unexpected script result %T", raw)
+	}
+
+	count, ok1 := toInt64(vals[0])
+	ttlMS, ok2 := toInt64(vals[1])
+
+	if !ok1 || !ok2 {
+		return Result{}, fmt.Errorf("ratelimit: non-integer script result")
+	}
+
+	remaining := limit - int(count)
+	if remaining < 0 {
+		remaining = 0
+	}
+
+	retry := time.Duration(ttlMS) * time.Millisecond
+	if ttlMS < 0 {
+		retry = dur
+	}
+
+	return Result{
+		Allowed:    int(count) <= limit,
+		Limit:      limit,
+		Remaining:  remaining,
+		RetryAfter: retry,
+	}, nil
+}
+
+// Redis integers arrive as int64 through go-redis, but a fake or a future
+// client version may hand back int. Accept both rather than panicking.
+func toInt64(v any) (int64, bool) {
+	switch n := v.(type) {
+	case int64:
+		return n, true
+	case int:
+		return int64(n), true
+	default:
+		return 0, false
+	}
+}

internal/ratelimit/redis_adapter.go

@@ -0,0 +1,20 @@
+package ratelimit
+
+import (
+	"context"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// GoRedis bridges *redis.Client to the ScriptRunner interface.
+//
+// Kept in its own file, and as thin as possible, because it is the one piece
+// here that cannot be exercised without a live Redis: everything in redis.go is
+// tested against a fake.
+type GoRedis struct {
+	Client *redis.Client
+}
+
+func (g GoRedis) Eval(ctx context.Context, script string, keys []string, args ...any) (any, error) {
+	return g.Client.Eval(ctx, script, keys, args...).Result()
+}