feat: 集成限流器到服务

- main.go 初始化限流器(根据 Redis 可用性选择内存/Redis 实现)
- WebSocket handler 添加 query 消息限流(按 userID)
- Auth API 添加登录/注册限流(按 IP)
- refresh 和 logout 不限流(避免影响正常用户操作)
- 修复所有测试(传递 nil limiter 参数)
- 所有测试通过(包括 ws 和 api 集成测试)
This commit is contained in:
hhs
2026-06-21 00:00:24 +08:00
parent ea00939c13
commit 7adf81c6e5
6 changed files with 240 additions and 44 deletions

View File

@@ -3,6 +3,7 @@ package ws
import (
"context"
"encoding/json"
"fmt"
"net/http"
"sync"
"time"
@@ -17,6 +18,7 @@ import (
"github.com/hhs/camtalk/internal/logger"
"github.com/hhs/camtalk/internal/models"
"github.com/hhs/camtalk/internal/orchestrator"
"github.com/hhs/camtalk/internal/ratelimit"
"github.com/hhs/camtalk/internal/session"
)
@@ -93,19 +95,19 @@ func (w *WSClient) SendError(err models.WsError) error {
}
// ServeWS 处理 WebSocket 升级请求。
func ServeWS(sessionMgr session.Manager, orch orchestrator.Orchestrator, cfg *config.Config, tokenMgr *auth.TokenManager) gin.HandlerFunc {
func ServeWS(sessionMgr session.Manager, orch orchestrator.Orchestrator, cfg *config.Config, tokenMgr *auth.TokenManager, limiter ratelimit.Limiter) gin.HandlerFunc {
upgrader := newUpgrader(cfg)
heartbeatInterval := time.Duration(cfg.Server.HeartbeatInterval) * time.Second
heartbeatTimeout := time.Duration(cfg.Server.HeartbeatTimeout) * time.Second
version := cfg.App.Version
return func(c *gin.Context) {
serveWS(c, sessionMgr, orch, upgrader, heartbeatInterval, heartbeatTimeout, version, tokenMgr)
serveWS(c, sessionMgr, orch, upgrader, heartbeatInterval, heartbeatTimeout, version, tokenMgr, limiter)
}
}
func serveWS(c *gin.Context, sessionMgr session.Manager, orch orchestrator.Orchestrator,
upgrader websocket.Upgrader, heartbeatInterval, heartbeatTimeout time.Duration, version string, tokenMgr *auth.TokenManager) {
upgrader websocket.Upgrader, heartbeatInterval, heartbeatTimeout time.Duration, version string, tokenMgr *auth.TokenManager, limiter ratelimit.Limiter) {
// --- JWT 认证upgrade 前完成,失败直接返回 HTTP 错误) ---
token := c.Query("token")
@@ -225,6 +227,18 @@ func serveWS(c *gin.Context, sessionMgr session.Manager, orch orchestrator.Orche
}
logger.Log.Infow("query received", "session", sessionID, "request", msg.RequestID)
// 限流检查
if limiter != nil {
key := fmt.Sprintf("%s:query", userID)
allowed, retryAfter := limiter.Allow(context.Background(), key)
if !allowed {
logger.Log.Warnw("rate limited", "user_id", userID, "retry_after", retryAfter)
errors.SendWSError(client, errors.CodeRateLimited, msg.RequestID,
fmt.Errorf("rate limited, retry after %s", retryAfter.Round(time.Second)))
continue
}
}
// 刷新会话 TTL
if err := client.sessionMgr.Touch(context.Background(), sessionID); err != nil {
logger.Log.Warnw("touch session failed", "session", sessionID, "error", err)