chore: establish local platform baseline

internal/ratelimit/limiter.go

@@ -0,0 +1,115 @@
+package ratelimit
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+type Policy struct {
+	RatePerSecond float64
+	Burst         int
+}
+
+type Result struct {
+	Allowed    bool
+	Remaining  int
+	RetryAfter time.Duration
+}
+
+type Limiter struct {
+	client   *redis.Client
+	prefix   string
+	policies map[string]Policy
+	script   *redis.Script
+}
+
+func New(client *redis.Client, prefix string, policy Policy) (*Limiter, error) {
+	if client == nil || prefix == "" || policy.RatePerSecond <= 0 || policy.Burst <= 0 {
+		return nil, errors.New("redis client, prefix and positive rate policy are required")
+	}
+	return NewPolicies(client, prefix, map[string]Policy{"*": policy})
+}
+
+// NewPolicies builds a limiter whose token-bucket policy is selected by scope.
+// A "*" policy is used as the fallback for unrecognised scopes.
+func NewPolicies(client *redis.Client, prefix string, policies map[string]Policy) (*Limiter, error) {
+	if client == nil || prefix == "" || len(policies) == 0 {
+		return nil, errors.New("redis client, prefix and rate policies are required")
+	}
+	cloned := make(map[string]Policy, len(policies))
+	for scope, policy := range policies {
+		if scope == "" || policy.RatePerSecond <= 0 || policy.Burst <= 0 {
+			return nil, errors.New("rate policy scopes, rates and bursts must be positive")
+		}
+		cloned[scope] = policy
+	}
+	return &Limiter{client: client, prefix: prefix, policies: cloned, script: redis.NewScript(tokenBucketScript)}, nil
+}
+
+func (l *Limiter) Allow(ctx context.Context, scope, identifier string) (Result, error) {
+	if scope == "" || identifier == "" {
+		return Result{}, errors.New("rate limit scope and identifier are required")
+	}
+	policy, ok := l.policies[scope]
+	if !ok {
+		policy, ok = l.policies["*"]
+	}
+	if !ok {
+		return Result{}, fmt.Errorf("rate limit policy for scope %q is not configured", scope)
+	}
+	key := fmt.Sprintf("%s:ratelimit:%s:%s", l.prefix, scope, IdentifierHash(identifier))
+	values, err := l.script.Run(ctx, l.client, []string{key}, policy.RatePerSecond, policy.Burst).Slice()
+	if err != nil {
+		return Result{}, fmt.Errorf("execute rate limit: %w", err)
+	}
+	if len(values) != 3 {
+		return Result{}, errors.New("unexpected rate limit response")
+	}
+	allowed, ok1 := values[0].(int64)
+	remaining, ok2 := values[1].(int64)
+	retryMilliseconds, ok3 := values[2].(int64)
+	if !ok1 || !ok2 || !ok3 {
+		return Result{}, errors.New("invalid rate limit response types")
+	}
+	return Result{
+		Allowed: allowed == 1, Remaining: int(remaining),
+		RetryAfter: time.Duration(retryMilliseconds) * time.Millisecond,
+	}, nil
+}
+
+func IdentifierHash(identifier string) string {
+	digest := sha256.Sum256([]byte(identifier))
+	return hex.EncodeToString(digest[:16])
+}
+
+const tokenBucketScript = `
+local current_time = redis.call('TIME')
+local now = tonumber(current_time[1]) + tonumber(current_time[2]) / 1000000
+local rate = tonumber(ARGV[1])
+local burst = tonumber(ARGV[2])
+local values = redis.call('HMGET', KEYS[1], 'tokens', 'updated_at')
+local tokens = tonumber(values[1])
+local updated_at = tonumber(values[2])
+if tokens == nil then tokens = burst end
+if updated_at == nil then updated_at = now end
+if now > updated_at then
+  tokens = math.min(burst, tokens + (now - updated_at) * rate)
+end
+local allowed = 0
+local retry_ms = 0
+if tokens >= 1 then
+  allowed = 1
+  tokens = tokens - 1
+else
+  retry_ms = math.ceil(((1 - tokens) / rate) * 1000)
+end
+redis.call('HSET', KEYS[1], 'tokens', tokens, 'updated_at', now)
+redis.call('PEXPIRE', KEYS[1], math.ceil((burst / rate) * 2000) + 1000)
+return {allowed, math.floor(tokens), retry_ms}
+`