优化连接效率

server/fwd/analysis.go

@@ -19,7 +19,7 @@ func analysisAndLog(conn *core.Conn, reader io.Reader) error {
 	if err != nil {
 		err = errors.Wrap(err, "analysis sniffing error")
 	} else {
-		slog.Info(
+		slog.Debug(
 			"用户访问记录",
 			slog.Uint64("uid", uint64(conn.Auth.Payload.ID)),
 			slog.String("user", conn.RemoteAddr().String()),

server/fwd/core/auth.go

@@ -40,7 +40,7 @@ func CheckIp(conn net.Conn) (*AuthContext, error) {
 	_, localPort, err := net.SplitHostPort(localAddr)
 
 	// 查询权限记录
-	slog.Info("用户 " + remoteHost + " 请求连接到 " + localPort)
+	slog.Debug("用户 " + remoteHost + " 请求连接到 " + localPort)
 	var channels []models.Channel
 	err = orm.DB.
 		Joins("INNER JOIN public.nodes n ON channels.node_id = n.id AND n.name = ?", localPort).

server/fwd/ctrl.go

@@ -0,0 +1,238 @@
+package fwd
+
+import (
+	"bufio"
+	"context"
+	"log/slog"
+	"net"
+	"proxy-server/pkg/utils"
+	"proxy-server/server/fwd/core"
+	"proxy-server/server/fwd/dispatcher"
+	"proxy-server/server/models"
+	"proxy-server/server/pkg/env"
+	"proxy-server/server/pkg/orm"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/pkg/errors"
+)
+
+type CtrlCmd struct {
+	conn net.Conn
+	buf  []byte
+}
+
+var ctrlCmdChan = make(chan CtrlCmd, 1024)
+
+func (s *Service) startCtrlTun() error {
+	ctrlPort := env.AppCtrlPort
+	slog.Debug("监听控制通道", slog.Uint64("port", uint64(ctrlPort)))
+
+	// 监听端口
+	ls, err := net.Listen("tcp", ":"+strconv.Itoa(int(ctrlPort)))
+	if err != nil {
+		return errors.Wrap(err, "监听控制通道失败")
+	}
+	defer utils.Close(ls)
+
+	// 处理连接
+	connCh := utils.ChanConnAccept(s.ctx, ls)
+	err = nil
+	for loop := true; loop; {
+		select {
+		case <-s.ctx.Done():
+			loop = false
+		case conn, ok := <-connCh:
+			if !ok {
+				err = errors.New("获取连接失败")
+				loop = false
+			}
+			s.ctrlConnWg.Add(1)
+			go func() {
+				defer s.ctrlConnWg.Done()
+				defer utils.Close(conn)
+				err := s.processCtrlConn(conn)
+				if err != nil {
+					slog.Error("处理控制通道连接失败", "err", err)
+				}
+			}()
+		}
+	}
+
+	return err
+}
+
+func (s *Service) processCtrlConn(conn net.Conn) error {
+	reader := bufio.NewReader(conn)
+
+	// version
+	version, err := reader.ReadByte()
+	if err != nil {
+		_ = ctrlResp(conn, CtrlFail)
+		return errors.Wrap(err, "获取版本号失败")
+	}
+
+	// name
+	nameLen, err := reader.ReadByte()
+	if err != nil {
+		_ = ctrlResp(conn, CtrlFail)
+		return errors.Wrap(err, "获取 name 失败")
+	}
+	nameBuf, err := utils.ReadBuffer(reader, int(nameLen))
+	if err != nil {
+		_ = ctrlResp(conn, CtrlFail)
+		return errors.Wrap(err, "获取 name 失败")
+	}
+	name := string(nameBuf)
+
+	if name == "" {
+		_ = ctrlResp(conn, CtrlFail)
+		return errors.New("客户端名称不能为空")
+	}
+
+	// 检查客户端
+	var node models.Node
+	err = orm.DB.First(&node, &models.Node{
+		Name: name,
+	}).Error
+	if err != nil {
+		_ = ctrlResp(conn, CtrlFail)
+		return errors.Wrap(err, "查询客户端失败")
+	}
+
+	if version != node.Version {
+		_ = ctrlResp(conn, CtrlFail)
+		return errors.New("客户端版本不匹配")
+	}
+
+	err = ctrlResp(conn, CtrlDone)
+	if err != nil {
+		return errors.Wrap(err, "向客户端发送响应失败")
+	}
+
+	port := node.FwdPort
+	slog.Info("监听转发端口", "port", port, "client", name)
+
+	// 启动转发服务
+	proxy, err := dispatcher.New(port)
+	if err != nil {
+		return errors.Wrap(err, "创建 socks 转发服务失败")
+	}
+	defer proxy.Close()
+
+	s.fwdLesWg.Add(1)
+	go func() {
+		defer s.fwdLesWg.Done()
+		err := proxy.Run()
+		if err != nil {
+			slog.Error("代理服务运行失败", "err", err)
+		}
+	}()
+
+	// 监听控制通道连接
+	errCh := make(chan error, 1)
+	go func() {
+		defer close(errCh)
+		_, err := reader.ReadByte()
+		errCh <- err
+	}()
+
+	// 批量同步写入
+	go func() {
+		for {
+			select {
+			case <-s.ctx.Done():
+				return
+			case cmd := <-ctrlCmdChan:
+				_, err := cmd.conn.Write(cmd.buf)
+				if err != nil {
+					slog.Error("批量写入失败", "err", err)
+					utils.Close(cmd.conn)
+				}
+			}
+		}
+	}()
+
+	// 处理连接
+	for {
+		select {
+		case <-s.ctx.Done():
+			return nil
+		case err := <-errCh:
+			switch {
+			case strings.Contains(err.Error(), "An existing connection was forcibly closed by the remote host."):
+				slog.Debug("客户端主动断开连接")
+				return nil
+			case err == nil:
+				return errors.New("客户端握手失败")
+			default:
+				return errors.Wrap(err, "客户端意外断开连接")
+			}
+		case user := <-proxy.Conn:
+			s.userConnWg.Add(1)
+			go func() {
+				defer s.userConnWg.Done()
+				err := s.processUserConn(user, conn)
+				if err != nil {
+					slog.Error("处理用户连接失败", "err", err)
+					utils.Close(user)
+				}
+			}()
+		}
+	}
+}
+
+func (s *Service) processUserConn(user *core.Conn, ctrl net.Conn) error {
+
+	// 组织写入信息
+	dst := user.DestAddr().String()
+	dstLen := len(dst)
+
+	tag := user.Tag
+	tagLen := len(tag)
+
+	writeBuf := make([]byte, 2+dstLen+tagLen)
+	writeBuf[0] = byte(dstLen)
+	copy(writeBuf[1:], dst)
+	writeBuf[1+dstLen] = byte(tagLen)
+	copy(writeBuf[2+dstLen:], tag)
+
+	// 异步写入命令
+	ctrlCmdChan <- CtrlCmd{
+		conn: ctrl,
+		buf:  writeBuf,
+	}
+
+	// 记录用户连接
+	s.userConnMap.Store(user.Tag, user)
+
+	// 如果限定时间内没有建立数据通道，则关闭连接
+	timeout, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	select {
+	case <-s.ctx.Done():
+		// 服务会在退出时统一关闭未消费的连接
+	case <-timeout.Done():
+		storedUser, ok := s.userConnMap.LoadAndDelete(user.Tag)
+		if ok {
+			slog.Debug("建立数据通道超时", "tag", user.Tag)
+			utils.Close(storedUser)
+		}
+	}
+
+	return nil
+}
+
+type CtrlResult byte
+
+const (
+	CtrlFail CtrlResult = iota
+	CtrlDone
+)
+
+func ctrlResp(conn net.Conn, result CtrlResult) error {
+	_, err := conn.Write([]byte{byte(result)})
+	return err
+}

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
+}

server/fwd/dispatcher/dispatch.go

@@ -77,8 +77,9 @@ func (s *Server) Run() error {
 		}
 	}()
 
-	errCh := make(chan error)
+	errCh := make(chan error, 1)
 	go func() {
+		defer close(errCh)
 		err = m.Serve()
 		if err != nil {
 			err = errors.Wrap(err, "dispatcher serve error")

server/fwd/fwd.go

@@ -1,22 +1,11 @@
 package fwd
 
 import (
-	"bufio"
 	"context"
-	"encoding/binary"
-	"io"
 	"log/slog"
-	"net"
 	"proxy-server/pkg/utils"
 	"proxy-server/server/fwd/core"
-	"proxy-server/server/fwd/dispatcher"
-	"proxy-server/server/pkg/env"
-	"strconv"
-	"strings"
 	"sync"
-	"time"
-
-	"github.com/pkg/errors"
 )
 
 type Config struct {
@@ -122,256 +111,3 @@ func (s *Service) Run() {
 	wg.Wait()
 	slog.Info("fwd 服务已退出")
 }
-
-func (s *Service) startCtrlTun() error {
-	ctrlPort := env.AppCtrlPort
-	slog.Debug("监听控制通道", slog.Uint64("port", uint64(ctrlPort)))
-
-	// 监听端口
-	ls, err := net.Listen("tcp", ":"+strconv.Itoa(int(ctrlPort)))
-	if err != nil {
-		return errors.Wrap(err, "监听控制通道失败")
-	}
-	defer utils.Close(ls)
-
-	// 处理连接
-	connCh := utils.ChanConnAccept(s.ctx, ls)
-	for {
-		select {
-		case <-s.ctx.Done():
-			return nil
-		case conn, ok := <-connCh:
-			if !ok {
-				return errors.New("获取连接失败")
-			}
-			s.ctrlConnWg.Add(1)
-			go func() {
-				defer s.ctrlConnWg.Done()
-				err := s.processCtrlConn(conn)
-				if err != nil {
-					slog.Error("处理控制通道连接失败", "err", err)
-					utils.Close(conn)
-				}
-			}()
-		}
-	}
-}
-
-func (s *Service) processCtrlConn(conn net.Conn) error {
-	reader := bufio.NewReader(conn)
-
-	// 获取转发端口
-	portBuf, err := utils.ReadBuffer(reader, 2)
-	if err != nil {
-		return errors.Wrap(err, "获取转发端口失败")
-	}
-	port := binary.BigEndian.Uint16(portBuf)
-
-	slog.Info("客户端注册", "addr", conn.RemoteAddr().String(), "port", port)
-
-	// 启动转发服务
-	proxy, err := dispatcher.New(port)
-	if err != nil {
-		return errors.Wrap(err, "创建 socks 转发服务失败")
-	}
-	defer proxy.Close()
-
-	s.fwdLesWg.Add(1)
-	go func() {
-		defer s.fwdLesWg.Done()
-		err := proxy.Run()
-		if err != nil {
-			slog.Error("代理服务运行失败", "err", err)
-			return
-		}
-	}()
-
-	// 监听客户端连接
-	errCh := make(chan error)
-	defer close(errCh)
-	go func() {
-		_, err := reader.ReadByte()
-		errCh <- err
-	}()
-
-	// 处理连接
-	for {
-		select {
-		case <-s.ctx.Done():
-			return nil
-		case err := <-errCh:
-			switch {
-			case strings.Contains(err.Error(), "An existing connection was forcibly closed by the remote host."):
-				slog.Debug("客户端主动断开连接")
-				return nil
-			case err == nil:
-				return errors.New("客户端握手失败")
-			default:
-				return errors.Wrap(err, "客户端意外断开连接")
-			}
-		case user := <-proxy.Conn:
-			s.userConnWg.Add(1)
-			go func() {
-				defer s.userConnWg.Done()
-				err := s.processUserConn(user, conn)
-				if err != nil {
-					slog.Error("处理用户连接失败", "err", err)
-					utils.Close(user)
-				}
-			}()
-		}
-	}
-}
-
-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 {
-
-	// 读取 tag
-	var tag string
-	select {
-	case <-s.ctx.Done():
-		return nil
-	default:
-		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("查找用户连接失败")
-	}
-	defer utils.Close(user)
-
-	// 发送目标地址
-	select {
-	case <-s.ctx.Done():
-		return nil
-	default:
-		dest := user.Dest.String()
-		destLen := len(dest)
-		destBuf := make([]byte, 1+destLen)
-		destBuf[0] = byte(destLen)
-		copy(destBuf[1:], dest)
-		_, err := client.Write(destBuf)
-		if err != nil {
-			return errors.Wrap(err, "向客户端发送目标地址失败")
-		}
-	}
-
-	// 数据转发
-	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, "errType")
-		}
-	}()
-
-	select {
-	case <-s.ctx.Done():
-	case <-utils.ChanWgWait(s.ctx, &wg):
-	}
-
-	return nil
-}
-
-func (s *Service) processUserConn(user *core.Conn, ctrl net.Conn) error {
-
-	// 发送 tag
-	tag := user.Tag
-	tagLen := len(tag)
-	tagBuf := make([]byte, 1+tagLen)
-	tagBuf[0] = byte(tagLen)
-	copy(tagBuf[1:], tag)
-
-	_, err := ctrl.Write(tagBuf)
-	if err != nil {
-		return errors.Wrap(err, "向控制通道发送 tag 失败")
-	}
-
-	// 记录用户连接
-	s.userConnMap.Store(user.Tag, user)
-
-	// 如果限定时间内没有建立数据通道，则关闭连接
-	timeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-
-	select {
-	case <-s.ctx.Done():
-		// 服务会在退出时统一关闭未消费的连接
-	case <-timeout.Done():
-		storedUser, ok := s.userConnMap.LoadAndDelete(user.Tag)
-		if ok {
-			slog.Debug("用户连接超时，关闭连接", "tag", user.Tag)
-			utils.Close(storedUser)
-		}
-	}
-
-	return nil
-}

server/models/node.go

@@ -5,10 +5,11 @@ import "gorm.io/gorm"
 // Node 客户端模型
 type Node struct {
 	gorm.Model
-	Name      string
-	Provider  string
-	Location  string
-	IPAddress string
+	Name     string
+	Version  byte
+	FwdPort  uint16
+	Provider string
+	Location string
 
 	Channels []Channel `gorm:"foreignKey:NodeId"`
 }

server/pkg/log/logs.go

@@ -15,11 +15,13 @@ func Init() {
 		mode = "dev"
 	}
 
+	level := slog.LevelInfo
+
 	switch mode {
 	case "dev":
 		writer := colorable.NewColorable(os.Stdout)
 		logger := slog.New(tint.NewHandler(writer, &tint.Options{
-			Level:      slog.LevelDebug,
+			Level:      level,
 			TimeFormat: time.RFC3339,
 			ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
 				err, ok := attr.Value.Any().(error)
@@ -32,7 +34,7 @@ func Init() {
 		slog.SetDefault(logger)
 	case "test":
 		logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
-			Level: slog.LevelInfo,
+			Level: level,
 		}))
 		slog.SetDefault(logger)
 	default: