package client import ( "net/url" "strings" "sync" "time" "git.noahlan.cn/noahlan/nnet/v2/pkg/client" protocolpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/protocol" unpackerpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/unpacker" "github.com/gorilla/websocket" ) // websocketClient WebSocket客户端实现 type websocketClient struct { config *client.Config conn *websocket.Conn mu sync.RWMutex requestMu sync.Mutex readMu sync.Mutex writeMu sync.Mutex connected bool protocol protocolpkg.Protocol unpacker unpackerpkg.Unpacker pending [][]byte } // NewWebSocketClient 创建WebSocket客户端 func NewWebSocketClient(config *client.Config) client.Client { if config == nil { config = client.DefaultConfig() } protocol := newApplicationProtocol(config.ApplicationProtocol) return &websocketClient{ config: config, connected: false, protocol: protocol, unpacker: newApplicationUnpacker(protocol), } } // Connect 连接服务器 func (c *websocketClient) Connect() error { c.mu.Lock() defer c.mu.Unlock() if c.connected { return nil } // 解析地址 addr := c.config.Addr scheme := "ws" if strings.HasPrefix(addr, "ws://") { addr = strings.TrimPrefix(addr, "ws://") scheme = "ws" } else if strings.HasPrefix(addr, "wss://") { addr = strings.TrimPrefix(addr, "wss://") scheme = "wss" } else if c.config.TLSEnabled { scheme = "wss" } // 构建URL u := url.URL{Scheme: scheme, Host: addr, Path: "/"} // 建立WebSocket连接 dialer := websocket.Dialer{ HandshakeTimeout: c.config.ConnectTimeout, } conn, _, err := dialer.Dial(u.String(), nil) if err != nil { return client.NewErrorf("failed to dial WebSocket: %v", err) } c.conn = conn c.connected = true c.unpacker = newApplicationUnpacker(c.protocol) c.pending = nil return nil } // Disconnect 断开连接 func (c *websocketClient) Disconnect() error { c.mu.Lock() defer c.mu.Unlock() if !c.connected { return nil } if c.conn != nil { c.conn.Close() c.conn = nil } c.connected = false c.pending = nil return nil } // Send 发送数据 func (c *websocketClient) Send(data []byte) error { frame, err := encodeApplicationMessage(c.protocol, data) if err != nil { return err } c.mu.RLock() if !c.connected || c.conn == nil { c.mu.RUnlock() return client.NewError("not connected") } conn := c.conn writeTimeout := c.config.WriteTimeout c.mu.RUnlock() c.writeMu.Lock() defer c.writeMu.Unlock() if writeTimeout > 0 { conn.SetWriteDeadline(time.Now().Add(writeTimeout)) } if err := conn.WriteMessage(websocket.BinaryMessage, frame); err != nil { c.markDisconnected() return err } return nil } // Receive 接收数据 func (c *websocketClient) Receive() ([]byte, error) { return c.receiveWithTimeout(0) } func (c *websocketClient) receiveWithTimeout(timeout time.Duration) ([]byte, error) { c.readMu.Lock() defer c.readMu.Unlock() if len(c.pending) > 0 { msg := c.pending[0] c.pending[0] = nil c.pending = c.pending[1:] return msg, nil } c.mu.RLock() if !c.connected || c.conn == nil { c.mu.RUnlock() return nil, client.NewError("not connected") } conn := c.conn readTimeout := c.config.ReadTimeout if timeout > 0 { readTimeout = timeout } protocol := c.protocol c.mu.RUnlock() // 设置读取超时 for { if readTimeout > 0 { conn.SetReadDeadline(time.Now().Add(readTimeout)) } _, data, err := conn.ReadMessage() if err != nil { c.markDisconnected() return nil, err } if protocol == nil { msg := make([]byte, len(data)) copy(msg, data) return msg, nil } if c.unpacker == nil { return decodeApplicationMessage(protocol, data) } frames, _, _, err := c.unpacker.Unpack(data) if err != nil { c.markDisconnected() return nil, err } if len(frames) == 0 { continue } decoded := make([][]byte, 0, len(frames)) for _, frame := range frames { msg, err := decodeApplicationMessage(protocol, frame) if err != nil { return nil, err } decoded = append(decoded, msg) } c.pending = append(c.pending, decoded[1:]...) return decoded[0], nil } } // Request 请求-响应(带超时) func (c *websocketClient) Request(data []byte, timeout time.Duration) ([]byte, error) { c.requestMu.Lock() defer c.requestMu.Unlock() // 发送请求 if err := c.Send(data); err != nil { return nil, err } // 接收响应 return c.receiveWithTimeout(timeout) } // IsConnected 检查是否已连接 func (c *websocketClient) IsConnected() bool { c.mu.RLock() defer c.mu.RUnlock() return c.connected } // Close 关闭客户端 func (c *websocketClient) Close() error { return c.Disconnect() } func (c *websocketClient) markDisconnected() { c.mu.Lock() defer c.mu.Unlock() if c.conn != nil { _ = c.conn.Close() c.conn = nil } c.connected = false c.pending = nil }