优化连接处理逻辑，增加超时设置；重构命令读取与错误处理；新增公共工具函数以简化错误处理

gateway/env/env.go

@@ -14,13 +14,16 @@ import (
 var (
 	RunMode = "dev" // 运行模式，dev: 开发模式，prod: 生产模式
 
-	AppCtrlPort    uint16 = 18080
-	AppDataPort    uint16 = 18081
-	AppWebPort     uint16 = 8848
-	AppLogMode            = "dev"
-	AppExitTimeout        = 5  // 等待服务停止的超时时间
-	AppDataTimeout        = 10 // 等待数据通道连接的超时时间
-	AppUserTimeout        = 10 // 等待用户发送数据的超时时间（端口复用需要分析协议，如果用户长期不发送数据，将会阻塞分析协程）
+	AppCtrlPort      uint16 = 18080
+	AppDataPort      uint16 = 18081
+	AppWebPort       uint16 = 8848
+	AppLogMode              = "dev"
+	AppExitTimeout          = 5  // 等待服务停止的超时时间
+	AppDataTimeout          = 10 // 等待数据通道连接的超时时间
+	AppUserRWTimeout        = 10 // 等待用户连接读写超时时间
+	AppDataRWTimeout        = 10 // 等待数据通道读写超时时间
+	AppCtrlRWTimeout        = 10 // 等待控制通道读写超时时间
+	AppCtrlHBTimeout        = 30 // 控制通道心跳超时时间（断开连接等待时间为：心跳等待时间 * 2 + 读写等待时间）
 
 	AuthWhitelist []net.IP // 全局白名单，可以将白名单 IP 视为一个可信任代理
 
@@ -103,13 +106,40 @@ func Init() {
 		AppDataTimeout = appDataTimeout
 	}
 
-	value = os.Getenv("APP_USER_TIMEOUT")
+	value = os.Getenv("APP_USER_RW_TIMEOUT")
 	if value != "" {
 		appUserTimeout, err := strconv.Atoi(value)
 		if err != nil {
-			panic(fmt.Sprintf("环境变量 APP_USER_TIMEOUT 格式错误: %v", err))
+			panic(fmt.Sprintf("环境变量 APP_USER_RW_TIMEOUT 格式错误: %v", err))
 		}
-		AppUserTimeout = appUserTimeout
+		AppUserRWTimeout = appUserTimeout
+	}
+
+	value = os.Getenv("APP_DATA_RW_TIMEOUT")
+	if value != "" {
+		appDataRWTimeout, err := strconv.Atoi(value)
+		if err != nil {
+			panic(fmt.Sprintf("环境变量 APP_DATA_RW_TIMEOUT 格式错误: %v", err))
+		}
+		AppDataRWTimeout = appDataRWTimeout
+	}
+
+	value = os.Getenv("APP_CTRL_RW_TIMEOUT")
+	if value != "" {
+		appCtrlRWTimeout, err := strconv.Atoi(value)
+		if err != nil {
+			panic(fmt.Sprintf("环境变量 APP_CTRL_RW_TIMEOUT 格式错误: %v", err))
+		}
+		AppCtrlRWTimeout = appCtrlRWTimeout
+	}
+
+	value = os.Getenv("APP_CTRL_HB_TIMEOUT")
+	if value != "" {
+		appCtrlHBTimeout, err := strconv.Atoi(value)
+		if err != nil {
+			panic(fmt.Sprintf("环境变量 APP_CTRL_HB_TIMEOUT 格式错误: %v", err))
+		}
+		AppCtrlHBTimeout = appCtrlHBTimeout
 	}
 
 	value = os.Getenv("AUTH_WHITELIST")

gateway/fwd/ctrl.go

@@ -14,7 +14,7 @@ import (
 	"proxy-server/gateway/report"
 	"proxy-server/utils"
 	"strconv"
-	"syscall"
+	"time"
 )
 
 type CtrlCmdType int
@@ -37,8 +37,7 @@ func ListenCtrl(ctx context.Context) error {
 	}
 	defer utils.Close(ls)
 
-	// 处理连接
-	// 异步等待连接
+	// 异步等待处理连接
 	var connCh = make(chan net.Conn)
 	go func() {
 		for {
@@ -80,77 +79,88 @@ func ListenCtrl(ctx context.Context) error {
 }
 
 func processCtrlConn(_ctx context.Context, conn net.Conn) (err error) {
-	// 通道上下文
-	ctx, cancel := context.WithCancel(_ctx)
+	reader := bufio.NewReader(conn)
+
+	// 上下文与通道信息
+	ctx, cancel := context.WithCancel(_ctx)
+	defer cancel()
 
-	// 结束后清理资源
 	var fwdPort uint16
-	defer func() {
-		slog.Debug("关闭控制通道", "port", fwdPort)
-		app.DelEdge(fwdPort)
+
+	// 处理连接命令
+	var errCh = make(chan error)
+	go func() {
+		var err error
+		for {
+			// 读取命令
+			var timeout = time.Duration(env.AppCtrlHBTimeout*2+env.AppCtrlRWTimeout) * time.Second
+			err = conn.SetReadDeadline(time.Now().Add(timeout))
+			if err != nil {
+				errCh <- fmt.Errorf("设置读取超时失败: %w", err)
+				return
+			}
+
+			cmd, err := reader.ReadByte()
+			if err := utils.WarpConnErr(err); err != nil {
+				errCh <- err
+				return
+			}
+
+			// 处理节点命令
+			err = conn.SetReadDeadline(time.Now().Add(time.Duration(env.AppCtrlRWTimeout) * time.Second))
+			if err != nil {
+				errCh <- fmt.Errorf("设置读取超时失败: %w", err)
+				return
+			}
+			switch CtrlCmdType(cmd) {
+
+			// 连接建立命令
+			case CtrlCmdOpen:
+				var recv = make([]byte, 4)
+				_, err = io.ReadFull(reader, recv)
+				if err != nil {
+					errCh <- fmt.Errorf("读取节点 ID 失败: %w", err)
+					return
+				}
+				var client = int32(binary.BigEndian.Uint32(recv))
+				fwdPort, err = onOpen(ctx, conn, client)
+				if err != nil {
+					errCh <- fmt.Errorf("处理连接建立命令失败: %w", err)
+					return
+				}
+
+			// 心跳命令
+			case CtrlCmdPing:
+				err = onPing(conn)
+				if err != nil {
+					errCh <- fmt.Errorf("处理心跳命令失败: %w", err)
+					return
+				}
+
+			// 连接关闭命令
+			case CtrlCmdClose:
+				err = onClose(conn)
+				if err != nil {
+					errCh <- fmt.Errorf("处理关闭命令失败: %w", err)
+					return
+				}
+
+			// 忽略其他不应该由节点发起的命令
+			default:
+				errCh <- fmt.Errorf("无法处理控制命令: %d", cmd)
+				return
+			}
+		}
 	}()
 
-	// 处理控制命令
-	defer cancel()
-	reader := bufio.NewReader(conn)
-	for {
-		// 循环等待直到服务关闭
-		select {
-		case <-ctx.Done():
-			return nil
-		default:
-		}
-
-		// 读取命令
-		cmd, err := reader.ReadByte()
-		if errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.WSAECONNRESET) {
-			slog.Debug("节点重置了控制通道连接(WSAECONNRESET)")
-			return nil
-		}
-		if errors.Is(err, io.EOF) {
-			slog.Debug("节点关闭了控制通道")
-			return nil
-		}
-		if err != nil {
-			return fmt.Errorf("读取节点命令失败: %w", err)
-		}
-
-		// 处理节点命令
-		switch CtrlCmdType(cmd) {
-
-		// 连接建立命令
-		case CtrlCmdOpen:
-			var recv = make([]byte, 4)
-			_, err = io.ReadFull(reader, recv)
-			if err != nil {
-				return fmt.Errorf("读取节点 ID 失败: %w", err)
-			}
-			var client = int32(binary.BigEndian.Uint32(recv))
-			fwdPort, err = onOpen(ctx, conn, client)
-			if err != nil {
-				return fmt.Errorf("处理连接建立命令失败: %w", err)
-			}
-
-		// 心跳命令
-		case CtrlCmdPing:
-			err = onPing(conn)
-			if err != nil {
-				return fmt.Errorf("处理心跳命令失败: %w", err)
-			}
-
-		// 连接关闭命令
-		case CtrlCmdClose:
-			err = onClose(conn)
-			if err != nil {
-				return fmt.Errorf("处理关闭命令失败: %w", err)
-			}
-			return nil
-
-		// 忽略其他不应该由节点发起的命令
-		default:
-			return fmt.Errorf("无法处理控制命令: %d", cmd)
-		}
+	// 等待处理结束
+	select {
+	case <-ctx.Done():
+	case err = <-errCh:
 	}
+
+	app.DelEdge(fwdPort)
+	return
 }
 
 func onOpen(ctx context.Context, writer io.Writer, edge int32) (port uint16, err error) {

gateway/fwd/data.go

@@ -73,8 +73,8 @@ func ListenData(ctx context.Context) error {
 	}
 }
 
-func processDataConn(ctx context.Context, client net.Conn) error {
-	var reader = bufio.NewReader(client)
+func processDataConn(ctx context.Context, edge net.Conn) error {
+	var reader = bufio.NewReader(edge)
 
 	// 接收连接结果
 	var buf = make([]byte, 17)
@@ -133,7 +133,7 @@ func processDataConn(ctx context.Context, client net.Conn) error {
 	// 复制用户数据到节点
 	var waitUser = make(chan error)
 	go func() {
-		_, err := io.Copy(client, teeUser)
+		_, err := io.Copy(edge, teeUser)
 		switch {
 		case errors.Is(err, net.ErrClosed):
 			slog.Debug("用户连接意外关闭")

gateway/fwd/user.go

@@ -17,7 +17,7 @@ import (
 )
 
 func ListenUser(ctx context.Context, port uint16, ctrl io.Writer) error {
-	dspt, err := dispatcher.New(port, time.Duration(env.AppUserTimeout)*time.Second)
+	dspt, err := dispatcher.New(port, time.Duration(env.AppUserRWTimeout)*time.Second)
 	if err != nil {
 		return err
 	}