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) conn, err := dialer.DialContext(ctx, network, address)
duration := time.Since(start) duration := time.Since(start)
d.stats.mu.Lock()
d.stats.LastConnectTime = start d.stats.LastConnectTime = start
d.stats.mu.Unlock()
if err != nil { if err != nil {
atomic.AddInt64(&d.stats.FailedConnections, 1) atomic.AddInt64(&d.stats.FailedConnections, 1)
d.stats.mu.Lock()
d.stats.LastError = err.Error() d.stats.LastError = err.Error()
d.stats.mu.Unlock()
return nil, NewProxyError(ErrTypeConnection, ErrMsgDirectConnFailed, ErrCodeDirectConnFailed, err) return nil, NewProxyError(ErrTypeConnection, ErrMsgDirectConnFailed, ErrCodeDirectConnFailed, err)
} }
@@ -323,15 +328,22 @@ func (s *socks5Dialer) DialContext(ctx context.Context, network, address string)
select { select {
case <-dialCtx.Done(): case <-dialCtx.Done():
atomic.AddInt64(&s.stats.FailedConnections, 1) atomic.AddInt64(&s.stats.FailedConnections, 1)
s.stats.mu.Lock()
s.stats.LastError = dialCtx.Err().Error() s.stats.LastError = dialCtx.Err().Error()
s.stats.mu.Unlock()
return nil, NewProxyError(ErrTypeTimeout, ErrMsgSOCKS5ConnTimeout, ErrCodeSOCKS5ConnTimeout, dialCtx.Err()) return nil, NewProxyError(ErrTypeTimeout, ErrMsgSOCKS5ConnTimeout, ErrCodeSOCKS5ConnTimeout, dialCtx.Err())
case result := <-connChan: case result := <-connChan:
duration := time.Since(start) duration := time.Since(start)
s.stats.mu.Lock()
s.stats.LastConnectTime = start s.stats.LastConnectTime = start
s.stats.mu.Unlock()
if result.err != nil { if result.err != nil {
atomic.AddInt64(&s.stats.FailedConnections, 1) atomic.AddInt64(&s.stats.FailedConnections, 1)
s.stats.mu.Lock()
s.stats.LastError = result.err.Error() s.stats.LastError = result.err.Error()
s.stats.mu.Unlock()
return nil, NewProxyError(ErrTypeConnection, ErrMsgSOCKS5ConnFailed, ErrCodeSOCKS5ConnFailed, result.err) return nil, NewProxyError(ErrTypeConnection, ErrMsgSOCKS5ConnFailed, ErrCodeSOCKS5ConnFailed, result.err)
} }
@@ -347,7 +359,8 @@ func (s *socks5Dialer) DialContext(ctx context.Context, network, address string)
// updateAverageConnectTime 更新平均连接时间 // updateAverageConnectTime 更新平均连接时间
func (d *directDialer) updateAverageConnectTime(duration time.Duration) { func (d *directDialer) updateAverageConnectTime(duration time.Duration) {
// 简单的移动平均 d.stats.mu.Lock()
defer d.stats.mu.Unlock()
if d.stats.AverageConnectTime == 0 { if d.stats.AverageConnectTime == 0 {
d.stats.AverageConnectTime = duration d.stats.AverageConnectTime = duration
} else { } else {
@@ -356,7 +369,8 @@ func (d *directDialer) updateAverageConnectTime(duration time.Duration) {
} }
func (s *socks5Dialer) 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 { if s.stats.AverageConnectTime == 0 {
s.stats.AverageConnectTime = duration s.stats.AverageConnectTime = duration
} else { } else {
+2
View File
@@ -4,6 +4,7 @@ import (
"context" "context"
"crypto/tls" "crypto/tls"
"net" "net"
"sync"
"time" "time"
) )
@@ -98,6 +99,7 @@ type ProxyStats struct {
TotalConnections int64 `json:"total_connections"` TotalConnections int64 `json:"total_connections"`
ActiveConnections int64 `json:"active_connections"` ActiveConnections int64 `json:"active_connections"`
FailedConnections int64 `json:"failed_connections"` FailedConnections int64 `json:"failed_connections"`
mu sync.Mutex `json:"-"`
AverageConnectTime time.Duration `json:"average_connect_time"` AverageConnectTime time.Duration `json:"average_connect_time"`
LastConnectTime time.Time `json:"last_connect_time"` LastConnectTime time.Time `json:"last_connect_time"`
LastError string `json:"last_error,omitempty"` 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 // 全量模式:忽略指纹和CDN/WAF检测,直接扫描所有POC
target := info.Target() target := info.Target()
common.LogDebug(fmt.Sprintf("WebPOC %s 全量扫描模式", target)) common.LogDebug(fmt.Sprintf("WebPOC %s 全量扫描模式", target))
WebScan.WebScan(info, config) WebScan.WebScan(ctx, info, config)
return &WebScanResult{ return &WebScanResult{
Type: plugins.ResultTypeWeb, 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 return title, statusCode, contentLen, server, fingerprints, displayURL, nil
} }
@@ -188,20 +188,20 @@ func (p *WebTitlePlugin) resolveRedirectURL(baseURL, location string) string {
} }
// identifyFingerprintsMulti 识别多个响应的指纹并合并 // 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) fingerprints := WebScan.InfoCheck(baseURL, &checkDataList)
// 非全量模式下,基于指纹触发POC扫描 // 非全量模式下,基于指纹触发POC扫描
if !config.POC.Full && !config.POC.Disabled { if !config.POC.Full && !config.POC.Disabled {
p.triggerPocScan(info, fingerprints, config) p.triggerPocScan(ctx, info, fingerprints, config)
} }
return fingerprints return fingerprints
} }
// triggerPocScan 基于指纹触发POC扫描 // 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() target := info.Target()
// 无指纹,跳过 // 无指纹,跳过
@@ -219,7 +219,7 @@ func (p *WebTitlePlugin) triggerPocScan(info *common.HostInfo, fingerprints []st
// 基于指纹执行POC扫描 // 基于指纹执行POC扫描
common.LogDebug(fmt.Sprintf("WebTitle %s 触发指纹POC扫描: %v", target, fingerprints)) common.LogDebug(fmt.Sprintf("WebTitle %s 触发指纹POC扫描: %v", target, fingerprints))
info.Info = fingerprints info.Info = fingerprints
WebScan.WebScan(info, config) WebScan.WebScan(ctx, info, config)
} }
// formatHeaders 将 HTTP Header 格式化为字符串 // formatHeaders 将 HTTP Header 格式化为字符串
+16 -9
View File
@@ -40,21 +40,25 @@ var (
//go:embed pocs //go:embed pocs
var pocsFS embed.FS var pocsFS embed.FS
var ( var (
once sync.Once pocMu sync.Mutex
pocLoaded bool
allPocs []*lib.Poc allPocs []*lib.Poc
cachedPocPath string // 缓存POC路径,用于initPocs cachedPocPath string
) )
// WebScan 执行Web漏洞扫描 // WebScan 执行Web漏洞扫描
func WebScan(info *common.HostInfo, cfg *common.Config) { func WebScan(ctx context.Context, info *common.HostInfo, cfg *common.Config) {
// 初始化POC配置(用于CEL回调函数) // 初始化POC配置(用于CEL回调函数)
lib.InitPOCConfig(cfg.DNSLog) lib.InitPOCConfig(cfg.DNSLog)
// 缓存POC路径供initPocs使用 // 加载POC(互斥保护,避免并发 race)
pocMu.Lock()
if !pocLoaded {
cachedPocPath = cfg.POC.PocPath cachedPocPath = cfg.POC.PocPath
initPocs()
// 初始化POC pocLoaded = true
once.Do(initPocs) }
pocMu.Unlock()
// 验证输入 // 验证输入
if info == nil { if info == nil {
@@ -74,9 +78,12 @@ func WebScan(info *common.HostInfo, cfg *common.Config) {
return return
} }
// 使用带超时的上下文 // 超时兜底:如果调用方 ctx 没有 deadline,加一个默认超时
ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout) if _, hasDeadline := ctx.Deadline(); !hasDeadline {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, defaultTimeout)
defer cancel() defer cancel()
}
// 根据扫描策略执行POC // 根据扫描策略执行POC
if cfg.POC.PocName == "" && len(info.Info) == 0 { if cfg.POC.PocName == "" && len(info.Info) == 0 {