package ws import ( "encoding/json" "errors" "log" "net/http" "strings" "github.com/edu-cloud/push-gateway/internal/hub" "github.com/gin-gonic/gin" "github.com/golang-jwt/jwt/v5" "github.com/gorilla/websocket" ) var ( errMissingToken = errors.New("missing token") errInvalidToken = errors.New("invalid token") errInvalidClaims = errors.New("invalid claims") errMissingUserID = errors.New("missing user id") ) var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true // P5 骨架,生产环境需校验 origin }, } // internalAPIKeyHeader 是内部 API 鉴权请求头名称 const internalAPIKeyHeader = "X-Internal-Key" // Handler 处理 WebSocket 升级与内部推送 API type Handler struct { hub *hub.Hub jwtSecret string devMode bool internalAPIKey string } // NewHandler 创建 Handler 实例 func NewHandler(h *hub.Hub, jwtSecret string, devMode bool, internalAPIKey string) *Handler { return &Handler{hub: h, jwtSecret: jwtSecret, devMode: devMode, internalAPIKey: internalAPIKey} } // HandleWebSocket 升级 HTTP 为 WebSocket,保持长连接并处理心跳 func (h *Handler) HandleWebSocket(c *gin.Context) { userID, err := h.authenticate(c) if err != nil { c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()}) return } conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { log.Printf("[ws] upgrade failed, user=%s: %v", userID, err) return } defer conn.Close() connPtr := h.hub.Register(userID, conn) defer h.hub.Unregister(connPtr) // 写协程:从 send chan 读取并写入 WebSocket,避免并发写 go func() { for msg := range connPtr.Outgoing() { if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil { log.Printf("[ws] write failed, user=%s: %v", userID, err) return } } }() // 读取循环(保持连接,处理心跳) for { _, msg, err := conn.ReadMessage() if err != nil { if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { log.Printf("[ws] read failed, user=%s: %v", userID, err) } break } if strings.ToLower(string(msg)) == "ping" { connPtr.Send([]byte("pong")) } } } // PushHandler 接收来自 Msg 服务的定向推送请求 func (h *Handler) PushHandler(c *gin.Context) { if !h.checkInternalAPIKey(c) { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "success": false, "error": gin.H{"code": "UNAUTHORIZED", "message": "invalid or missing internal api key"}, }) return } var req struct { UserID string `json:"userId"` Event string `json:"event"` Data map[string]any `json:"data"` } if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_REQUEST", "message": err.Error()}}) return } message, err := buildMessage(req.Event, req.Data) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_PAYLOAD", "message": err.Error()}}) return } if err := h.hub.SendToUser(req.UserID, message); err != nil { c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": gin.H{"code": "PUSH_FAILED", "message": err.Error()}}) return } c.JSON(http.StatusOK, gin.H{"success": true}) } // BroadcastHandler 接收来自 Msg 服务的广播请求 func (h *Handler) BroadcastHandler(c *gin.Context) { if !h.checkInternalAPIKey(c) { c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ "success": false, "error": gin.H{"code": "UNAUTHORIZED", "message": "invalid or missing internal api key"}, }) return } var req struct { Event string `json:"event"` Data map[string]any `json:"data"` } if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_REQUEST", "message": err.Error()}}) return } message, err := buildMessage(req.Event, req.Data) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_PAYLOAD", "message": err.Error()}}) return } h.hub.Broadcast(message) c.JSON(http.StatusOK, gin.H{"success": true}) } // checkInternalAPIKey 校验内部 API 请求头 X-Internal-Key。 // DevMode 下跳过校验;非 DevMode 下要求 INTERNAL_API_KEY 已配置且与请求头匹配。 func (h *Handler) checkInternalAPIKey(c *gin.Context) bool { if h.devMode { return true } if h.internalAPIKey == "" { log.Println("[ws] internal api key not configured, rejecting request") return false } return c.GetHeader(internalAPIKeyHeader) == h.internalAPIKey } // authenticate 校验 WebSocket 连接的 JWT 或 dev token func (h *Handler) authenticate(c *gin.Context) (string, error) { tokenStr := c.Query("token") if tokenStr == "" { if auth := c.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") { tokenStr = strings.TrimPrefix(auth, "Bearer ") } } if tokenStr == "" { return "", errMissingToken } // DEV_MODE 下接受 dev-token,便于本地联调 if h.devMode && tokenStr == "dev-token" { return "dev-user", nil } token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (any, error) { if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok { return nil, jwt.ErrSignatureInvalid } return []byte(h.jwtSecret), nil }) if err != nil || !token.Valid { return "", errInvalidToken } claims, ok := token.Claims.(jwt.MapClaims) if !ok { return "", errInvalidClaims } userID, ok := claims["sub"].(string) if !ok { return "", errMissingUserID } return userID, nil } // buildMessage 构造推送 JSON 消息体 func buildMessage(event string, data map[string]any) ([]byte, error) { return json.Marshal(map[string]any{ "event": event, "data": data, }) }