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

@@ -19,6 +19,7 @@ import (
"github.com/hhs/camtalk/internal/config"
eino "github.com/hhs/camtalk/internal/eino"
"github.com/hhs/camtalk/internal/logger"
"github.com/hhs/camtalk/internal/ratelimit"
"github.com/hhs/camtalk/internal/session"
"github.com/hhs/camtalk/internal/store"
"github.com/hhs/camtalk/internal/ws"
@@ -195,6 +196,23 @@ func main() {
)
authService := auth.NewAuthService(tokenMgr, userRepo)
// 初始化限流器
var limiter ratelimit.Limiter
if cfg.RateLimit.Enabled {
if rdb != nil {
// 多实例:使用 Redis 令牌桶
limiter = ratelimit.NewRedisLimiter(rdb, cfg.RateLimit)
logger.Log.Info("rate limiter initialized with Redis backend")
} else {
// 单实例:使用内存令牌桶
limiter = ratelimit.NewMemoryLimiter(cfg.RateLimit)
logger.Log.Info("rate limiter initialized with in-memory backend")
}
defer limiter.Stop()
} else {
logger.Log.Info("rate limiter disabled")
}
// Gin 模式
if cfg.App.Env == "prod" {
gin.SetMode(gin.ReleaseMode)
@@ -215,14 +233,14 @@ func main() {
// Auth REST 端点
authHandler := api.NewAuthHandler(authService, tokenMgr)
authHandler.RegisterRoutes(apiGroup)
authHandler.RegisterRoutes(apiGroup, limiter)
// Conversation REST 端点
convHandler := api.NewConversationHandler(sessionMgr, tokenMgr, msgRepo)
convHandler.RegisterRoutes(apiGroup)
// WebSocket
r.GET("/ws", ws.ServeWS(sessionMgr, orch, cfg, tokenMgr))
r.GET("/ws", ws.ServeWS(sessionMgr, orch, cfg, tokenMgr, limiter))
// HTTP Server
srv := &http.Server{