优化连接处理逻辑,增加超时设置;重构命令读取与错误处理;新增公共工具函数以简化错误处理
This commit is contained in:
@@ -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) {
|
||||
|
||||
@@ -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("用户连接意外关闭")
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user