150 lines
3.4 KiB
Go
150 lines
3.4 KiB
Go
// Package websocket 提供WebSocket通信功能
|
|
// 实现中继设备和平台之间的双向通信
|
|
package websocket
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.huangwc.com/pig/pig-farm-controller/internal/logs"
|
|
"git.huangwc.com/pig/pig-farm-controller/internal/service"
|
|
"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
|
|
}
|
|
|
|
// NewManager 创建WebSocket管理器实例
|
|
func NewManager(websocketService *service.WebSocketService) *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),
|
|
}
|
|
}
|
|
|
|
// 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()
|
|
|
|
wm.logger.Info("设备 " + deviceID + " 已连接")
|
|
|
|
// 发送连接成功消息
|
|
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("读取设备 " + deviceID + " 消息失败: " + err.Error())
|
|
break
|
|
}
|
|
|
|
// 只处理文本消息
|
|
if messageType != websocket.TextMessage {
|
|
continue
|
|
}
|
|
|
|
// 处理设备消息
|
|
if err := wm.websocketService.HandleMessage(deviceID, message); err != nil {
|
|
wm.logger.Error("处理设备 " + deviceID + " 消息失败: " + err.Error())
|
|
continue
|
|
}
|
|
}
|
|
|
|
// 连接断开时清理
|
|
wm.mutex.Lock()
|
|
delete(wm.connections, deviceID)
|
|
wm.mutex.Unlock()
|
|
|
|
conn.Close()
|
|
wm.logger.Info("设备 " + deviceID + " 已断开连接")
|
|
}
|
|
|
|
// 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 {
|
|
return 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
|
|
}
|