feat: 去除 webTerminal 的 websocket 依赖

pkg/grpcx/io_stream_wrapper.go

@@ -0,0 +1,65 @@
+package grpcx
+
+import (
+	"context"
+	"io"
+	"sync/atomic"
+
+	"github.com/naiba/nezha/proto"
+)
+
+var _ io.ReadWriteCloser = &IOStreamWrapper{}
+
+type IOStream interface {
+	Recv() (*proto.IOStreamData, error)
+	Send(*proto.IOStreamData) error
+	Context() context.Context
+}
+
+type IOStreamWrapper struct {
+	IOStream
+	dataBuf []byte
+	closed  *atomic.Bool
+	closeCh chan struct{}
+}
+
+func NewIOStreamWrapper(stream IOStream) *IOStreamWrapper {
+	return &IOStreamWrapper{
+		IOStream: stream,
+		closeCh:  make(chan struct{}),
+		closed:   new(atomic.Bool),
+	}
+}
+
+func (iw *IOStreamWrapper) Read(p []byte) (n int, err error) {
+	if len(iw.dataBuf) > 0 {
+		n := copy(p, iw.dataBuf)
+		iw.dataBuf = iw.dataBuf[n:]
+		return n, nil
+	}
+	var data *proto.IOStreamData
+	if data, err = iw.Recv(); err != nil {
+		return 0, err
+	}
+	n = copy(p, data.Data)
+	if n < len(data.Data) {
+		iw.dataBuf = data.Data[n:]
+	}
+	return n, nil
+}
+
+func (iw *IOStreamWrapper) Write(p []byte) (n int, err error) {
+	err = iw.Send(&proto.IOStreamData{Data: p})
+	return len(p), err
+}
+
+func (iw *IOStreamWrapper) Close() error {
+	if iw.closed.CompareAndSwap(false, true) {
+		close(iw.closeCh)
+	}
+	return nil
+}
+
+func (iw *IOStreamWrapper) Wait() {
+	<-iw.closeCh
+}

pkg/websocketx/safe_conn.go

@@ -4,22 +4,44 @@ import (
 	"sync"
 
 	"github.com/gorilla/websocket"
-	"github.com/samber/lo"
 )
 
 type Conn struct {
 	*websocket.Conn
-	writeLock sync.Mutex
+	writeLock *sync.Mutex
+	dataBuf   []byte
 }
 
-func (conn *Conn) WriteMessage(msgType int, data []byte) error {
+func NewConn(conn *websocket.Conn) *Conn {
+	return &Conn{Conn: conn, writeLock: new(sync.Mutex)}
+}
+
+func (conn *Conn) Write(data []byte) (int, error) {
 	conn.writeLock.Lock()
 	defer conn.writeLock.Unlock()
-	var err error
-	lo.TryCatchWithErrorValue(func() error {
-		return conn.Conn.WriteMessage(msgType, data)
-	}, func(res any) {
-		err = res.(error)
-	})
-	return err
+	if err := conn.Conn.WriteMessage(websocket.BinaryMessage, data); err != nil {
+		return 0, err
+	}
+	return len(data), nil
+}
+
+func (conn *Conn) Read(data []byte) (int, error) {
+	if len(conn.dataBuf) > 0 {
+		n := copy(data, conn.dataBuf)
+		conn.dataBuf = conn.dataBuf[n:]
+		return n, nil
+	}
+	mType, innerData, err := conn.Conn.ReadMessage()
+	if err != nil {
+		return 0, err
+	}
+	// 将文本消息转换为命令输入
+	if mType == websocket.TextMessage {
+		innerData = append([]byte{0}, innerData...)
+	}
+	n := copy(data, innerData)
+	if n < len(innerData) {
+		conn.dataBuf = innerData[n:]
+	}
+	return n, nil
 }