增加websocket支持

internal/websocket/manager.go

@@ -0,0 +1,149 @@
+// 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
+}