feat: 去除 webTerminal 的 websocket 依赖

service/rpc/auth.go

@@ -10,19 +10,19 @@ import (
 	"github.com/naiba/nezha/service/singleton"
 )
 
-type AuthHandler struct {
+type authHandler struct {
 	ClientSecret string
 }
 
-func (a *AuthHandler) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
+func (a *authHandler) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
 	return map[string]string{"client_secret": a.ClientSecret}, nil
 }
 
-func (a *AuthHandler) RequireTransportSecurity() bool {
+func (a *authHandler) RequireTransportSecurity() bool {
 	return false
 }
 
-func (a *AuthHandler) Check(ctx context.Context) (uint64, error) {
+func (a *authHandler) Check(ctx context.Context) (uint64, error) {
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
 		return 0, status.Errorf(codes.Unauthenticated, "获取 metaData 失败")

service/rpc/io_stream.go

@@ -0,0 +1,140 @@
+package rpc
+
+import (
+	"errors"
+	"io"
+	"sync/atomic"
+	"time"
+)
+
+type ioStreamContext struct {
+	userIo           io.ReadWriteCloser
+	agentIo          io.ReadWriteCloser
+	userIoConnectCh  chan struct{}
+	agentIoConnectCh chan struct{}
+}
+
+func (s *NezhaHandler) CreateStream(streamId string) {
+	s.ioStreamMutex.Lock()
+	defer s.ioStreamMutex.Unlock()
+
+	s.ioStreams[streamId] = &ioStreamContext{
+		userIoConnectCh:  make(chan struct{}),
+		agentIoConnectCh: make(chan struct{}),
+	}
+}
+
+func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
+	s.ioStreamMutex.RLock()
+	defer s.ioStreamMutex.RUnlock()
+
+	if ctx, ok := s.ioStreams[streamId]; ok {
+		return ctx, nil
+	}
+
+	return nil, errors.New("stream not found")
+}
+
+func (s *NezhaHandler) CloseStream(streamId string) error {
+	s.ioStreamMutex.Lock()
+	defer s.ioStreamMutex.Unlock()
+
+	if ctx, ok := s.ioStreams[streamId]; ok {
+		if ctx.userIo != nil {
+			ctx.userIo.Close()
+		}
+		if ctx.agentIo != nil {
+			ctx.agentIo.Close()
+		}
+		delete(s.ioStreams, streamId)
+	}
+
+	return nil
+}
+
+func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
+	stream, err := s.GetStream(streamId)
+	if err != nil {
+		return err
+	}
+
+	stream.userIo = userIo
+	close(stream.userIoConnectCh)
+
+	return nil
+}
+
+func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error {
+	stream, err := s.GetStream(streamId)
+	if err != nil {
+		return err
+	}
+
+	stream.agentIo = agentIo
+	close(stream.agentIoConnectCh)
+
+	return nil
+}
+
+func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error {
+	stream, err := s.GetStream(streamId)
+	if err != nil {
+		return err
+	}
+
+	timeoutTimer := time.NewTimer(timeout)
+
+LOOP:
+	for {
+		select {
+		case <-stream.userIoConnectCh:
+			if stream.agentIo != nil {
+				timeoutTimer.Stop()
+				break LOOP
+			}
+		case <-stream.agentIoConnectCh:
+			if stream.userIo != nil {
+				timeoutTimer.Stop()
+				break LOOP
+			}
+		case <-time.After(timeout):
+			break LOOP
+		}
+	}
+
+	if stream.userIo == nil && stream.agentIo == nil {
+		return errors.New("timeout: no connection established")
+	}
+	if stream.userIo == nil {
+		return errors.New("timeout: user connection not established")
+	}
+	if stream.agentIo == nil {
+		return errors.New("timeout: agent connection not established")
+	}
+
+	isDone := new(atomic.Bool)
+	endCh := make(chan struct{})
+
+	go func() {
+		_, innerErr := io.Copy(stream.userIo, stream.agentIo)
+		if innerErr != nil {
+			err = innerErr
+		}
+		if isDone.CompareAndSwap(false, true) {
+			close(endCh)
+		}
+	}()
+	go func() {
+		_, innerErr := io.Copy(stream.agentIo, stream.userIo)
+		if innerErr != nil {
+			err = innerErr
+		}
+		if isDone.CompareAndSwap(false, true) {
+			close(endCh)
+		}
+	}()
+
+	<-endCh
+
+	return err
+}

service/rpc/nezha.go

@@ -3,11 +3,14 @@ package rpc
 import (
 	"context"
 	"fmt"
-	"github.com/naiba/nezha/pkg/ddns"
-	"github.com/naiba/nezha/pkg/utils"
 	"log"
+	"sync"
 	"time"
 
+	"github.com/naiba/nezha/pkg/ddns"
+	"github.com/naiba/nezha/pkg/grpcx"
+	"github.com/naiba/nezha/pkg/utils"
+
 	"github.com/jinzhu/copier"
 	"github.com/nicksnyder/go-i18n/v2/i18n"
 
@@ -16,8 +19,20 @@ import (
 	"github.com/naiba/nezha/service/singleton"
 )
 
+var NezhaHandlerSingleton *NezhaHandler
+
 type NezhaHandler struct {
-	Auth *AuthHandler
+	Auth          *authHandler
+	ioStreams     map[string]*ioStreamContext
+	ioStreamMutex *sync.RWMutex
+}
+
+func NewNezhaHandler() *NezhaHandler {
+	return &NezhaHandler{
+		Auth:          &authHandler{},
+		ioStreamMutex: new(sync.RWMutex),
+		ioStreams:     make(map[string]*ioStreamContext),
+	}
 }
 
 func (s *NezhaHandler) ReportTask(c context.Context, r *pb.TaskResult) (*pb.Receipt, error) {
@@ -177,3 +192,28 @@ func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Rece
 	singleton.ServerList[clientID].Host = &host
 	return &pb.Receipt{Proced: true}, nil
 }
+
+func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
+	if _, err := s.Auth.Check(stream.Context()); err != nil {
+		return err
+	}
+	id, err := stream.Recv()
+	if err != nil {
+		return err
+	}
+	if id == nil || len(id.Data) < 4 || (id.Data[0] != 0xff && id.Data[1] != 0x05 && id.Data[2] != 0xff && id.Data[3] == 0x05) {
+		return fmt.Errorf("invalid stream id")
+	}
+
+	streamId := string(id.Data[4:])
+
+	if _, err := s.GetStream(streamId); err != nil {
+		return err
+	}
+	iw := grpcx.NewIOStreamWrapper(stream)
+	if err := s.AgentConnected(streamId, iw); err != nil {
+		return err
+	}
+	iw.Wait()
+	return nil
+}