laic-backend/websocket/hub.go

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
		}
	}
}