完善客户端与服务端注册，端口分配和协议交互逻辑

server/fwd/ctrl.go

@@ -3,20 +3,22 @@ package fwd
 import (
 	"bufio"
 	"context"
+	"encoding/binary"
+	"fmt"
+	"io"
 	"log/slog"
 	"net"
 	"proxy-server/pkg/utils"
 	"proxy-server/server/fwd/core"
 	"proxy-server/server/fwd/dispatcher"
 	"proxy-server/server/fwd/metrics"
-	"proxy-server/server/fwd/repo"
 	"proxy-server/server/pkg/env"
-	"proxy-server/server/pkg/orm"
+	"proxy-server/server/report"
 	"strconv"
 	"strings"
 	"time"
 
-	"github.com/pkg/errors"
+	"errors"
 )
 
 type CtrlCmd struct {
@@ -33,7 +35,7 @@ func (s *Service) startCtrlTun() error {
 	// 监听端口
 	ls, err := net.Listen("tcp", ":"+strconv.Itoa(int(ctrlPort)))
 	if err != nil {
-		return errors.Wrap(err, "监听控制通道失败")
+		return fmt.Errorf("监听控制通道失败: %w", err)
 	}
 	defer utils.Close(ls)
 
@@ -67,58 +69,49 @@ func (s *Service) startCtrlTun() error {
 func (s *Service) processCtrlConn(conn net.Conn) error {
 	reader := bufio.NewReader(conn)
 
-	// version
-	version, err := reader.ReadByte()
+	var recv = make([]byte, 4)
+	_, err := io.ReadFull(reader, recv)
 	if err != nil {
-		_ = ctrlResp(conn, CtrlFail)
-		return errors.Wrap(err, "获取版本号失败")
+		return fmt.Errorf("读取客户端 ID 失败: %w", err)
+	}
+	var clientId = int32(binary.BigEndian.Uint32(recv))
+
+	// 分配端口
+	var minim uint16 = 20000
+	var maxim uint16 = 60000
+	var fwdPort uint16
+	for i := minim; i < maxim; i++ {
+		var _, ok = s.fwdPortMap[i]
+		if !ok {
+			fwdPort = i
+			s.fwdPortMap[i] = clientId
+			break
+		}
+	}
+	if fwdPort == 0 {
+		return errors.New("没有可用的端口")
 	}
 
-	// name
-	nameLen, err := reader.ReadByte()
+	// 报告端口分配
+	if s.Config.Id == nil || *s.Config.Id == 0 {
+		return errors.New("转发服务未成功注册，无法提供服务")
+	}
+	err = report.Assigned(s.ctx, *s.Config.Id, clientId, fwdPort)
 	if err != nil {
-		_ = ctrlResp(conn, CtrlFail)
-		return errors.Wrap(err, "获取 name 失败")
+		return fmt.Errorf("报告端口分配失败: %w", err)
 	}
-	nameBuf, err := utils.ReadBuffer(reader, int(nameLen))
+
+	// 响应客户端
+	_, err = conn.Write([]byte{1})
 	if err != nil {
-		_ = ctrlResp(conn, CtrlFail)
-		return errors.Wrap(err, "获取 name 失败")
+		return fmt.Errorf("响应客户端失败: %w", err)
 	}
-	name := string(nameBuf)
-
-	if name == "" {
-		_ = ctrlResp(conn, CtrlFail)
-		return errors.New("客户端名称不能为空")
-	}
-
-	// 检查客户端
-	var node repo.Node
-	err = orm.DB.Take(&node, &repo.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)
+	slog.Info("监听转发端口", "port", fwdPort, "client", clientId)
+	proxy, err := dispatcher.New(fwdPort)
 	if err != nil {
-		return errors.Wrap(err, "创建 socks 转发服务失败")
+		return err
 	}
 	defer proxy.Close()
 
@@ -168,7 +161,7 @@ func (s *Service) processCtrlConn(conn net.Conn) error {
 			case err == nil:
 				return errors.New("客户端握手失败")
 			default:
-				return errors.Wrap(err, "客户端意外断开连接")
+				return fmt.Errorf("客户端意外断开连接: %w", err)
 			}
 		case user := <-proxy.Conn:
 			metrics.TimerAuth.Store(user.Conn, time.Now())
@@ -226,15 +219,3 @@ func (s *Service) processUserConn(user *core.Conn, ctrl net.Conn) error {
 
 	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/fwd.go

@@ -9,6 +9,7 @@ import (
 )
 
 type Config struct {
+	Id *int32
 }
 
 type Service struct {
@@ -22,6 +23,8 @@ type Service struct {
 	ctrlConnWg utils.CountWaitGroup
 	dataConnWg utils.CountWaitGroup
 	userConnWg utils.CountWaitGroup
+
+	fwdPortMap map[uint16]int32 // 转发端口映射，key 为端口号，value 为边缘节点 ID
 }
 
 func New(config *Config) *Service {
@@ -31,9 +34,10 @@ func New(config *Config) *Service {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Service{
-		Config: config,
-		ctx:    ctx,
-		cancel: cancel,
+		Config:     config,
+		ctx:        ctx,
+		cancel:     cancel,
+		fwdPortMap: make(map[uint16]int32),
 	}
 }
 

server/report/report.go

@@ -47,7 +47,7 @@ func Offline(ctx context.Context, name string) (err error) {
 	return err
 }
 
-func Assigned(ctx context.Context, id int32, edgeId int32, port int16) (err error) {
+func Assigned(ctx context.Context, id int32, edgeId int32, port uint16) (err error) {
 	_, err = repeat(ctx, env.EndpointAssigned, map[string]any{
 		"proxy": id,
 		"edge":  edgeId,
@@ -85,7 +85,7 @@ func repeat(ctx context.Context, endpoint string, body any) (string, error) {
 		default:
 		}
 
-		slog.Warn("服务注册失败，五秒后重试", "err", err)
+		slog.Warn("服务调用失败，五秒后重试", "err", err)
 		time.Sleep(5 * time.Second)
 	}
 }

server/server.go

@@ -100,6 +100,7 @@ func (s *server) Run() (err error) {
 		id, err := report.Online(ctx, s.name)
 		if err != nil {
 			reportErrCh <- err
+			return
 		}
 		s.id = id
 	}()
@@ -190,7 +191,9 @@ func (s *server) restore() error {
 }
 
 func (s *server) startFwd(ctx context.Context) error {
-	server := fwd.New(nil)
+	server := fwd.New(&fwd.Config{
+		Id: &s.id,
+	})
 	go func() {
 		<-ctx.Done()
 		server.Stop()