优化数据连接处理逻辑,避免数据通道连接空等待问题;授权部分添加全局白名单支持;现在节点丢失连接后不会清空授权数据

This commit is contained in:
2025-05-23 18:59:53 +08:00
parent 12038f0af7
commit 52907b3fae
10 changed files with 170 additions and 92 deletions

View File

@@ -3,19 +3,19 @@ package fwd
import (
"bufio"
"context"
"encoding/hex"
"errors"
"fmt"
"github.com/google/uuid"
"io"
"log/slog"
"net"
"proxy-server/gateway/app"
"proxy-server/gateway/core"
"proxy-server/gateway/debug"
"proxy-server/gateway/env"
"proxy-server/gateway/fwd/metrics"
"proxy-server/utils"
"strconv"
"sync"
"time"
)
@@ -60,7 +60,10 @@ func ListenData(ctx context.Context) error {
app.DataConnWg.Add(1)
go func() {
defer app.DataConnWg.Done()
defer utils.Close(conn)
defer func() {
utils.Close(conn)
slog.Debug("关闭数据通道连接")
}()
err := processDataConn(ctx, conn)
if err != nil {
slog.Error("处理数据通道连接失败", "err", err)
@@ -80,88 +83,107 @@ func processDataConn(ctx context.Context, client net.Conn) error {
return fmt.Errorf("从节点获取连接结果失败: %w", err)
}
tag := buf[0:16]
tag := hex.EncodeToString(buf[0:16])
status := buf[16]
// 加载用户连接
var tagStr = uuid.UUID(tag).String()
user, ok := app.UserConnMap.LoadAndDelete(tagStr)
user, ok := app.UserConnMap.LoadAndDelete(tag)
if !ok {
return fmt.Errorf("用户连接已关闭tag%s", tagStr)
return fmt.Errorf("用户连接已关闭tag%s", tag)
}
defer utils.Close(user)
defer func() {
utils.Close(user)
slog.Debug("关闭用户连接")
}()
// 检查状态
if status != 1 {
return errors.New("目标地址建立连接失败")
}
// 转发数据
data := time.Now()
userPipeReader, userPipeWriter := io.Pipe()
defer utils.Close(userPipeWriter)
// 复制用户流量进行访问目标分析
userCopyFrom, userCopyTo := io.Pipe()
defer utils.Close(userCopyTo)
teeUser := io.TeeReader(user, userPipeWriter)
teeUser := io.TeeReader(user, userCopyTo)
go func() {
err := analysisAndLog(user, userPipeReader)
err := analysisAndLog(user, userCopyFrom)
if err != nil {
slog.Error("数据解析失败", "err", err)
}
}()
wg := sync.WaitGroup{}
wg.Add(2)
// 复制节点数据到用户
var waitEdge = make(chan error)
go func() {
defer wg.Done()
_, err := io.Copy(client, teeUser)
if err != nil {
slog.Error("数据转发失败 user->client", "err", err)
}
}()
go func() {
defer wg.Done()
_, err := io.Copy(user, reader)
if err != nil {
slog.Error("数据转发失败 client->user", "err", err)
switch {
case errors.Is(err, net.ErrClosed):
slog.Debug("节点连接意外关闭")
case err != nil:
slog.Error("读取节点数据失败", "err", err)
default:
slog.Debug("节点数据读取完成")
}
waitEdge <- err
}()
// 复制用户数据到节点
var waitUser = make(chan error)
go func() {
_, err := io.Copy(client, teeUser)
switch {
case errors.Is(err, net.ErrClosed):
slog.Debug("用户连接意外关闭")
case err != nil:
slog.Error("读取用户数据失败", "err", err)
default:
slog.Debug("用户数据读取完成")
}
waitUser <- err
}()
// 等待数据转发完成,关闭数据通道的时机:
select {
case <-ctx.Done():
return nil
slog.Debug("服务关闭")
case <-waitEdge:
case <-waitUser:
storeConnMatrics(user, data)
}
case <-utils.WgWait(&wg):
proxy := time.Now()
return nil
}
start, startOk := metrics.TimerStart.Load(user.Conn)
auth, authOk := metrics.TimerAuth.Load(user.Conn)
func storeConnMatrics(user *core.Conn, data time.Time) {
proxy := time.Now()
var authDuration time.Duration
if startOk && authOk {
authDuration = auth.(time.Time).Sub(start.(time.Time))
}
start, startOk := metrics.TimerStart.Load(user.Conn)
auth, authOk := metrics.TimerAuth.Load(user.Conn)
var dataDuration time.Duration
if authOk {
dataDuration = data.Sub(auth.(time.Time))
}
var authDuration time.Duration
if startOk && authOk {
authDuration = auth.(time.Time).Sub(start.(time.Time))
}
proxyDuration := proxy.Sub(data)
var dataDuration time.Duration
if authOk {
dataDuration = data.Sub(auth.(time.Time))
}
var totalDuration time.Duration
if startOk {
totalDuration = proxy.Sub(start.(time.Time))
}
proxyDuration := proxy.Sub(data)
debug.ConsumingCh <- debug.Consuming{
Auth: authDuration,
Data: dataDuration,
Proxy: proxyDuration,
Total: totalDuration,
}
var totalDuration time.Duration
if startOk {
totalDuration = proxy.Sub(start.(time.Time))
}
return nil
debug.ConsumingCh <- debug.Consuming{
Auth: authDuration,
Data: dataDuration,
Proxy: proxyDuration,
Total: totalDuration,
}
}