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

edge/edge.go

@@ -21,6 +21,8 @@ import (
 )
 
 func Start() error {
+	var ctx, cancel = signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
+	defer cancel()
 
 	// 初始化环境变量
 	slog.Debug("初始化环境变量...")
@@ -44,36 +46,17 @@ func Start() error {
 	}
 
 	// 连接到网关
-	var ctx, cancel = signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
-	defer cancel()
-
 	var errCh = make(chan error)
 	go func() {
-		for {
-			err = ctrl(ctx, id, host)
-			if err == nil {
-				errCh <- nil
-				return
-			}
-
-			select {
-			case <-ctx.Done():
-				return
-			default:
-				slog.Error("建立控制通道失败", "err", err)
-				slog.Info(fmt.Sprintf("%d 秒后重试", core.RetryInterval))
-			}
-
-			select {
-			case <-ctx.Done():
-				return
-			case <-time.After(time.Duration(core.RetryInterval) * time.Second):
-			}
+		err = ctrl(ctx, id, host)
+		if err == nil {
+			errCh <- err
 		}
 	}()
 
 	// 等待退出
 	select {
+	case <-ctx.Done():
 	case err := <-errCh:
 		if err != nil {
 			slog.Error("控制通道发生错误", "err", err)
@@ -102,7 +85,7 @@ func ctrl(ctx context.Context, id int32, host string) error {
 	defer utils.Close(conn)
 	var reader = bufio.NewReader(conn)
 
-	// 发送节点连接命令
+	// 发送开启连接
 	slog.Debug("发送节点连接命令")
 	err = sendOpen(conn, id)
 	if err != nil {
@@ -126,78 +109,55 @@ func ctrl(ctx context.Context, id int32, host string) error {
 		}
 	}()
 
-	// 异步等待连接命令
-	slog.Info("等待用户连接")
-	var cmdCh = make(chan ConnCmd)
+	// 异步读取节点命令
 	var errCh = make(chan error)
 	go func() {
 		for {
+			// 读取命令
 			cmd, err := reader.ReadByte()
-			if err != nil {
-				switch {
-				case errors.Is(err, net.ErrClosed):
-					err = fmt.Errorf("控制通道关闭: %w", err)
-				case errors.Is(err, io.EOF):
-					err = fmt.Errorf("网关关闭了控制通道: %w", err)
-				default:
-					err = fmt.Errorf("读取命令失败: %w", err)
-				}
+			if err := utils.WarpConnErr(err); err != nil {
 				errCh <- err
 				return
 			}
-
 			switch cmd {
+
+			// pong 命令，忽略
 			case 1:
-				// 忽略网关响应的 pong 命令
+
+			// 代理命令
 			case 5:
-				tag, addr, err := onConn(reader)
+				err := onConn(reader, dataAddr)
 				if err != nil {
-					slog.Error("接收连接命令失败", "err", err)
+					errCh <- fmt.Errorf("处理代理命令失败: %w", err)
 					return
 				}
-				cmdCh <- ConnCmd{
-					Tag:  tag,
-					Addr: addr,
-				}
 			}
 		}
 	}()
 
 	// 等待建立数据通道
-	for loop := true; loop; {
-		select {
-		case <-ctx.Done():
-			loop = false
-		case err = <-errCh:
-			slog.Error("读取控制命令失败", "err", err)
-			loop = false
-		case cmd := <-cmdCh:
-			slog.Debug("建立数据通道", "tag", cmd.Tag, "addr", cmd.Addr)
-			go func() {
-				err := data(dataAddr, cmd.Addr, cmd.Tag)
-				if err != nil {
-					slog.Error("建立数据通道失败", "err", err)
-				}
-			}()
-		}
+	select {
+	case <-ctx.Done():
+	case err = <-errCh:
 	}
 
-	// 发送关闭连接（不 return err，否则会重新连接）
+	// 发送关闭连接
 	slog.Debug("发送关闭连接")
 	err = sendClose(conn)
 	if err != nil {
-		slog.Error("发送关闭连接失败", "err", err)
+		return fmt.Errorf("发送关闭连接失败: %w", err)
 	}
 
 	return nil
 }
 
-func data(proxy string, dest string, tag [16]byte) error {
+func data(proxy string, destination string, tag [16]byte) error {
+	slog.Debug("建立数据通道", "tag", tag, "addr", destination)
 
 	// 向目标地址建立连接
 	var result = 1
 	var dstErr error
-	dst, err := net.Dial("tcp", dest)
+	dst, err := net.Dial("tcp", destination)
 	if err != nil {
 		dstErr = fmt.Errorf("连接目标地址失败: %w", dstErr)
 		result = 0
@@ -295,24 +255,34 @@ func sendPing(writer io.Writer) error {
 	return nil
 }
 
-func onConn(reader io.Reader) (tag [16]byte, addr string, err error) {
+func onConn(reader io.Reader, dataAddr string) (err error) {
+
+	// 读取连接命令
 	var buf = make([]byte, 16+2)
 	_, err = io.ReadFull(reader, buf)
 	if err != nil {
-		return [16]byte{}, "", err
+		return err
 	}
 
-	tag = [16]byte(buf[0:16])
+	var tag = [16]byte(buf[0:16])
 
 	var addrLen = binary.BigEndian.Uint16(buf[16:18])
 	var addrBuf = make([]byte, addrLen)
 	_, err = io.ReadFull(reader, addrBuf)
 	if err != nil {
-		return [16]byte{}, "", err
+		return err
 	}
+	var addr = string(addrBuf)
 
-	addr = string(addrBuf)
-	return tag, addr, nil
+	// 异步建立数据通道
+	go func() {
+		err := data(dataAddr, addr, tag)
+		if err != nil {
+			slog.Error("建立数据通道失败", "err", err)
+		}
+	}()
+
+	return nil
 }
 
 type ConnCmd struct {