package socket import ( "context" "encoding/json" "fmt" "net/url" "strings" "time" "github.com/go-gost/x/config" "github.com/gorilla/websocket" "github.com/shirou/gopsutil/v3/cpu" "github.com/shirou/gopsutil/v3/host" "github.com/shirou/gopsutil/v3/mem" psnet "github.com/shirou/gopsutil/v3/net" ) // SystemInfo 系统信息结构体 type SystemInfo struct { Uptime uint64 `json:"uptime"` // 开机时间 (秒) BytesReceived uint64 `json:"bytes_received"` // 接收字节数 BytesTransmitted uint64 `json:"bytes_transmitted"` // 发送字节数 CPUUsage float64 `json:"cpu_usage"` // CPU使用率(百分比) MemoryUsage float64 `json:"memory_usage"` // 内存使用率(百分比) } // NetworkStats 网络统计信息 type NetworkStats struct { BytesReceived uint64 `json:"bytes_received"` // 接收字节数 BytesTransmitted uint64 `json:"bytes_transmitted"` // 发送字节数 } // CPUInfo CPU信息 type CPUInfo struct { Usage float64 `json:"usage"` // CPU使用率(百分比) } // MemoryInfo 内存信息 type MemoryInfo struct { Usage float64 `json:"usage"` // 内存使用率(百分比) } // CommandMessage 命令消息结构体 type CommandMessage struct { Type string `json:"type"` Data interface{} `json:"data"` RequestId string `json:"requestId,omitempty"` } // CommandResponse 命令响应结构体 type CommandResponse struct { Type string `json:"type"` Success bool `json:"success"` Message string `json:"message"` Data interface{} `json:"data,omitempty"` RequestId string `json:"requestId,omitempty"` } type WebSocketReporter struct { url string conn *websocket.Conn reconnectTime time.Duration pingInterval time.Duration configInterval time.Duration ctx context.Context cancel context.CancelFunc connected bool } // NewWebSocketReporter 创建一个新的WebSocket报告器 func NewWebSocketReporter(serverURL string) *WebSocketReporter { ctx, cancel := context.WithCancel(context.Background()) return &WebSocketReporter{ url: serverURL, reconnectTime: 5 * time.Second, // 重连间隔 pingInterval: 2 * time.Second, // 发送间隔改为2秒 configInterval: 10 * time.Minute, // 配置上报间隔 ctx: ctx, cancel: cancel, connected: false, } } // Start 启动WebSocket报告器 func (w *WebSocketReporter) Start() { go w.run() } // Stop 停止WebSocket报告器 func (w *WebSocketReporter) Stop() { w.cancel() if w.conn != nil { w.conn.Close() } } // run 主运行循环 func (w *WebSocketReporter) run() { for { select { case <-w.ctx.Done(): return default: if err := w.connect(); err != nil { fmt.Printf("❌ WebSocket连接失败: %v,%v后重试\n", err, w.reconnectTime) select { case <-time.After(w.reconnectTime): continue case <-w.ctx.Done(): return } } // 连接成功,开始发送消息 w.handleConnection() } } } // connect 建立WebSocket连接 func (w *WebSocketReporter) connect() error { u, err := url.Parse(w.url) if err != nil { return fmt.Errorf("解析URL失败: %v", err) } dialer := websocket.DefaultDialer dialer.HandshakeTimeout = 10 * time.Second conn, _, err := dialer.Dial(u.String(), nil) if err != nil { return fmt.Errorf("连接WebSocket失败: %v", err) } w.conn = conn w.connected = true // 设置关闭处理器来检测连接状态 w.conn.SetCloseHandler(func(code int, text string) error { w.connected = false return nil }) return nil } // handleConnection 处理WebSocket连接 func (w *WebSocketReporter) handleConnection() { defer func() { if w.conn != nil { w.conn.Close() w.conn = nil } w.connected = false }() // 启动消息接收goroutine go w.receiveMessages() // 启动配置上报goroutine go w.reportConfig() // 主发送循环 ticker := time.NewTicker(w.pingInterval) defer ticker.Stop() for { select { case <-w.ctx.Done(): return case <-ticker.C: // 检查连接状态 if !w.connected { return } // 获取系统信息并发送 sysInfo := w.collectSystemInfo() if err := w.sendSystemInfo(sysInfo); err != nil { fmt.Printf("❌ 发送系统信息失败: %v,准备重连\n", err) return } } } } // collectSystemInfo 收集系统信息 func (w *WebSocketReporter) collectSystemInfo() SystemInfo { networkStats := getNetworkStats() cpuInfo := getCPUInfo() memoryInfo := getMemoryInfo() return SystemInfo{ Uptime: getUptime(), BytesReceived: networkStats.BytesReceived, BytesTransmitted: networkStats.BytesTransmitted, CPUUsage: cpuInfo.Usage, MemoryUsage: memoryInfo.Usage, } } // sendSystemInfo 发送系统信息 func (w *WebSocketReporter) sendSystemInfo(sysInfo SystemInfo) error { if w.conn == nil || !w.connected { return fmt.Errorf("连接未建立") } // 转换为JSON jsonData, err := json.Marshal(sysInfo) if err != nil { return fmt.Errorf("序列化系统信息失败: %v", err) } // 设置写入超时 w.conn.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err := w.conn.WriteMessage(websocket.TextMessage, jsonData); err != nil { w.connected = false // 标记连接已断开 return fmt.Errorf("写入消息失败: %v", err) } return nil } // receiveMessages 接收服务端发送的消息 func (w *WebSocketReporter) receiveMessages() { for { select { case <-w.ctx.Done(): return default: if w.conn == nil || !w.connected { return } // 设置读取超时 w.conn.SetReadDeadline(time.Now().Add(30 * time.Second)) messageType, message, err := w.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { fmt.Printf("❌ WebSocket读取消息错误: %v\n", err) } w.connected = false return } // 处理接收到的消息 w.handleReceivedMessage(messageType, message) } } } // handleReceivedMessage 处理接收到的消息 func (w *WebSocketReporter) handleReceivedMessage(messageType int, message []byte) { switch messageType { case websocket.TextMessage: // 解析命令消息 var cmdMsg CommandMessage if err := json.Unmarshal(message, &cmdMsg); err != nil { fmt.Printf("❌ 解析命令消息失败: %v\n", err) w.sendErrorResponse("ParseError", fmt.Sprintf("解析命令失败: %v", err)) return } if cmdMsg.Type != "call" { w.routeCommand(cmdMsg) } default: fmt.Printf("📨 收到未知类型消息: %d\n", messageType) } } // routeCommand 路由命令到对应的处理函数 func (w *WebSocketReporter) routeCommand(cmd CommandMessage) { var err error var response CommandResponse // 传递 requestId response.RequestId = cmd.RequestId switch cmd.Type { // Service 相关命令 case "AddService": err = w.handleAddService(cmd.Data) response.Type = "AddServiceResponse" case "UpdateService": err = w.handleUpdateService(cmd.Data) response.Type = "UpdateServiceResponse" case "DeleteService": err = w.handleDeleteService(cmd.Data) response.Type = "DeleteServiceResponse" case "PauseService": err = w.handlePauseService(cmd.Data) response.Type = "PauseServiceResponse" case "ResumeService": err = w.handleResumeService(cmd.Data) response.Type = "ResumeServiceResponse" // Chain 相关命令 case "AddChains": err = w.handleAddChain(cmd.Data) response.Type = "AddChainsResponse" case "UpdateChains": err = w.handleUpdateChain(cmd.Data) response.Type = "UpdateChainsResponse" case "DeleteChains": err = w.handleDeleteChain(cmd.Data) response.Type = "DeleteChainsResponse" // Limiter 相关命令 case "AddLimiters": err = w.handleAddLimiter(cmd.Data) response.Type = "AddLimitersResponse" case "UpdateLimiters": err = w.handleUpdateLimiter(cmd.Data) response.Type = "UpdateLimitersResponse" case "DeleteLimiters": err = w.handleDeleteLimiter(cmd.Data) response.Type = "DeleteLimitersResponse" default: err = fmt.Errorf("未知命令类型: %s", cmd.Type) response.Type = "UnknownCommandResponse" } // 发送响应 if err != nil { saveConfig() response.Success = false response.Message = err.Error() } else { saveConfig() response.Success = true response.Message = "OK" } w.sendResponse(response) } // Service 命令处理函数 func (w *WebSocketReporter) handleAddService(data interface{}) error { // 将 interface{} 转换为 JSON 再解析为具体类型 jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var services []config.ServiceConfig if err := json.Unmarshal(jsonData, &services); err != nil { return fmt.Errorf("解析服务配置失败: %v", err) } req := createServicesRequest{Data: services} return createServices(req) } func (w *WebSocketReporter) handleUpdateService(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var services []config.ServiceConfig if err := json.Unmarshal(jsonData, &services); err != nil { return fmt.Errorf("解析服务配置失败: %v", err) } req := updateServicesRequest{Data: services} return updateServices(req) } func (w *WebSocketReporter) handleDeleteService(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var req deleteServicesRequest if err := json.Unmarshal(jsonData, &req); err != nil { return fmt.Errorf("解析删除请求失败: %v", err) } return deleteServices(req) } func (w *WebSocketReporter) handlePauseService(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var req pauseServicesRequest if err := json.Unmarshal(jsonData, &req); err != nil { return fmt.Errorf("解析暂停请求失败: %v", err) } return pauseServices(req) } func (w *WebSocketReporter) handleResumeService(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var req resumeServicesRequest if err := json.Unmarshal(jsonData, &req); err != nil { return fmt.Errorf("解析恢复请求失败: %v", err) } return resumeServices(req) } // Chain 命令处理函数 func (w *WebSocketReporter) handleAddChain(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var chainConfig config.ChainConfig if err := json.Unmarshal(jsonData, &chainConfig); err != nil { return fmt.Errorf("解析链配置失败: %v", err) } req := createChainRequest{Data: chainConfig} return createChain(req) } func (w *WebSocketReporter) handleUpdateChain(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } // 对于更新操作,Java端发送的格式可能是: {"chain": "name", "data": {...}} var updateReq struct { Chain string `json:"chain"` Data config.ChainConfig `json:"data"` } // 尝试解析为更新请求格式 if err := json.Unmarshal(jsonData, &updateReq); err != nil { // 如果失败,可能是直接的ChainConfig,从name字段获取chain名称 var chainConfig config.ChainConfig if err := json.Unmarshal(jsonData, &chainConfig); err != nil { return fmt.Errorf("解析链配置失败: %v", err) } updateReq.Chain = chainConfig.Name updateReq.Data = chainConfig } req := updateChainRequest{ Chain: updateReq.Chain, Data: updateReq.Data, } return updateChain(req) } func (w *WebSocketReporter) handleDeleteChain(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } // 删除操作可能是: {"chain": "name"} 或者直接是链名称字符串 var deleteReq deleteChainRequest // 尝试解析为删除请求格式 if err := json.Unmarshal(jsonData, &deleteReq); err != nil { // 如果失败,可能是字符串格式的名称 var chainName string if err := json.Unmarshal(jsonData, &chainName); err != nil { return fmt.Errorf("解析链删除请求失败: %v", err) } deleteReq.Chain = chainName } return deleteChain(deleteReq) } // Limiter 命令处理函数 func (w *WebSocketReporter) handleAddLimiter(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } var limiterConfig config.LimiterConfig if err := json.Unmarshal(jsonData, &limiterConfig); err != nil { return fmt.Errorf("解析限流器配置失败: %v", err) } req := createLimiterRequest{Data: limiterConfig} return createLimiter(req) } func (w *WebSocketReporter) handleUpdateLimiter(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } // 对于更新操作,Java端发送的格式可能是: {"limiter": "name", "data": {...}} var updateReq struct { Limiter string `json:"limiter"` Data config.LimiterConfig `json:"data"` } // 尝试解析为更新请求格式 if err := json.Unmarshal(jsonData, &updateReq); err != nil { // 如果失败,可能是直接的LimiterConfig,从name字段获取limiter名称 var limiterConfig config.LimiterConfig if err := json.Unmarshal(jsonData, &limiterConfig); err != nil { return fmt.Errorf("解析限流器配置失败: %v", err) } updateReq.Limiter = limiterConfig.Name updateReq.Data = limiterConfig } req := updateLimiterRequest{ Limiter: updateReq.Limiter, Data: updateReq.Data, } return updateLimiter(req) } func (w *WebSocketReporter) handleDeleteLimiter(data interface{}) error { jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化数据失败: %v", err) } // 删除操作可能是: {"limiter": "name"} 或者直接是限流器名称字符串 var deleteReq deleteLimiterRequest // 尝试解析为删除请求格式 if err := json.Unmarshal(jsonData, &deleteReq); err != nil { // 如果失败,可能是字符串格式的名称 var limiterName string if err := json.Unmarshal(jsonData, &limiterName); err != nil { return fmt.Errorf("解析限流器删除请求失败: %v", err) } deleteReq.Limiter = limiterName } return deleteLimiter(deleteReq) } // handleCall 处理服务端的call回调消息 func (w *WebSocketReporter) handleCall(data interface{}) error { // 解析call数据 jsonData, err := json.Marshal(data) if err != nil { return fmt.Errorf("序列化call数据失败: %v", err) } // 可以根据call的具体内容进行不同的处理 var callData map[string]interface{} if err := json.Unmarshal(jsonData, &callData); err != nil { return fmt.Errorf("解析call数据失败: %v", err) } fmt.Printf("🔔 收到服务端call回调: %v\n", callData) // 根据call的类型执行不同的操作 if callType, exists := callData["type"]; exists { switch callType { case "ping": fmt.Printf("📡 收到ping,发送pong回应\n") // 可以在这里发送pong响应 case "info_request": fmt.Printf("📊 服务端请求额外信息\n") // 可以在这里发送额外的系统信息 case "command": fmt.Printf("⚡ 服务端发送执行命令\n") // 可以在这里执行特定命令 default: fmt.Printf("❓ 未知的call类型: %v\n", callType) } } // 简单返回成功,表示call已被处理 return nil } // reportConfig 定时上报配置信息 func (w *WebSocketReporter) reportConfig() { // 立即发送一次配置 w.sendConfigReport() // 启动定时器 ticker := time.NewTicker(w.configInterval) defer ticker.Stop() for { select { case <-w.ctx.Done(): return case <-ticker.C: if w.connected { w.sendConfigReport() } } } } // sendConfigReport 发送配置报告 func (w *WebSocketReporter) sendConfigReport() { if w.conn == nil || !w.connected { return } // 获取配置数据 configData, err := getConfig() if err != nil { fmt.Printf("❌ 获取配置失败: %v\n", err) return } // 构建配置报告消息 configMsg := struct { Type string `json:"type"` Data interface{} `json:"data"` }{ Type: "config_report", Data: json.RawMessage(configData), } // 转换为JSON jsonData, err := json.Marshal(configMsg) if err != nil { fmt.Printf("❌ 序列化配置报告失败: %v\n", err) return } // 设置写入超时 w.conn.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err := w.conn.WriteMessage(websocket.TextMessage, jsonData); err != nil { fmt.Printf("❌ 发送配置报告失败: %v\n", err) w.connected = false return } } // sendResponse 发送响应消息到服务端 func (w *WebSocketReporter) sendResponse(response CommandResponse) { if w.conn == nil || !w.connected { fmt.Printf("❌ 无法发送响应:连接未建立\n") return } jsonData, err := json.Marshal(response) if err != nil { fmt.Printf("❌ 序列化响应失败: %v\n", err) return } w.conn.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err := w.conn.WriteMessage(websocket.TextMessage, jsonData); err != nil { fmt.Printf("❌ 发送响应失败: %v\n", err) w.connected = false } } // sendErrorResponse 发送错误响应 func (w *WebSocketReporter) sendErrorResponse(responseType, message string) { response := CommandResponse{ Type: responseType, Success: false, Message: message, } w.sendResponse(response) } // getUptime 获取系统开机时间(秒) func getUptime() uint64 { uptime, err := host.Uptime() if err != nil { return 0 } return uptime } // getNetworkStats 获取网络统计信息 func getNetworkStats() NetworkStats { var stats NetworkStats ioCounters, err := psnet.IOCounters(true) if err != nil { fmt.Printf("获取网络统计失败: %v\n", err) return stats } // 汇总所有非回环接口的流量 for _, io := range ioCounters { // 跳过回环接口 if io.Name == "lo" || strings.HasPrefix(io.Name, "lo") { continue } stats.BytesReceived += io.BytesRecv stats.BytesTransmitted += io.BytesSent } return stats } // getCPUInfo 获取CPU信息 func getCPUInfo() CPUInfo { var cpuInfo CPUInfo // 获取CPU使用率 percentages, err := cpu.Percent(time.Second, false) if err == nil && len(percentages) > 0 { cpuInfo.Usage = percentages[0] } return cpuInfo } // getMemoryInfo 获取内存信息 func getMemoryInfo() MemoryInfo { var memInfo MemoryInfo vmStat, err := mem.VirtualMemory() if err != nil { return memInfo } memInfo.Usage = vmStat.UsedPercent return memInfo } // StartWebSocketReporterWithConfig 使用配置启动WebSocket报告器 func StartWebSocketReporterWithConfig(Addr string, Secret string) *WebSocketReporter { // 构建包含本机IP的WebSocket URL var fullURL = "ws://" + Addr + "/system-info?type=1&secret=" + Secret fmt.Printf("🔗 WebSocket连接URL: %s\n", fullURL) reporter := NewWebSocketReporter(fullURL) reporter.Start() return reporter }