// Package websocket 提供WebSocket通信功能 // 实现中继设备和平台之间的双向通信 package websocket import ( "fmt" "net/http" "sync" "time" "git.huangwc.com/pig/pig-farm-controller/internal/logs" "git.huangwc.com/pig/pig-farm-controller/internal/service" "git.huangwc.com/pig/pig-farm-controller/internal/storage/repository" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" ) // Manager WebSocket管理器 type Manager struct { // websocketService WebSocket服务 websocketService *service.WebSocketService // logger 日志记录器 logger *logs.Logger // upgrader WebSocket升级器 upgrader websocket.Upgrader // mutex 互斥锁 mutex sync.RWMutex // connections 设备连接映射 connections map[string]*websocket.Conn // deviceRepo 设备仓库 deviceRepo repository.DeviceRepo } // NewManager 创建WebSocket管理器实例 func NewManager(websocketService *service.WebSocketService, deviceRepo repository.DeviceRepo) *Manager { return &Manager{ websocketService: websocketService, logger: logs.NewLogger(), upgrader: websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { // 允许所有跨域请求 return true }, }, connections: make(map[string]*websocket.Conn), deviceRepo: deviceRepo, } } // getDeviceDisplayName 获取设备显示名称 func (wm *Manager) getDeviceDisplayName(deviceID string) string { if wm.deviceRepo != nil { if device, err := wm.deviceRepo.FindByIDString(deviceID); err == nil && device != nil { return fmt.Sprintf("%s(id:%s)", device.Name, deviceID) } } return fmt.Sprintf("未知设备(id:%s)", deviceID) } // HandleConnection 处理WebSocket连接 func (wm *Manager) HandleConnection(c *gin.Context) { // 升级HTTP连接到WebSocket conn, err := wm.upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { wm.logger.Error("WebSocket连接升级失败: " + err.Error()) return } // 获取设备ID deviceID := c.Query("device_id") if deviceID == "" { wm.logger.Error("缺少设备ID参数") conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "缺少设备ID参数")) conn.Close() return } // 添加连接到映射 wm.mutex.Lock() wm.connections[deviceID] = conn wm.mutex.Unlock() deviceName := wm.getDeviceDisplayName(deviceID) wm.logger.Info("设备 " + deviceName + " 已连接") // 发送连接成功消息 successMsg := service.WebSocketMessage{ Type: "system", Command: "connected", Timestamp: time.Now(), } conn.WriteJSON(successMsg) // 处理消息循环 for { // 读取消息 messageType, message, err := conn.ReadMessage() if err != nil { wm.logger.Error("读取设备 " + deviceName + " 消息失败: " + err.Error()) break } // 只处理文本消息 if messageType != websocket.TextMessage { continue } // 处理设备消息 if err := wm.websocketService.HandleMessage(deviceID, message); err != nil { wm.logger.Error("处理设备 " + deviceName + " 消息失败: " + err.Error()) continue } } // 连接断开时清理 wm.mutex.Lock() delete(wm.connections, deviceID) wm.mutex.Unlock() conn.Close() wm.logger.Info("设备 " + deviceName + " 已断开连接") } // SendCommand 向指定设备发送指令 func (wm *Manager) SendCommand(deviceID, command string, data interface{}) error { wm.mutex.RLock() conn, exists := wm.connections[deviceID] wm.mutex.RUnlock() if !exists { return wm.websocketService.SendCommand(deviceID, command, data) } // 构造消息 msg := service.WebSocketMessage{ Type: service.MessageTypeCommand, Command: command, Data: data, Timestamp: time.Now(), } // 发送消息 if err := conn.WriteJSON(msg); err != nil { deviceName := wm.getDeviceDisplayName(deviceID) return fmt.Errorf("向设备 %s 发送指令失败: %v", deviceName, err) } return nil } // GetConnectedDevices 获取已连接的设备列表 func (wm *Manager) GetConnectedDevices() []string { wm.mutex.RLock() defer wm.mutex.RUnlock() devices := make([]string, 0, len(wm.connections)) for deviceID := range wm.connections { devices = append(devices, deviceID) } return devices }