nnet/internal/client/websocket_client.go

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
}