From ba3265b0f2cc4c2844831d62f744e150e9fe4e8f Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 4 Dec 2025 21:45:06 +0800 Subject: [PATCH] fix: checkpool deadlock --- core/pool/brutepool.go | 4 +-- core/pool/checkpool.go | 60 +++++++++++++++++++----------------------- core/pool/pool.go | 27 +++++++++---------- spray.go | 11 +++++--- 4 files changed, 49 insertions(+), 53 deletions(-) diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index 14786c4..56e117e 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -48,9 +48,9 @@ func NewBrutePool(ctx context.Context, config *Config) (*BrutePool, error) { Timeout: config.Timeout, ProxyClient: config.ProxyClient, }), - additionCh: make(chan *Unit, config.Thread), + additionCh: make(chan *Unit, config.Thread*10), closeCh: make(chan struct{}), - processCh: make(chan *baseline.Baseline, config.Thread), + processCh: make(chan *baseline.Baseline, config.Thread*2), wg: &sync.WaitGroup{}, }, base: u.Scheme + "://" + u.Host, diff --git a/core/pool/checkpool.go b/core/pool/checkpool.go index 136d460..554b20a 100644 --- a/core/pool/checkpool.go +++ b/core/pool/checkpool.go @@ -31,9 +31,9 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) { ProxyClient: config.ProxyClient, }), wg: &sync.WaitGroup{}, - additionCh: make(chan *Unit, 1024), + additionCh: make(chan *Unit, config.Thread*10), closeCh: make(chan struct{}), - processCh: make(chan *baseline.Baseline, config.Thread), + processCh: make(chan *baseline.Baseline, config.Thread*2), }, } pool.Headers.Set("Connection", "close") @@ -93,8 +93,10 @@ Loop: case <-pool.closeCh: break Loop case <-ctx.Done(): + // 手动退出，不等待任务完成，直接退出 break Loop case <-pool.ctx.Done(): + // 手动退出，不等待任务完成，直接退出 break Loop } } @@ -156,22 +158,20 @@ func (pool *CheckPool) Invoke(v interface{}) { bl.ReqDepth = unit.depth bl.Source = unit.source bl.Spended = time.Since(start).Milliseconds() + if bl.RedirectURL != "" { + pool.doRedirect(bl, bl.ReqDepth) + } pool.processCh <- bl } func (pool *CheckPool) Handler() { for bl := range pool.processCh { if bl.IsValid { - if bl.RedirectURL != "" { - pool.doRedirect(bl, bl.ReqDepth) - pool.putToOutput(bl) - } else { - params := map[string]interface{}{ - "current": bl, - } - if pool.MatchExpr != nil && pkg.CompareWithExpr(pool.MatchExpr, params) { - bl.IsValid = true - } + params := map[string]interface{}{ + "current": bl, + } + if pool.MatchExpr != nil && pkg.CompareWithExpr(pool.MatchExpr, params) { + bl.IsValid = true } } if bl.Source == parsers.CheckSource { @@ -196,17 +196,14 @@ func (pool *CheckPool) doRedirect(bl *baseline.Baseline, depth int) { reURL = pkg.BaseURL(bl.Url) + pkg.FormatURL(pkg.BaseURL(bl.Url), bl.RedirectURL) } - pool.wg.Add(1) - go func() { - pool.additionCh <- &Unit{ - path: reURL, - parent: bl.Number, - source: parsers.RedirectSource, - frontUrl: bl.UrlString, - depth: depth + 1, - from: bl.Source, - } - }() + pool.addAddition(&Unit{ + path: reURL, + parent: bl.Number, + source: parsers.RedirectSource, + frontUrl: bl.UrlString, + depth: depth + 1, + from: bl.Source, + }) } // tcp与400进行协议转换 @@ -214,20 +211,17 @@ func (pool *CheckPool) doUpgrade(bl *baseline.Baseline) { if bl.ReqDepth >= 1 { return } - pool.wg.Add(1) var reurl string if strings.HasPrefix(bl.UrlString, "https") { reurl = strings.Replace(bl.UrlString, "https", "http", 1) } else { reurl = strings.Replace(bl.UrlString, "http", "https", 1) } - go func() { - pool.additionCh <- &Unit{ - path: reurl, - parent: bl.Number, - source: parsers.UpgradeSource, - depth: bl.ReqDepth + 1, - from: bl.Source, - } - }() + pool.addAddition(&Unit{ + path: reurl, + parent: bl.Number, + source: parsers.UpgradeSource, + depth: bl.ReqDepth + 1, + from: bl.Source, + }) } diff --git a/core/pool/pool.go b/core/pool/pool.go index 36b58c0..6e44dd2 100644 --- a/core/pool/pool.go +++ b/core/pool/pool.go @@ -33,28 +33,27 @@ func (pool *BasePool) doRetry(bl *baseline.Baseline) { if bl.Retry >= pool.RetryLimit { return } - pool.wg.Add(1) - go func() { - defer pool.wg.Done() - pool.addAddition(&Unit{ - path: bl.Path, - parent: bl.Number, - host: bl.Host, - source: parsers.RetrySource, - from: bl.Source, - retry: bl.Retry + 1, - }) - }() + pool.addAddition(&Unit{ + path: bl.Path, + parent: bl.Number, + host: bl.Host, + source: parsers.RetrySource, + from: bl.Source, + retry: bl.Retry + 1, + }) } func (pool *BasePool) addAddition(u *Unit) { // 强行屏蔽报错, 防止goroutine泄露 - pool.wg.Add(1) defer func() { if err := recover(); err != nil { } }() - pool.additionCh <- u + pool.wg.Add(1) + select { + case pool.additionCh <- u: + case <-pool.ctx.Done(): + } } func (pool *BasePool) putToOutput(bl *baseline.Baseline) { diff --git a/spray.go b/spray.go index c57a737..57b1d3c 100644 --- a/spray.go +++ b/spray.go @@ -5,7 +5,8 @@ import ( "github.com/chainreactors/spray/cmd" "github.com/gookit/config/v2" "github.com/gookit/config/v2/yaml" - //_ "net/http/pprof" + "net/http" + _ "net/http/pprof" ) func init() { @@ -17,8 +18,10 @@ func init() { } func main() { - //f, _ := os.Create("cpu.txt") - //pprof.StartCPUProfile(f) - //defer pprof.StopCPUProfile() + // 启动 pprof HTTP 服务器 + go func() { + http.ListenAndServe("localhost:6060", nil) + }() + cmd.Spray() }