starter code

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
+	}
+}