fix: checkpool deadlock

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,

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,
+	})
 }

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) {

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()
 }