feat: 集成限流器到服务
- main.go 初始化限流器(根据 Redis 可用性选择内存/Redis 实现) - WebSocket handler 添加 query 消息限流(按 userID) - Auth API 添加登录/注册限流(按 IP) - refresh 和 logout 不限流(避免影响正常用户操作) - 修复所有测试(传递 nil limiter 参数) - 所有测试通过(包括 ws 和 api 集成测试)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -148,7 +148,7 @@ func setupTestServer(t *testing.T, orch orchestrator.Orchestrator) (*httptest.Se
|
||||
Server: config.ServerConfig{HeartbeatInterval: 30, HeartbeatTimeout: 60},
|
||||
Session: config.SessionConfig{MaxHistory: 20},
|
||||
}
|
||||
r.GET("/ws", ServeWS(sessionMgr, orch, cfg, tokenMgr))
|
||||
r.GET("/ws", ServeWS(sessionMgr, orch, cfg, tokenMgr, nil))
|
||||
|
||||
srv := httptest.NewServer(r)
|
||||
|
||||
@@ -591,7 +591,7 @@ func setupTestServerEx(t *testing.T, orch orchestrator.Orchestrator) (*httptest.
|
||||
Server: config.ServerConfig{HeartbeatInterval: 30, HeartbeatTimeout: 60},
|
||||
Session: config.SessionConfig{MaxHistory: 20},
|
||||
}
|
||||
r.GET("/ws", ServeWS(sessionMgr, orch, cfg, tokenMgr))
|
||||
r.GET("/ws", ServeWS(sessionMgr, orch, cfg, tokenMgr, nil))
|
||||
|
||||
srv := httptest.NewServer(r)
|
||||
return srv, tokenMgr, sessionMgr
|
||||
@@ -642,7 +642,7 @@ func TestWS_AuthExpiredToken(t *testing.T) {
|
||||
Server: config.ServerConfig{HeartbeatInterval: 30, HeartbeatTimeout: 60},
|
||||
Session: config.SessionConfig{MaxHistory: 20},
|
||||
}
|
||||
r.GET("/ws", ServeWS(sessionMgr, &MockOrchestrator{}, cfg, tokenMgr))
|
||||
r.GET("/ws", ServeWS(sessionMgr, &MockOrchestrator{}, cfg, tokenMgr, nil))
|
||||
srv := httptest.NewServer(r)
|
||||
defer srv.Close()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user