diff --git a/services/push-gateway/internal/config/config.go b/services/push-gateway/internal/config/config.go index 17e627f..f27f35d 100644 --- a/services/push-gateway/internal/config/config.go +++ b/services/push-gateway/internal/config/config.go @@ -2,18 +2,25 @@ package config import "os" +// Config 持有 push-gateway 运行时配置 type Config struct { Port string JWTSecret string + DevMode bool + RedisURL string } +// Load 从环境变量加载配置并提供默认值 func Load() *Config { return &Config{ Port: getEnv("PUSH_GATEWAY_PORT", "8081"), JWTSecret: getEnv("JWT_SECRET", "p1-dev-secret-change-in-production"), + DevMode: getEnv("DEV_MODE", "false") == "true", + RedisURL: getEnv("REDIS_URL", ""), } } +// getEnv 读取环境变量,缺失时返回 fallback func getEnv(key, fallback string) string { if v := os.Getenv(key); v != "" { return v diff --git a/services/push-gateway/internal/hub/hub.go b/services/push-gateway/internal/hub/hub.go index 1b9bb67..08da7c0 100644 --- a/services/push-gateway/internal/hub/hub.go +++ b/services/push-gateway/internal/hub/hub.go @@ -6,54 +6,94 @@ import ( "github.com/gorilla/websocket" ) -// Hub 管理 WebSocket 客户端连接 +// Connection 包装单个 WebSocket 连接及其异步发送通道 +type Connection struct { + UserID string + conn *websocket.Conn + send chan []byte +} + +// Send 将消息投递到该连接的发送通道(非阻塞,通道满则丢弃) +func (c *Connection) Send(message []byte) { + select { + case c.send <- message: + default: + // 通道满,丢弃消息避免阻塞 hub + } +} + +// Outgoing 返回该连接的发送通道,供写协程消费 +func (c *Connection) Outgoing() <-chan []byte { + return c.send +} + +// Hub 管理所有在线 WebSocket 连接,按 userID 索引 type Hub struct { mu sync.RWMutex - clients map[string]map[*websocket.Conn]bool // userID -> connections + clients map[string]map[*Connection]bool } +// NewHub 创建 Hub 实例 func NewHub() *Hub { return &Hub{ - clients: make(map[string]map[*websocket.Conn]bool), + clients: make(map[string]map[*Connection]bool), } } -func (h *Hub) Register(userID string, conn *websocket.Conn) { +// Register 注册一个用户连接,返回 Connection 供调用方持有 +func (h *Hub) Register(userID string, conn *websocket.Conn) *Connection { + c := &Connection{ + UserID: userID, + conn: conn, + send: make(chan []byte, 64), + } h.mu.Lock() defer h.mu.Unlock() - if h.clients[userID] == nil { - h.clients[userID] = make(map[*websocket.Conn]bool) + h.clients[userID] = make(map[*Connection]bool) } - h.clients[userID][conn] = true + h.clients[userID][c] = true + return c } -func (h *Hub) Unregister(userID string, conn *websocket.Conn) { +// Unregister 注销一个用户连接并关闭其发送通道 +func (h *Hub) Unregister(c *Connection) { h.mu.Lock() defer h.mu.Unlock() - - if conns, ok := h.clients[userID]; ok { - delete(conns, conn) - if len(conns) == 0 { - delete(h.clients, userID) - } + conns, ok := h.clients[c.UserID] + if !ok { + return + } + if _, exists := conns[c]; exists { + delete(conns, c) + close(c.send) + } + if len(conns) == 0 { + delete(h.clients, c.UserID) } } -// SendToUser 向指定用户的所有连接推送消息 +// SendToUser 向指定用户的所有在线连接推送消息 func (h *Hub) SendToUser(userID string, message []byte) error { h.mu.RLock() defer h.mu.RUnlock() - conns, ok := h.clients[userID] if !ok { return nil } - - for conn := range conns { - if err := conn.WriteMessage(websocket.TextMessage, message); err != nil { - return err - } + for c := range conns { + c.Send(message) } return nil } + +// Broadcast 向所有在线连接广播消息 +func (h *Hub) Broadcast(message []byte) { + h.mu.RLock() + defer h.mu.RUnlock() + for _, conns := range h.clients { + for c := range conns { + c.Send(message) + } + } +} diff --git a/services/push-gateway/internal/ws/handler.go b/services/push-gateway/internal/ws/handler.go index e04267b..7f005a1 100644 --- a/services/push-gateway/internal/ws/handler.go +++ b/services/push-gateway/internal/ws/handler.go @@ -1,6 +1,8 @@ package ws import ( + "encoding/json" + "errors" "net/http" "strings" @@ -10,50 +12,36 @@ import ( "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 }, } +// Handler 处理 WebSocket 升级与内部推送 API type Handler struct { hub *hub.Hub jwtSecret string + devMode bool } -func NewHandler(h *hub.Hub, jwtSecret string) *Handler { - return &Handler{hub: h, jwtSecret: jwtSecret} +// NewHandler 创建 Handler 实例 +func NewHandler(h *hub.Hub, jwtSecret string, devMode bool) *Handler { + return &Handler{hub: h, jwtSecret: jwtSecret, devMode: devMode} } +// HandleWebSocket 升级 HTTP 为 WebSocket,保持长连接并处理心跳 func (h *Handler) HandleWebSocket(c *gin.Context) { - // 从 query 参数获取 token(WebSocket 无法设置 Authorization 头) - tokenStr := c.Query("token") - if tokenStr == "" { - c.JSON(http.StatusUnauthorized, gin.H{"error": "missing token"}) - return - } - - token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (interface{}, error) { - if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok { - return nil, jwt.ErrSignatureInvalid - } - return []byte(h.jwtSecret), nil - }) - - if err != nil || !token.Valid { - c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"}) - return - } - - claims, ok := token.Claims.(jwt.MapClaims) - if !ok { - c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid claims"}) - return - } - - userID, ok := claims["sub"].(string) - if !ok { - c.JSON(http.StatusUnauthorized, gin.H{"error": "missing user id"}) + userID, err := h.authenticate(c) + if err != nil { + c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()}) return } @@ -63,8 +51,17 @@ func (h *Handler) HandleWebSocket(c *gin.Context) { } defer conn.Close() - h.hub.Register(userID, conn) - defer h.hub.Unregister(userID, conn) + 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 { + return + } + } + }() // 读取循环(保持连接,处理心跳) for { @@ -72,28 +69,102 @@ func (h *Handler) HandleWebSocket(c *gin.Context) { if err != nil { break } - // 处理心跳 ping if strings.ToLower(string(msg)) == "ping" { - conn.WriteMessage(websocket.TextMessage, []byte("pong")) + connPtr.Send([]byte("pong")) } } } -// PushHandler 接收来自 Msg 服务的推送请求 +// PushHandler 接收来自 Msg 服务的定向推送请求 func (h *Handler) PushHandler(c *gin.Context) { var req struct { - UserID string `json:"user_id"` - Message string `json:"message"` + 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 } - if err := h.hub.SendToUser(req.UserID, []byte(req.Message)); err != nil { + 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) { + 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}) +} + +// 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, + }) +} diff --git a/services/push-gateway/main.go b/services/push-gateway/main.go index fae9104..9d44948 100644 --- a/services/push-gateway/main.go +++ b/services/push-gateway/main.go @@ -20,7 +20,7 @@ func main() { gin.SetMode(gin.ReleaseMode) h := hub.NewHub() - wsHandler := ws.NewHandler(h, cfg.JWTSecret) + wsHandler := ws.NewHandler(h, cfg.JWTSecret, cfg.DevMode) r := gin.New() r.Use(gin.Recovery()) @@ -35,6 +35,7 @@ func main() { // 内部推送 API(Msg 服务调用) api := r.Group("/internal") api.POST("/push", wsHandler.PushHandler) + api.POST("/broadcast", wsHandler.BroadcastHandler) srv := &http.Server{ Addr: ":" + cfg.Port,