优化连接效率

server/fwd/data.go

@@ -0,0 +1,120 @@
+package fwd
+
+import (
+	"io"
+	"log/slog"
+	"net"
+	"proxy-server/pkg/utils"
+	"proxy-server/server/pkg/env"
+	"strconv"
+	"sync"
+
+	"github.com/pkg/errors"
+)
+
+func (s *Service) startDataTun() error {
+	dataPort := env.AppDataPort
+	slog.Debug("监听数据通道", slog.Uint64("port", uint64(dataPort)))
+
+	// 监听端口
+	ls, err := net.Listen("tcp", ":"+strconv.Itoa(int(dataPort)))
+	if err != nil {
+		return errors.Wrap(err, "监听数据通道失败")
+	}
+	defer utils.Close(ls)
+
+	go func() {
+		<-s.ctx.Done()
+		utils.Close(ls)
+	}()
+
+	for {
+		conn, err := ls.Accept()
+		if err != nil {
+			return errors.Wrap(err, "监听数据通道失败")
+		}
+
+		select {
+		case <-s.ctx.Done():
+			utils.Close(conn)
+			return nil
+		default:
+			s.dataConnWg.Add(1)
+			go func() {
+				defer s.dataConnWg.Done()
+				defer utils.Close(conn)
+				err := s.processDataConn(conn)
+				if err != nil {
+					slog.Error("建立数据通道失败失败", "err", err)
+				}
+			}()
+		}
+	}
+}
+
+func (s *Service) processDataConn(client net.Conn) error {
+
+	// 接收 status
+	status, err := utils.ReadByte(client)
+	if err != nil {
+		return errors.Wrap(err, "从客户端获取 status 失败")
+	}
+
+	// 接收 tag
+	tagLen, err := utils.ReadByte(client)
+	if err != nil {
+		return errors.Wrap(err, "从客户端获取 tag 失败")
+	}
+	tagBuf, err := utils.ReadBuffer(client, int(tagLen))
+	if err != nil {
+		return errors.Wrap(err, "从客户端获取 tag 失败")
+	}
+	tag := string(tagBuf)
+
+	// 找到用户连接
+	user, ok := s.userConnMap.LoadAndDelete(tag)
+	if !ok {
+		return errors.New("用户连接已关闭，tag：" + tag)
+	}
+	defer utils.Close(user)
+
+	// 检查状态
+	if status != 1 {
+		return errors.New("目标地址建立连接失败")
+	}
+
+	// 数据转发
+	userPipeReader, userPipeWriter := io.Pipe()
+	defer utils.Close(userPipeWriter)
+	teeUser := io.TeeReader(user, userPipeWriter)
+	go func() {
+		err := analysisAndLog(user, userPipeReader)
+		if err != nil {
+			slog.Error("数据解析失败", "err", err)
+		}
+	}()
+
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(client, teeUser)
+		if err != nil {
+			slog.Error("数据转发失败 user->client", "err", err)
+		}
+	}()
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(user, client)
+		if err != nil {
+			slog.Error("数据转发失败 client->user", "err", err)
+		}
+	}()
+
+	select {
+	case <-s.ctx.Done():
+	case <-utils.ChanWgWait(s.ctx, &wg):
+	}
+
+	return nil
+}