增加websocket支持

internal/core/application.go

@@ -9,6 +9,7 @@ import (
 	"git.huangwc.com/pig/pig-farm-controller/internal/api"
 	"git.huangwc.com/pig/pig-farm-controller/internal/config"
 	"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/db"
 	"git.huangwc.com/pig/pig-farm-controller/internal/storage/repository"
 	"git.huangwc.com/pig/pig-farm-controller/internal/task"
@@ -38,6 +39,9 @@ type Application struct {
 	// DeviceRepo 设备仓库实例
 	DeviceRepo repository.DeviceRepo
 
+	// WebSocketService WebSocket服务实例
+	WebSocketService *service.WebSocketService
+
 	// Config 应用配置
 	Config *config.Config
 
@@ -71,8 +75,11 @@ func NewApplication(cfg *config.Config) *Application {
 	// 初始化设备仓库
 	deviceRepo := repository.NewDeviceRepo(store.GetDB())
 
+	// 初始化WebSocket服务
+	websocketService := service.NewWebSocketService()
+
 	// 初始化API组件
-	apiInstance := api.NewAPI(cfg, userRepo, operationHistoryRepo, deviceControlRepo, deviceRepo)
+	apiInstance := api.NewAPI(cfg, userRepo, operationHistoryRepo, deviceControlRepo, deviceRepo, websocketService)
 
 	// 初始化任务执行器组件(使用5个工作协程)
 	taskExecutor := task.NewExecutor(5)
@@ -85,6 +92,7 @@ func NewApplication(cfg *config.Config) *Application {
 		OperationHistoryRepo: operationHistoryRepo,
 		DeviceControlRepo:    deviceControlRepo,
 		DeviceRepo:           deviceRepo,
+		WebSocketService:     websocketService,
 		Config:               cfg,
 		logger:               logs.NewLogger(),
 	}

internal/core/websocket.go

@@ -0,0 +1,174 @@
+// Package core 提供WebSocket服务功能
+// 实现中继设备和平台之间的双向通信
+package core
+
+import (
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	"git.huangwc.com/pig/pig-farm-controller/internal/logs"
+	"git.huangwc.com/pig/pig-farm-controller/internal/model"
+	"github.com/gorilla/websocket"
+)
+
+// WebSocket消息类型常量
+const (
+	// MessageTypeCommand 平台向设备发送的指令
+	MessageTypeCommand = "command"
+
+	// MessageTypeResponse 设备向平台发送的响应
+	MessageTypeResponse = "response"
+
+	// MessageTypeHeartbeat 心跳消息
+	MessageTypeHeartbeat = "heartbeat"
+)
+
+// WebSocketMessage WebSocket消息结构
+type WebSocketMessage struct {
+	// Type 消息类型
+	Type string `json:"type"`
+
+	// DeviceID 设备ID
+	DeviceID string `json:"device_id,omitempty"`
+
+	// Command 指令内容
+	Command string `json:"command,omitempty"`
+
+	// Data 消息数据
+	Data interface{} `json:"data,omitempty"`
+
+	// Timestamp 时间戳
+	Timestamp time.Time `json:"timestamp"`
+}
+
+// DeviceConnection 设备连接信息
+type DeviceConnection struct {
+	// DeviceID 设备ID
+	DeviceID string
+
+	// Connection WebSocket连接
+	Connection *websocket.Conn
+
+	// LastHeartbeat 最后心跳时间
+	LastHeartbeat time.Time
+
+	// DeviceInfo 设备信息
+	DeviceInfo *model.Device
+}
+
+// WebSocketService WebSocket服务
+type WebSocketService struct {
+	// connections 设备连接映射
+	connections map[string]*DeviceConnection
+
+	// mutex 互斥锁
+	mutex sync.RWMutex
+
+	// logger 日志记录器
+	logger *logs.Logger
+}
+
+// NewWebSocketService 创建WebSocket服务实例
+func NewWebSocketService() *WebSocketService {
+	return &WebSocketService{
+		connections: make(map[string]*DeviceConnection),
+		logger:      logs.NewLogger(),
+	}
+}
+
+// AddConnection 添加设备连接
+func (ws *WebSocketService) AddConnection(deviceID string, conn *websocket.Conn) {
+	ws.mutex.Lock()
+	defer ws.mutex.Unlock()
+
+	ws.connections[deviceID] = &DeviceConnection{
+		DeviceID:      deviceID,
+		Connection:    conn,
+		LastHeartbeat: time.Now(),
+	}
+
+	ws.logger.Info(fmt.Sprintf("设备 %s 已连接", deviceID))
+}
+
+// RemoveConnection 移除设备连接
+func (ws *WebSocketService) RemoveConnection(deviceID string) {
+	ws.mutex.Lock()
+	defer ws.mutex.Unlock()
+
+	delete(ws.connections, deviceID)
+
+	ws.logger.Info(fmt.Sprintf("设备 %s 已断开连接", deviceID))
+}
+
+// SendCommand 向指定设备发送指令
+func (ws *WebSocketService) SendCommand(deviceID, command string, data interface{}) error {
+	ws.mutex.RLock()
+	deviceConn, exists := ws.connections[deviceID]
+	ws.mutex.RUnlock()
+
+	if !exists {
+		return fmt.Errorf("设备 %s 未连接", deviceID)
+	}
+
+	// 构造消息
+	msg := WebSocketMessage{
+		Type:      MessageTypeCommand,
+		Command:   command,
+		Data:      data,
+		Timestamp: time.Now(),
+	}
+
+	// 发送消息
+	if err := deviceConn.Connection.WriteJSON(msg); err != nil {
+		return fmt.Errorf("向设备 %s 发送指令失败: %v", deviceID, err)
+	}
+
+	return nil
+}
+
+// GetConnectedDevices 获取已连接的设备列表
+func (ws *WebSocketService) GetConnectedDevices() []string {
+	ws.mutex.RLock()
+	defer ws.mutex.RUnlock()
+
+	devices := make([]string, 0, len(ws.connections))
+	for deviceID := range ws.connections {
+		devices = append(devices, deviceID)
+	}
+
+	return devices
+}
+
+// HandleMessage 处理来自设备的消息
+func (ws *WebSocketService) HandleMessage(deviceID string, message []byte) error {
+	// 解析消息
+	var msg WebSocketMessage
+	if err := json.Unmarshal(message, &msg); err != nil {
+		return fmt.Errorf("解析设备 %s 消息失败: %v", deviceID, err)
+	}
+
+	// 更新心跳时间
+	if msg.Type == MessageTypeHeartbeat {
+		ws.mutex.Lock()
+		if deviceConn, exists := ws.connections[deviceID]; exists {
+			deviceConn.LastHeartbeat = time.Now()
+		}
+		ws.mutex.Unlock()
+	}
+
+	// 记录消息日志
+	ws.logger.Info(fmt.Sprintf("收到来自设备 %s 的消息: %v", deviceID, msg))
+
+	return nil
+}
+
+// GetDeviceConnection 获取设备连接信息
+func (ws *WebSocketService) GetDeviceConnection(deviceID string) (*DeviceConnection, bool) {
+	ws.mutex.RLock()
+	defer ws.mutex.RUnlock()
+
+	deviceConn, exists := ws.connections[deviceID]
+	return deviceConn, exists
+}