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

This commit is contained in:
2025-05-26 16:37:54 +08:00
parent 8c928a8321
commit c2dcae7af5
7 changed files with 205 additions and 154 deletions

View File

@@ -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) {