diff --git a/backend/cmd/server/main.go b/backend/cmd/server/main.go index 6100c19..2a6275b 100644 --- a/backend/cmd/server/main.go +++ b/backend/cmd/server/main.go @@ -145,7 +145,7 @@ func main() { convHandler.RegisterRoutes(apiGroup) // WebSocket - r.GET("/ws", ws.ServeWS(sessionMgr, orch, cfg)) + r.GET("/ws", ws.ServeWS(sessionMgr, orch, cfg, tokenMgr)) // HTTP Server srv := &http.Server{ diff --git a/backend/internal/ws/handler.go b/backend/internal/ws/handler.go index 25d0f97..7274b68 100644 --- a/backend/internal/ws/handler.go +++ b/backend/internal/ws/handler.go @@ -10,6 +10,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + "github.com/hhs/camtalk/internal/auth" "github.com/hhs/camtalk/internal/config" "github.com/hhs/camtalk/internal/errors" "github.com/hhs/camtalk/internal/logger" @@ -91,7 +92,7 @@ func (w *WSClient) SendError(err models.WsError) error { } // ServeWS 处理 WebSocket 升级请求。 -func ServeWS(sessionMgr session.Manager, orch orchestrator.Orchestrator, cfg *config.Config) gin.HandlerFunc { +func ServeWS(sessionMgr session.Manager, orch orchestrator.Orchestrator, cfg *config.Config, tokenMgr *auth.TokenManager) gin.HandlerFunc { upgrader := newUpgrader(cfg) heartbeatInterval := time.Duration(cfg.Server.HeartbeatInterval) * time.Second heartbeatTimeout := time.Duration(cfg.Server.HeartbeatTimeout) * time.Second @@ -100,12 +101,12 @@ func ServeWS(sessionMgr session.Manager, orch orchestrator.Orchestrator, cfg *co maxHistory := cfg.Session.MaxHistory return func(c *gin.Context) { - serveWS(c, sessionMgr, orch, upgrader, heartbeatInterval, heartbeatTimeout, version, maxHistory) + serveWS(c, sessionMgr, orch, upgrader, heartbeatInterval, heartbeatTimeout, version, maxHistory, tokenMgr) } } func serveWS(c *gin.Context, sessionMgr session.Manager, orch orchestrator.Orchestrator, - upgrader websocket.Upgrader, heartbeatInterval, heartbeatTimeout time.Duration, version string, maxHistory int) { + upgrader websocket.Upgrader, heartbeatInterval, heartbeatTimeout time.Duration, version string, maxHistory int, tokenMgr *auth.TokenManager) { conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { logger.Log.Errorw("websocket upgrade failed", "error", err) diff --git a/backend/internal/ws/handler_test.go b/backend/internal/ws/handler_test.go index 2872d5c..7fc43cb 100644 --- a/backend/internal/ws/handler_test.go +++ b/backend/internal/ws/handler_test.go @@ -13,6 +13,7 @@ import ( "github.com/stretchr/testify/require" "context" + "github.com/hhs/camtalk/internal/auth" "github.com/hhs/camtalk/internal/config" "github.com/hhs/camtalk/internal/logger" "github.com/hhs/camtalk/internal/models" @@ -144,7 +145,8 @@ 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 := auth.NewTokenManager("test-secret", 15*time.Minute, 7*24*time.Hour) + r.GET("/ws", ServeWS(sessionMgr, orch, cfg, tokenMgr)) srv := httptest.NewServer(r)