laic-backend/websocket/handler.go

package websocket

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	"laic-backend/common"
	"laic-backend/logger"
	"laic-backend/model"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		return origin == "" || origin == "http://localhost:5173" || origin == "http://127.0.0.1:5173"
	},
}

// HandleMonitor 前端实时监控 WS 端点（设备状态 + 告警推送）
func HandleMonitor(c *gin.Context) {
	userID := common.GetUserId(c)
	isAdmin := common.IsAdmin(c)
	if userID == 0 {
		c.Status(http.StatusUnauthorized)
		return
	}
	docks := make(map[string]bool)
	if !isAdmin {
		var owned []model.Dock
		if err := common.DB.Select("dock_id").Where("user_id = ?", userID).Find(&owned).Error; err != nil {
			c.Status(http.StatusInternalServerError)
			return
		}
		for _, dock := range owned {
			docks[dock.DockID] = true
		}
	}
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.ERROR("WS 升级失败", err)
		return
	}
	client := &Client{conn: conn, send: make(chan []byte, 256), userID: userID, isAdmin: isAdmin, docks: docks}
	DefaultHub.add(client)
	defer DefaultHub.remove(client)

	go client.writePump()
	client.readPump()
}

// readPump 阻塞读取以探测断开，出错即退出
func (c *Client) readPump() {
	_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	})
	for {
		if _, _, err := c.conn.ReadMessage(); err != nil {
			return
		}
	}
}