package api import ( "fmt" "net" "net/http" "net/url" "sync" "time" "webhelper/global" "github.com/gorilla/websocket" ) // WebSocketClient WebSocket 客户端 type WebSocketClient struct { conn *websocket.Conn isConnected bool mu sync.RWMutex onMessage func(message string) onError func(err error) onClose func() } // WebSocketConfig WebSocket 配置 type WebSocketConfig struct { URL string `json:"url"` Headers map[string]string `json:"headers"` Insecure bool `json:"insecure"` } // WebSocketMessage WebSocket 消息 type WebSocketMessage struct { Type string `json:"type"` // text, binary, ping, pong, error, system Content string `json:"content"` Time string `json:"time"` } var wsClient *WebSocketClient // NewWebSocketClient 创建 WebSocket 客户端 func NewWebSocketClient() *WebSocketClient { return &WebSocketClient{ isConnected: false, } } // Connect 连接 WebSocket func (c *WebSocketClient) Connect(config *WebSocketConfig, onMessage func(string), onError func(error), onClose func()) error { c.mu.Lock() defer c.mu.Unlock() if c.isConnected { return fmt.Errorf("已连接") } // 解析 URL u, err := url.Parse(config.URL) if err != nil { return fmt.Errorf("URL 格式错误: %v", err) } // 创建 HTTP Header(只添加用户自定义头) header := http.Header{} for key, value := range config.Headers { header.Set(key, value) } // 创建 Dialer(最简配置) dialer := &websocket.Dialer{ HandshakeTimeout: 10 * time.Second, } // 连接 global.Log.Infof("正在连接 WebSocket: %s", u.String()) conn, _, err := dialer.Dial(u.String(), header) if err != nil { return fmt.Errorf("连接失败: %v", err) } c.conn = conn c.isConnected = true c.onMessage = onMessage c.onError = onError c.onClose = onClose global.Log.Infof("WebSocket 连接成功: %s", u.String()) // 启动消息监听 go c.readMessages() return nil } // readMessages 读取消息 func (c *WebSocketClient) readMessages() { defer func() { c.mu.Lock() c.isConnected = false c.mu.Unlock() if c.onClose != nil { c.onClose() } }() for { _, message, err := c.conn.ReadMessage() if err != nil { if c.onError != nil { c.onError(err) } return } if c.onMessage != nil { c.onMessage(string(message)) } } } // Send 发送消息 func (c *WebSocketClient) Send(message string) error { c.mu.RLock() defer c.mu.RUnlock() if !c.isConnected || c.conn == nil { return fmt.Errorf("未连接") } return c.conn.WriteMessage(websocket.TextMessage, []byte(message)) } // Close 关闭连接 func (c *WebSocketClient) Close() error { c.mu.Lock() defer c.mu.Unlock() if c.conn != nil { err := c.conn.Close() c.isConnected = false return err } return nil } // IsConnected 是否已连接 func (c *WebSocketClient) IsConnected() bool { c.mu.RLock() defer c.mu.RUnlock() return c.isConnected } // GetWebSocketClient 获取 WebSocket 客户端 func GetWebSocketClient() *WebSocketClient { if wsClient == nil { wsClient = NewWebSocketClient() } return wsClient } // ConnectWebSocket 连接 WebSocket func ConnectWebSocket(config *WebSocketConfig, onMessage func(string), onError func(error), onClose func()) error { client := GetWebSocketClient() return client.Connect(config, onMessage, onError, onClose) } // SendWebSocketMessage 发送 WebSocket 消息 func SendWebSocketMessage(message string) error { client := GetWebSocketClient() return client.Send(message) } // CloseWebSocket 关闭 WebSocket 连接 func CloseWebSocket() error { client := GetWebSocketClient() return client.Close() } // IsWebSocketConnected WebSocket 是否已连接 func IsWebSocketConnected() bool { client := GetWebSocketClient() return client.IsConnected() } // TCPClient TCP 客户端 type TCPClient struct { conn net.Conn isConnected bool mu sync.RWMutex onMessage func(message string) onError func(err error) onClose func() } // TCPConfig TCP 配置 type TCPConfig struct { Host string `json:"host"` Port int `json:"port"` } // TCPMessage TCP 消息 type TCPMessage struct { Type string `json:"type"` Content string `json:"content"` Time string `json:"time"` } var tcpClient *TCPClient // NewTCPClient 创建 TCP 客户端 func NewTCPClient() *TCPClient { return &TCPClient{ isConnected: false, } } // Connect 连接 TCP func (c *TCPClient) Connect(config *TCPConfig, onMessage func(string), onError func(error), onClose func()) error { c.mu.Lock() defer c.mu.Unlock() if c.isConnected { return fmt.Errorf("已连接") } addr := fmt.Sprintf("%s:%d", config.Host, config.Port) conn, err := net.DialTimeout("tcp", addr, 10*time.Second) if err != nil { return fmt.Errorf("连接失败: %v", err) } c.conn = conn c.isConnected = true c.onMessage = onMessage c.onError = onError c.onClose = onClose // 启动消息监听 go c.readMessages() return nil } // readMessages 读取消息 func (c *TCPClient) readMessages() { defer func() { c.mu.Lock() c.isConnected = false c.mu.Unlock() if c.onClose != nil { c.onClose() } }() buf := make([]byte, 4096) for { n, err := c.conn.Read(buf) if err != nil { if c.onError != nil { c.onError(err) } return } if c.onMessage != nil { c.onMessage(string(buf[:n])) } } } // Send 发送消息 func (c *TCPClient) Send(message string) error { c.mu.RLock() defer c.mu.RUnlock() if !c.isConnected || c.conn == nil { return fmt.Errorf("未连接") } _, err := c.conn.Write([]byte(message)) return err } // Close 关闭连接 func (c *TCPClient) Close() error { c.mu.Lock() defer c.mu.Unlock() if c.conn != nil { err := c.conn.Close() c.isConnected = false return err } return nil } // IsConnected 是否已连接 func (c *TCPClient) IsConnected() bool { c.mu.RLock() defer c.mu.RUnlock() return c.isConnected } // GetTCPClient 获取 TCP 客户端 func GetTCPClient() *TCPClient { if tcpClient == nil { tcpClient = NewTCPClient() } return tcpClient } // ConnectTCP 连接 TCP func ConnectTCP(config *TCPConfig, onMessage func(string), onError func(error), onClose func()) error { client := GetTCPClient() return client.Connect(config, onMessage, onError, onClose) } // SendTCPMessage 发送 TCP 消息 func SendTCPMessage(message string) error { client := GetTCPClient() return client.Send(message) } // CloseTCP 关闭 TCP 连接 func CloseTCP() error { client := GetTCPClient() return client.Close() } // IsTCPConnected TCP 是否已连接 func IsTCPConnected() bool { client := GetTCPClient() return client.isConnected }