diff --git a/common/proxy/manager.go b/common/proxy/manager.go index be69619..c25d816 100644 --- a/common/proxy/manager.go +++ b/common/proxy/manager.go @@ -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 { diff --git a/common/proxy/types.go b/common/proxy/types.go index 6614129..adb8c50 100644 --- a/common/proxy/types.go +++ b/common/proxy/types.go @@ -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"` diff --git a/plugins/web/webpoc.go b/plugins/web/webpoc.go index 273089d..1bbb5a5 100644 --- a/plugins/web/webpoc.go +++ b/plugins/web/webpoc.go @@ -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, diff --git a/plugins/web/webtitle.go b/plugins/web/webtitle.go index 9d0ef7b..d4cf3a8 100644 --- a/plugins/web/webtitle.go +++ b/plugins/web/webtitle.go @@ -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 格式化为字符串 diff --git a/webscan/web_scan.go b/webscan/web_scan.go index 25f695f..0246a5d 100644 --- a/webscan/web_scan.go +++ b/webscan/web_scan.go @@ -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 {