chore: establish local platform baseline

2026-08-26 06:41:54 +08:00
commit 29bdf4900f
575 changed files with 95644 additions and 0 deletions
@@ -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}
+`
@@ -0,0 +1,40 @@
+package ratelimit
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+func TestLimiterEnforcesBurstAndRefills(t *testing.T) {
+	redisAddr := os.Getenv("TEST_REDIS_ADDR")
+	if redisAddr == "" {
+		t.Skip("TEST_REDIS_ADDR is not set")
+	}
+	client := redis.NewClient(&redis.Options{Addr: redisAddr})
+	defer func() { _ = client.Close() }()
+	limiter, err := New(client, "integration", Policy{RatePerSecond: 10, Burst: 2})
+	if err != nil {
+		t.Fatal(err)
+	}
+	identifier := fmt.Sprintf("test-%d", time.Now().UnixNano())
+	for attempt := 0; attempt < 2; attempt++ {
+		result, err := limiter.Allow(context.Background(), "user", identifier)
+		if err != nil || !result.Allowed {
+			t.Fatalf("attempt %d result=%#v err=%v", attempt, result, err)
+		}
+	}
+	limited, err := limiter.Allow(context.Background(), "user", identifier)
+	if err != nil || limited.Allowed || limited.RetryAfter <= 0 {
+		t.Fatalf("limited result=%#v err=%v", limited, err)
+	}
+	time.Sleep(150 * time.Millisecond)
+	refilled, err := limiter.Allow(context.Background(), "user", identifier)
+	if err != nil || !refilled.Allowed {
+		t.Fatalf("refilled result=%#v err=%v", refilled, err)
+	}
+}
@@ -0,0 +1,26 @@
+package ratelimit
+
+import "testing"
+
+func TestIdentifierHashDoesNotExposeIdentifier(t *testing.T) {
+	identifier := "192.0.2.10"
+	first := IdentifierHash(identifier)
+	if len(first) != 32 || first == identifier || first != IdentifierHash(identifier) {
+		t.Fatalf("unexpected identifier hash: %s", first)
+	}
+	if first == IdentifierHash("192.0.2.11") {
+		t.Fatal("different identifiers should not share a hash")
+	}
+}
+
+func TestNewRejectsInvalidPolicy(t *testing.T) {
+	if _, err := New(nil, "", Policy{}); err == nil {
+		t.Fatal("invalid limiter configuration should fail")
+	}
+}
+
+func TestNewPoliciesRejectsInvalidScopedPolicy(t *testing.T) {
+	if _, err := NewPolicies(nil, "prefix", map[string]Policy{"device": {RatePerSecond: 1, Burst: 1}}); err == nil {
+		t.Fatal("nil Redis client should fail")
+	}
+}