端口分配时加锁；网关上线后保存平台恢复的节点与授权数据；现在新节点连接后会按需向平台报告更新

gateway/fwd/ctrl.go

@@ -11,7 +11,6 @@ import (
 	"net"
 	"proxy-server/gateway/app"
 	"proxy-server/gateway/env"
-	"proxy-server/gateway/report"
 	"proxy-server/utils"
 	"strconv"
 	"time"
@@ -85,7 +84,14 @@ func processCtrlConn(_ctx context.Context, conn net.Conn) (err error) {
 	ctx, cancel := context.WithCancel(_ctx)
 	defer cancel()
 
-	var fwdPort uint16
+	// 结束时清理
+	var edgeId int32
+	defer func() {
+		var edge, ok = app.Edges.Load(edgeId)
+		if ok {
+			*edge.Status = 0
+		}
+	}()
 
 	// 处理连接命令
 	var errCh = make(chan error)
@@ -122,8 +128,8 @@ func processCtrlConn(_ctx context.Context, conn net.Conn) (err error) {
 					errCh <- fmt.Errorf("读取节点 ID 失败: %w", err)
 					return
 				}
-				var client = int32(binary.BigEndian.Uint32(recv))
-				fwdPort, err = onOpen(ctx, conn, client)
+				edgeId = int32(binary.BigEndian.Uint32(recv))
+				err = onOpen(ctx, conn, edgeId, conn.RemoteAddr())
 				if err != nil {
 					errCh <- fmt.Errorf("处理连接建立命令失败: %w", err)
 					return
@@ -159,54 +165,58 @@ func processCtrlConn(_ctx context.Context, conn net.Conn) (err error) {
 	case err = <-errCh:
 	}
 
-	app.DelEdge(fwdPort)
 	return
 }
 
-func onOpen(ctx context.Context, writer io.Writer, edge int32) (port uint16, err error) {
-	// open 命令全局只执行一次
-	_, ok := app.Edges.Load(edge)
-	if ok {
-		return 0, fmt.Errorf("节点 ID %d 已经连接", edge)
-	}
+func onOpen(ctx context.Context, writer io.Writer, edgeId int32, addr net.Addr) (err error) {
+	var port uint16
 
-	// 分配端口
-	var minim uint16 = 20000
-	var maxim uint16 = 60000
-	for i := minim; i < maxim; i++ {
-		var _, ok = app.Assigns.Load(i)
-		if !ok {
-			port = i
-			app.AddEdge(edge, port)
-			break
+	edge, ok := app.Edges.Load(edgeId)
+	if ok && edge.Port != nil {
+		port = *edge.Port
+	} else {
+		// 分配端口
+		app.LockPortAssign.Lock()
+
+		var minim uint16 = 20000
+		var maxim uint16 = 60000
+		for i := minim; i < maxim; i++ {
+			var _, ok = app.Assigns.Load(i)
+			if !ok {
+				port = i
+				break
+			}
+		}
+		if port == 0 {
+			return errors.New("没有可用的端口")
 		}
-	}
-	if port == 0 {
-		return 0, errors.New("没有可用的端口")
-	}
 
-	// 报告端口分配
-	if err = report.Assigned(edge, port); err != nil {
-		return 0, fmt.Errorf("报告端口分配失败: %w", err)
-	}
+		if tcpAddr, ok := addr.(*net.TCPAddr); ok {
+			app.NewEdge(edgeId, port, tcpAddr)
+		} else {
+			return fmt.Errorf("无效的地址类型: %T", addr)
+		}
 
-	// 响应节点
-	if err = sendPong(writer); err != nil {
-		return 0, fmt.Errorf("响应节点失败: %w", err)
+		app.LockPortAssign.Unlock()
 	}
 
 	// 启动转发服务
 	app.FwdLesWg.Add(1)
 	go func() {
 		defer app.FwdLesWg.Done()
-		slog.Info("监听转发端口", "port", port, "edge", edge)
+		slog.Info("监听转发端口", "port", port, "edge", edgeId)
 		err = ListenUser(ctx, port, writer)
 		if err != nil {
-			slog.Error("监听转发端口失败", "port", port, "edge", edge, "err", err)
+			slog.Error("监听转发端口失败", "port", port, "edge", edgeId, "err", err)
 		}
 	}()
 
-	return port, nil
+	// 响应节点
+	if err = sendPong(writer); err != nil {
+		return fmt.Errorf("响应节点失败: %w", err)
+	}
+
+	return nil
 }
 
 func onPing(writer io.Writer) (err error) {