fix: POC 扫描接入调用方 context,修复 cachedPocPath 竞争和 ProxyStats data race

- webscan/web_scan.go: WebScan 接受 ctx 参数,替换 context.Background();
  sync.Once 改为 sync.Mutex 保护 POC 加载,消除 cachedPocPath 并发写竞争
- webtitle.go: ctx 从 Scan 穿透到 identifyFingerprintsMulti → triggerPocScan → WebScan
- webpoc.go: 传递 ctx 到 WebScan
- proxy/types.go: ProxyStats 增加 sync.Mutex
- proxy/manager.go: LastConnectTime/LastError/AverageConnectTime 读写加锁
This commit is contained in:
ZacharyZcR
2026-04-27 19:24:59 +08:00
parent 3bf1cb0376
commit 9b8ed55eab
5 changed files with 46 additions and 23 deletions
+16 -2
View File
@@ -264,11 +264,16 @@ func (d *directDialer) DialContext(ctx context.Context, network, address string)
conn, err := dialer.DialContext(ctx, network, address)
duration := time.Since(start)
d.stats.mu.Lock()
d.stats.LastConnectTime = start
d.stats.mu.Unlock()
if err != nil {
atomic.AddInt64(&d.stats.FailedConnections, 1)
d.stats.mu.Lock()
d.stats.LastError = err.Error()
d.stats.mu.Unlock()
return nil, NewProxyError(ErrTypeConnection, ErrMsgDirectConnFailed, ErrCodeDirectConnFailed, err)
}
@@ -323,15 +328,22 @@ func (s *socks5Dialer) DialContext(ctx context.Context, network, address string)
select {
case <-dialCtx.Done():
atomic.AddInt64(&s.stats.FailedConnections, 1)
s.stats.mu.Lock()
s.stats.LastError = dialCtx.Err().Error()
s.stats.mu.Unlock()
return nil, NewProxyError(ErrTypeTimeout, ErrMsgSOCKS5ConnTimeout, ErrCodeSOCKS5ConnTimeout, dialCtx.Err())
case result := <-connChan:
duration := time.Since(start)
s.stats.mu.Lock()
s.stats.LastConnectTime = start
s.stats.mu.Unlock()
if result.err != nil {
atomic.AddInt64(&s.stats.FailedConnections, 1)
s.stats.mu.Lock()
s.stats.LastError = result.err.Error()
s.stats.mu.Unlock()
return nil, NewProxyError(ErrTypeConnection, ErrMsgSOCKS5ConnFailed, ErrCodeSOCKS5ConnFailed, result.err)
}
@@ -347,7 +359,8 @@ func (s *socks5Dialer) DialContext(ctx context.Context, network, address string)
// updateAverageConnectTime 更新平均连接时间
func (d *directDialer) updateAverageConnectTime(duration time.Duration) {
// 简单的移动平均
d.stats.mu.Lock()
defer d.stats.mu.Unlock()
if d.stats.AverageConnectTime == 0 {
d.stats.AverageConnectTime = duration
} else {
@@ -356,7 +369,8 @@ func (d *directDialer) updateAverageConnectTime(duration time.Duration) {
}
func (s *socks5Dialer) updateAverageConnectTime(duration time.Duration) {
// 简单的移动平均
s.stats.mu.Lock()
defer s.stats.mu.Unlock()
if s.stats.AverageConnectTime == 0 {
s.stats.AverageConnectTime = duration
} else {
+5 -3
View File
@@ -4,6 +4,7 @@ import (
"context"
"crypto/tls"
"net"
"sync"
"time"
)
@@ -95,9 +96,10 @@ type ProxyManager interface {
//
//nolint:revive // 保持与现有代码的向后兼容性
type ProxyStats struct {
TotalConnections int64 `json:"total_connections"`
ActiveConnections int64 `json:"active_connections"`
FailedConnections int64 `json:"failed_connections"`
TotalConnections int64 `json:"total_connections"`
ActiveConnections int64 `json:"active_connections"`
FailedConnections int64 `json:"failed_connections"`
mu sync.Mutex `json:"-"`
AverageConnectTime time.Duration `json:"average_connect_time"`
LastConnectTime time.Time `json:"last_connect_time"`
LastError string `json:"last_error,omitempty"`
+1 -1
View File
@@ -106,7 +106,7 @@ func (p *WebPocPlugin) Scan(ctx context.Context, info *common.HostInfo, config *
// 全量模式:忽略指纹和CDN/WAF检测,直接扫描所有POC
target := info.Target()
common.LogDebug(fmt.Sprintf("WebPOC %s 全量扫描模式", target))
WebScan.WebScan(info, config)
WebScan.WebScan(ctx, info, config)
return &WebScanResult{
Type: plugins.ResultTypeWeb,
+5 -5
View File
@@ -159,7 +159,7 @@ func (p *WebTitlePlugin) getWebTitle(ctx context.Context, info *common.HostInfo,
}
// 执行指纹识别(合并原始响应和跳转后响应的指纹)
fingerprints := p.identifyFingerprintsMulti(info, baseURL, checkDataList, config)
fingerprints := p.identifyFingerprintsMulti(ctx, info, baseURL, checkDataList, config)
return title, statusCode, contentLen, server, fingerprints, displayURL, nil
}
@@ -188,20 +188,20 @@ func (p *WebTitlePlugin) resolveRedirectURL(baseURL, location string) string {
}
// identifyFingerprintsMulti 识别多个响应的指纹并合并
func (p *WebTitlePlugin) identifyFingerprintsMulti(info *common.HostInfo, baseURL string, checkDataList []WebScan.CheckDatas, config *common.Config) []string {
func (p *WebTitlePlugin) identifyFingerprintsMulti(ctx context.Context, info *common.HostInfo, baseURL string, checkDataList []WebScan.CheckDatas, config *common.Config) []string {
// 调用指纹识别
fingerprints := WebScan.InfoCheck(baseURL, &checkDataList)
// 非全量模式下,基于指纹触发POC扫描
if !config.POC.Full && !config.POC.Disabled {
p.triggerPocScan(info, fingerprints, config)
p.triggerPocScan(ctx, info, fingerprints, config)
}
return fingerprints
}
// triggerPocScan 基于指纹触发POC扫描
func (p *WebTitlePlugin) triggerPocScan(info *common.HostInfo, fingerprints []string, config *common.Config) {
func (p *WebTitlePlugin) triggerPocScan(ctx context.Context, info *common.HostInfo, fingerprints []string, config *common.Config) {
target := info.Target()
// 无指纹,跳过
@@ -219,7 +219,7 @@ func (p *WebTitlePlugin) triggerPocScan(info *common.HostInfo, fingerprints []st
// 基于指纹执行POC扫描
common.LogDebug(fmt.Sprintf("WebTitle %s 触发指纹POC扫描: %v", target, fingerprints))
info.Info = fingerprints
WebScan.WebScan(info, config)
WebScan.WebScan(ctx, info, config)
}
// formatHeaders 将 HTTP Header 格式化为字符串
+19 -12
View File
@@ -40,21 +40,25 @@ var (
//go:embed pocs
var pocsFS embed.FS
var (
once sync.Once
allPocs []*lib.Poc
cachedPocPath string // 缓存POC路径,用于initPocs
pocMu sync.Mutex
pocLoaded bool
allPocs []*lib.Poc
cachedPocPath string
)
// WebScan 执行Web漏洞扫描
func WebScan(info *common.HostInfo, cfg *common.Config) {
func WebScan(ctx context.Context, info *common.HostInfo, cfg *common.Config) {
// 初始化POC配置(用于CEL回调函数)
lib.InitPOCConfig(cfg.DNSLog)
// 缓存POC路径供initPocs使用
cachedPocPath = cfg.POC.PocPath
// 初始化POC
once.Do(initPocs)
// 加载POC(互斥保护,避免并发 race
pocMu.Lock()
if !pocLoaded {
cachedPocPath = cfg.POC.PocPath
initPocs()
pocLoaded = true
}
pocMu.Unlock()
// 验证输入
if info == nil {
@@ -74,9 +78,12 @@ func WebScan(info *common.HostInfo, cfg *common.Config) {
return
}
// 使用带超时的上下文
ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
defer cancel()
// 超时兜底:如果调用方 ctx 没有 deadline,加一个默认超时
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, defaultTimeout)
defer cancel()
}
// 根据扫描策略执行POC
if cfg.POC.PocName == "" && len(info.Info) == 0 {