Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/middleware/rate_limiter.go

@@ -0,0 +1,205 @@
+package middleware
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+	"strconv"
+	"time"
+
+	ippkg "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
+
+	"github.com/gin-gonic/gin"
+	"github.com/redis/go-redis/v9"
+)
+
+// RateLimitFailureMode Redis 故障策略
+type RateLimitFailureMode int
+
+const (
+	RateLimitFailOpen RateLimitFailureMode = iota
+	RateLimitFailClose
+)
+
+// RateLimitOptions 限流可选配置
+type RateLimitOptions struct {
+	FailureMode RateLimitFailureMode
+}
+
+var rateLimitScript = redis.NewScript(`
+local current = redis.call('INCR', KEYS[1])
+local ttl = redis.call('PTTL', KEYS[1])
+local repaired = 0
+if current == 1 then
+  redis.call('PEXPIRE', KEYS[1], ARGV[1])
+elseif ttl == -1 then
+  redis.call('PEXPIRE', KEYS[1], ARGV[1])
+  repaired = 1
+end
+return {current, repaired}
+`)
+
+// rateLimitRun 允许测试覆写脚本执行逻辑
+var rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
+	values, err := rateLimitScript.Run(ctx, client, []string{key}, windowMillis).Slice()
+	if err != nil {
+		return 0, false, err
+	}
+	if len(values) < 2 {
+		return 0, false, fmt.Errorf("rate limit script returned %d values", len(values))
+	}
+	count, err := parseInt64(values[0])
+	if err != nil {
+		return 0, false, err
+	}
+	repaired, err := parseInt64(values[1])
+	if err != nil {
+		return 0, false, err
+	}
+	return count, repaired == 1, nil
+}
+
+// RateLimiter Redis 速率限制器
+type RateLimiter struct {
+	redis  *redis.Client
+	prefix string
+}
+
+// NewRateLimiter 创建速率限制器实例
+func NewRateLimiter(redisClient *redis.Client) *RateLimiter {
+	return &RateLimiter{
+		redis:  redisClient,
+		prefix: "rate_limit:",
+	}
+}
+
+// AllowResult 单次固定窗口限流判定结果。
+type AllowResult struct {
+	// Allowed 是否放行
+	Allowed bool
+	// Count 当前窗口内累计请求数（含本次）
+	Count int64
+	// RetryAfter 超限时距窗口重置的剩余时间（尽力而为；PTTL 不可用时回退为完整窗口）
+	RetryAfter time.Duration
+}
+
+// Allow 对给定 key（不含 "rate_limit:" 前缀）执行一次固定窗口计数判定。
+// 供需要自定义限流维度（如按用户 ID）的调用方使用；Redis 错误由调用方决定 fail-open/close。
+func (r *RateLimiter) Allow(ctx context.Context, key string, limit int, window time.Duration) (AllowResult, error) {
+	redisKey := r.prefix + key
+	windowMillis := windowTTLMillis(window)
+
+	count, repaired, err := rateLimitRun(ctx, r.redis, redisKey, windowMillis)
+	if err != nil {
+		return AllowResult{}, err
+	}
+	if repaired {
+		log.Printf("[RateLimit] ttl repaired: key=%s window_ms=%d", redisKey, windowMillis)
+	}
+
+	result := AllowResult{Allowed: count <= int64(limit), Count: count}
+	if !result.Allowed {
+		result.RetryAfter = window
+		if ttl, ttlErr := r.redis.PTTL(ctx, redisKey).Result(); ttlErr == nil && ttl > 0 {
+			result.RetryAfter = ttl
+		}
+	}
+	return result, nil
+}
+
+// clientIPForRateLimit 返回 IP 维度限流使用的客户端地址。
+// 与审计日志/会话绑定/API Key IP ACL 共用同一套安全客户端 IP 解析
+// （SessionBindingContext 快照：兼容开关开启时信任反代转发头，关闭时走
+// server.trusted_proxies 可信链）。避免默认反代部署下 Gin ClientIP 恒等于
+// 代理地址、所有用户坍缩进同一个限流桶造成整体误拦截。
+func clientIPForRateLimit(c *gin.Context) string {
+	if resolved := ippkg.GetSecurityClientIP(c, false); resolved != "" {
+		return resolved
+	}
+	return c.ClientIP()
+}
+
+// Limit 返回速率限制中间件
+// key: 限制类型标识
+// limit: 时间窗口内最大请求数
+// window: 时间窗口
+func (r *RateLimiter) Limit(key string, limit int, window time.Duration) gin.HandlerFunc {
+	return r.LimitWithOptions(key, limit, window, RateLimitOptions{})
+}
+
+// LimitWithOptions 返回速率限制中间件（带可选配置）
+func (r *RateLimiter) LimitWithOptions(key string, limit int, window time.Duration, opts RateLimitOptions) gin.HandlerFunc {
+	failureMode := opts.FailureMode
+	if failureMode != RateLimitFailClose {
+		failureMode = RateLimitFailOpen
+	}
+
+	return func(c *gin.Context) {
+		result, err := r.Allow(c.Request.Context(), key+":"+clientIPForRateLimit(c), limit, window)
+		if err != nil {
+			log.Printf("[RateLimit] redis error: key=%s mode=%s err=%v", r.prefix+key, failureModeLabel(failureMode), err)
+			if failureMode == RateLimitFailClose {
+				abortRateLimit(c, window)
+				return
+			}
+			// Redis 错误时放行，避免影响正常服务
+			c.Next()
+			return
+		}
+
+		// 超过限制
+		if !result.Allowed {
+			abortRateLimit(c, result.RetryAfter)
+			return
+		}
+
+		c.Next()
+	}
+}
+
+func windowTTLMillis(window time.Duration) int64 {
+	ttl := window.Milliseconds()
+	if ttl < 1 {
+		return 1
+	}
+	return ttl
+}
+
+func abortRateLimit(c *gin.Context, retryAfter time.Duration) {
+	if retryAfter > 0 {
+		seconds := int64(retryAfter / time.Second)
+		if retryAfter%time.Second > 0 {
+			seconds++
+		}
+		c.Header("Retry-After", strconv.FormatInt(seconds, 10))
+	}
+	c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
+		"error":   "rate limit exceeded",
+		"message": "Too many requests, please try again later",
+	})
+}
+
+func failureModeLabel(mode RateLimitFailureMode) string {
+	if mode == RateLimitFailClose {
+		return "fail-close"
+	}
+	return "fail-open"
+}
+
+func parseInt64(value any) (int64, error) {
+	switch v := value.(type) {
+	case int64:
+		return v, nil
+	case int:
+		return int64(v), nil
+	case string:
+		parsed, err := strconv.ParseInt(v, 10, 64)
+		if err != nil {
+			return 0, err
+		}
+		return parsed, nil
+	default:
+		return 0, fmt.Errorf("unexpected value type %T", value)
+	}
+}

backend/internal/middleware/rate_limiter_integration_test.go

@@ -0,0 +1,158 @@
+//go:build integration
+
+package middleware
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
+	tcredis "github.com/testcontainers/testcontainers-go/modules/redis"
+)
+
+const redisImageTag = "redis:8.4-alpine"
+
+func TestRateLimiterSetsTTLAndDoesNotRefresh(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	ctx := context.Background()
+	rdb := startRedis(t, ctx)
+	limiter := NewRateLimiter(rdb)
+
+	router := gin.New()
+	router.Use(limiter.Limit("ttl-test", 10, 2*time.Second))
+	router.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	recorder := performRequest(router)
+	require.Equal(t, http.StatusOK, recorder.Code)
+
+	redisKey := limiter.prefix + "ttl-test:127.0.0.1"
+	ttlBefore, err := rdb.PTTL(ctx, redisKey).Result()
+	require.NoError(t, err)
+	require.Greater(t, ttlBefore, time.Duration(0))
+	require.LessOrEqual(t, ttlBefore, 2*time.Second)
+
+	time.Sleep(50 * time.Millisecond)
+
+	recorder = performRequest(router)
+	require.Equal(t, http.StatusOK, recorder.Code)
+
+	ttlAfter, err := rdb.PTTL(ctx, redisKey).Result()
+	require.NoError(t, err)
+	require.Less(t, ttlAfter, ttlBefore)
+}
+
+func TestRateLimiterFixesMissingTTL(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	ctx := context.Background()
+	rdb := startRedis(t, ctx)
+	limiter := NewRateLimiter(rdb)
+
+	router := gin.New()
+	router.Use(limiter.Limit("ttl-missing", 10, 2*time.Second))
+	router.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	redisKey := limiter.prefix + "ttl-missing:127.0.0.1"
+	require.NoError(t, rdb.Set(ctx, redisKey, 5, 0).Err())
+
+	ttlBefore, err := rdb.PTTL(ctx, redisKey).Result()
+	require.NoError(t, err)
+	require.Less(t, ttlBefore, time.Duration(0))
+
+	recorder := performRequest(router)
+	require.Equal(t, http.StatusOK, recorder.Code)
+
+	ttlAfter, err := rdb.PTTL(ctx, redisKey).Result()
+	require.NoError(t, err)
+	require.Greater(t, ttlAfter, time.Duration(0))
+}
+
+func performRequest(router *gin.Engine) *httptest.ResponseRecorder {
+	req := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req.RemoteAddr = "127.0.0.1:1234"
+	recorder := httptest.NewRecorder()
+	router.ServeHTTP(recorder, req)
+	return recorder
+}
+
+func startRedis(t *testing.T, ctx context.Context) *redis.Client {
+	t.Helper()
+	ensureDockerAvailable(t)
+
+	redisContainer, err := tcredis.Run(ctx, redisImageTag)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = redisContainer.Terminate(ctx)
+	})
+
+	redisHost, err := redisContainer.Host(ctx)
+	require.NoError(t, err)
+	redisPort, err := redisContainer.MappedPort(ctx, "6379/tcp")
+	require.NoError(t, err)
+
+	rdb := redis.NewClient(&redis.Options{
+		Addr: fmt.Sprintf("%s:%d", redisHost, redisPort.Int()),
+		DB:   0,
+	})
+	require.NoError(t, rdb.Ping(ctx).Err())
+
+	t.Cleanup(func() {
+		_ = rdb.Close()
+	})
+
+	return rdb
+}
+
+func ensureDockerAvailable(t *testing.T) {
+	t.Helper()
+	if dockerAvailable() {
+		return
+	}
+	t.Skip("Docker 未启用，跳过依赖 testcontainers 的集成测试")
+}
+
+func dockerAvailable() bool {
+	if os.Getenv("DOCKER_HOST") != "" {
+		return true
+	}
+
+	socketCandidates := []string{
+		"/var/run/docker.sock",
+		filepath.Join(os.Getenv("XDG_RUNTIME_DIR"), "docker.sock"),
+		filepath.Join(userHomeDir(), ".docker", "run", "docker.sock"),
+		filepath.Join(userHomeDir(), ".docker", "desktop", "docker.sock"),
+		filepath.Join("/run/user", strconv.Itoa(os.Getuid()), "docker.sock"),
+	}
+
+	for _, socket := range socketCandidates {
+		if socket == "" {
+			continue
+		}
+		if _, err := os.Stat(socket); err == nil {
+			return true
+		}
+	}
+	return false
+}
+
+func userHomeDir() string {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return ""
+	}
+	return home
+}

backend/internal/middleware/rate_limiter_test.go

@@ -0,0 +1,226 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	ippkg "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
+
+	"github.com/gin-gonic/gin"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWindowTTLMillis(t *testing.T) {
+	require.Equal(t, int64(1), windowTTLMillis(500*time.Microsecond))
+	require.Equal(t, int64(1), windowTTLMillis(1500*time.Microsecond))
+	require.Equal(t, int64(2), windowTTLMillis(2500*time.Microsecond))
+}
+
+func TestRateLimiterFailureModes(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	rdb := redis.NewClient(&redis.Options{
+		Addr:         "127.0.0.1:1",
+		DialTimeout:  50 * time.Millisecond,
+		ReadTimeout:  50 * time.Millisecond,
+		WriteTimeout: 50 * time.Millisecond,
+	})
+	t.Cleanup(func() {
+		_ = rdb.Close()
+	})
+
+	limiter := NewRateLimiter(rdb)
+
+	failOpenRouter := gin.New()
+	failOpenRouter.Use(limiter.Limit("test", 1, time.Second))
+	failOpenRouter.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req.RemoteAddr = "127.0.0.1:1234"
+	recorder := httptest.NewRecorder()
+	failOpenRouter.ServeHTTP(recorder, req)
+	require.Equal(t, http.StatusOK, recorder.Code)
+
+	failCloseRouter := gin.New()
+	failCloseRouter.Use(limiter.LimitWithOptions("test", 1, time.Second, RateLimitOptions{
+		FailureMode: RateLimitFailClose,
+	}))
+	failCloseRouter.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	req = httptest.NewRequest(http.MethodGet, "/test", nil)
+	req.RemoteAddr = "127.0.0.1:1234"
+	recorder = httptest.NewRecorder()
+	failCloseRouter.ServeHTTP(recorder, req)
+	require.Equal(t, http.StatusTooManyRequests, recorder.Code)
+}
+
+func TestRateLimiterDifferentIPsIndependent(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	callCounts := make(map[string]int64)
+	originalRun := rateLimitRun
+	rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
+		callCounts[key]++
+		return callCounts[key], false, nil
+	}
+	t.Cleanup(func() {
+		rateLimitRun = originalRun
+	})
+
+	limiter := NewRateLimiter(redis.NewClient(&redis.Options{Addr: "127.0.0.1:1"}))
+
+	router := gin.New()
+	router.Use(limiter.Limit("api", 1, time.Second))
+	router.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	// 第一个 IP 的请求应通过
+	req1 := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req1.RemoteAddr = "10.0.0.1:1234"
+	rec1 := httptest.NewRecorder()
+	router.ServeHTTP(rec1, req1)
+	require.Equal(t, http.StatusOK, rec1.Code, "第一个 IP 的第一次请求应通过")
+
+	// 第二个 IP 的请求应独立通过（不受第一个 IP 的计数影响）
+	req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req2.RemoteAddr = "10.0.0.2:5678"
+	rec2 := httptest.NewRecorder()
+	router.ServeHTTP(rec2, req2)
+	require.Equal(t, http.StatusOK, rec2.Code, "第二个 IP 的第一次请求应独立通过")
+
+	// 第一个 IP 的第二次请求应被限流
+	req3 := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req3.RemoteAddr = "10.0.0.1:1234"
+	rec3 := httptest.NewRecorder()
+	router.ServeHTTP(rec3, req3)
+	require.Equal(t, http.StatusTooManyRequests, rec3.Code, "第一个 IP 的第二次请求应被限流")
+}
+
+func TestRateLimiterAllow(t *testing.T) {
+	originalRun := rateLimitRun
+	var gotKey string
+	count := int64(0)
+	rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
+		gotKey = key
+		count++
+		return count, false, nil
+	}
+	t.Cleanup(func() {
+		rateLimitRun = originalRun
+	})
+
+	// PTTL 走真实客户端（不可达地址）→ 失败后 RetryAfter 应回退为完整窗口
+	limiter := NewRateLimiter(redis.NewClient(&redis.Options{
+		Addr:         "127.0.0.1:1",
+		DialTimeout:  50 * time.Millisecond,
+		ReadTimeout:  50 * time.Millisecond,
+		WriteTimeout: 50 * time.Millisecond,
+	}))
+
+	res, err := limiter.Allow(context.Background(), "panel:global:user:42", 1, time.Minute)
+	require.NoError(t, err)
+	require.True(t, res.Allowed)
+	require.Equal(t, int64(1), res.Count)
+	require.Zero(t, res.RetryAfter)
+	require.Equal(t, "rate_limit:panel:global:user:42", gotKey)
+
+	res, err = limiter.Allow(context.Background(), "panel:global:user:42", 1, time.Minute)
+	require.NoError(t, err)
+	require.False(t, res.Allowed)
+	require.Equal(t, int64(2), res.Count)
+	require.Equal(t, time.Minute, res.RetryAfter)
+}
+
+func TestRateLimiterHonorsForwardedIPSnapshot(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	callCounts := make(map[string]int64)
+	originalRun := rateLimitRun
+	rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
+		callCounts[key]++
+		return callCounts[key], false, nil
+	}
+	t.Cleanup(func() {
+		rateLimitRun = originalRun
+	})
+
+	limiter := NewRateLimiter(redis.NewClient(&redis.Options{Addr: "127.0.0.1:1"}))
+
+	router := gin.New()
+	// 模拟 SessionBindingContext：开启转发 IP 兼容模式快照
+	router.Use(func(c *gin.Context) {
+		ippkg.SetForwardedIPSettings(c, true, nil)
+		c.Next()
+	})
+	router.Use(limiter.Limit("fwd", 1, time.Second))
+	router.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	send := func(xff string) int {
+		req := httptest.NewRequest(http.MethodGet, "/test", nil)
+		// 所有请求都来自同一个反代地址
+		req.RemoteAddr = "127.0.0.1:5678"
+		req.Header.Set("X-Forwarded-For", xff)
+		rec := httptest.NewRecorder()
+		router.ServeHTTP(rec, req)
+		return rec.Code
+	}
+
+	// 反代后两个不同的真实客户端应各自独立计数，不因共享代理地址被合并限流
+	require.Equal(t, http.StatusOK, send("198.51.100.1"))
+	require.Equal(t, http.StatusOK, send("198.51.100.2"))
+	// 同一真实客户端第二次请求应被限流
+	require.Equal(t, http.StatusTooManyRequests, send("198.51.100.1"))
+
+	require.Contains(t, callCounts, "rate_limit:fwd:198.51.100.1")
+	require.Contains(t, callCounts, "rate_limit:fwd:198.51.100.2")
+}
+
+func TestRateLimiterSuccessAndLimit(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	originalRun := rateLimitRun
+	counts := []int64{1, 2}
+	callIndex := 0
+	rateLimitRun = func(ctx context.Context, client *redis.Client, key string, windowMillis int64) (int64, bool, error) {
+		if callIndex >= len(counts) {
+			return counts[len(counts)-1], false, nil
+		}
+		value := counts[callIndex]
+		callIndex++
+		return value, false, nil
+	}
+	t.Cleanup(func() {
+		rateLimitRun = originalRun
+	})
+
+	limiter := NewRateLimiter(redis.NewClient(&redis.Options{Addr: "127.0.0.1:1"}))
+
+	router := gin.New()
+	router.Use(limiter.Limit("test", 1, time.Second))
+	router.GET("/test", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req.RemoteAddr = "127.0.0.1:1234"
+	recorder := httptest.NewRecorder()
+	router.ServeHTTP(recorder, req)
+	require.Equal(t, http.StatusOK, recorder.Code)
+
+	req = httptest.NewRequest(http.MethodGet, "/test", nil)
+	req.RemoteAddr = "127.0.0.1:1234"
+	recorder = httptest.NewRecorder()
+	router.ServeHTTP(recorder, req)
+	require.Equal(t, http.StatusTooManyRequests, recorder.Code)
+}