package websocket import ( "encoding/json" "sync" "github.com/gorilla/websocket" ) // Client 单个 WS 连接 type Client struct { conn *websocket.Conn send chan []byte userID int64 isAdmin bool docks map[string]bool } // Hub 管理所有 WS 客户端 type Hub struct { clients map[*Client]bool mu sync.RWMutex } // DefaultHub 全局单例 var DefaultHub = &Hub{clients: make(map[*Client]bool)} // Message WS 推送消息信封 type Message struct { Type string `json:"type"` Data any `json:"data"` } // BroadcastToDockJSON delivers a sanitized resource event only to authorized clients. func (h *Hub) BroadcastToDockJSON(dockID, msgType string, data any) { b, err := json.Marshal(Message{Type: msgType, Data: data}) if err != nil { return } h.broadcastToDock(dockID, b) } func (h *Hub) broadcastToDock(dockID string, msg []byte) { h.mu.RLock() defer h.mu.RUnlock() for c := range h.clients { if !c.isAdmin && !c.docks[dockID] { continue } select { case c.send <- msg: default: } } } // BroadcastJSON 向所有客户端广播 JSON 消息 func (h *Hub) BroadcastJSON(msgType string, data any) { b, err := json.Marshal(Message{Type: msgType, Data: data}) if err != nil { return } h.broadcast(b) } func (h *Hub) broadcast(msg []byte) { h.mu.RLock() defer h.mu.RUnlock() for c := range h.clients { select { case c.send <- msg: default: // 缓冲满则丢弃,避免阻塞 } } } func (h *Hub) add(c *Client) { h.mu.Lock() h.clients[c] = true h.mu.Unlock() } func (h *Hub) remove(c *Client) { h.mu.Lock() if _, ok := h.clients[c]; ok { delete(h.clients, c) close(c.send) } h.mu.Unlock() } func (c *Client) writePump() { defer c.conn.Close() for msg := range c.send { if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil { return } } }