fix: use ctx as unified shutdown signal to prevent channel close panics

- Replace additionClosed atomic with ctx-based guard in addAddition
- Add handlerDone chan to coordinate Handler goroutine lifecycle
- Add sendProcess method with select+ctx.Done protection
- Protect putToOutput/putToFuzzy/checkCh sends with select+ctx.Done
- Rewrite BrutePool.Close/CheckPool.Close with correct shutdown order:
  Cancel → Release pools → close(processCh) → <-handlerDone
- Fix doCrawl orphan wg.Add(1) with defer wg.Done()
- Fix analyzeDone inverted condition (replaced with handlerDone chan)
- Add ctx.Done check to Run monitor goroutine

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: M09Ic

core/pool/brutepool.go

@@ -52,10 +52,11 @@ func NewBrutePool(ctx context.Context, config *Config) (*BrutePool, error) {
 				Timeout:     config.Timeout,
 				ProxyClient: config.ProxyClient,
 			}),
-			additionCh: make(chan *Unit, config.Thread*10),
-			closeCh:    make(chan struct{}),
-			processCh:  make(chan *baseline.Baseline, config.Thread*2),
-			wg:         &sync.WaitGroup{},
+			additionCh:  make(chan *Unit, config.Thread*10),
+			closeCh:     make(chan struct{}),
+			processCh:   make(chan *baseline.Baseline, config.Thread*2),
+			wg:          &sync.WaitGroup{},
+			handlerDone: make(chan struct{}),
 		},
 		base:        u.Scheme + "://" + u.Host,
 		isDir:       strings.HasSuffix(u.Path, "/"),
@@ -105,7 +106,6 @@ type BrutePool struct {
 	urls        sync.Map
 	scopeurls   map[string]struct{}
 	uniques     map[uint16]struct{}
-	analyzeDone bool
 	limiter     *rate.Limiter
 	locker      sync.Mutex
 	scopeLocker sync.Mutex
@@ -197,6 +197,11 @@ func (pool *BrutePool) Run(offset, limit int) {
 				close(pool.closeCh)
 				return
 			}
+			select {
+			case <-pool.ctx.Done():
+				return
+			default:
+			}
 			time.Sleep(100 * time.Millisecond)
 		}
 	}()
@@ -376,7 +381,7 @@ func (pool *BrutePool) Invoke(v interface{}) {
 
 	case parsers.WordSource:
 		// 异步进行性能消耗较大的深度对比
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 		if int(pool.Statistor.ReqTotal)%pool.CheckPeriod == 0 {
 			// 间歇插入check waf的探针
 			pool.doCheck()
@@ -388,9 +393,9 @@ func (pool *BrutePool) Invoke(v interface{}) {
 		pool.Bar.Done()
 	case parsers.RedirectSource:
 		bl.FrontURL = unit.frontUrl
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 	default:
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 	}
 }
 
@@ -432,6 +437,7 @@ func (pool *BrutePool) NoScopeInvoke(v interface{}) {
 }
 
 func (pool *BrutePool) Handler() {
+	defer close(pool.handlerDone)
 	for bl := range pool.processCh {
 		if bl.IsValid {
 			pool.addFuzzyBaseline(bl)
@@ -516,8 +522,6 @@ func (pool *BrutePool) Handler() {
 		}
 		pool.wg.Done()
 	}
-
-	pool.analyzeDone = true
 }
 
 func (pool *BrutePool) checkRedirect(redirectURL string) bool {
@@ -674,16 +678,12 @@ func (pool *BrutePool) fallback() {
 }
 
 func (pool *BrutePool) Close() {
-	for pool.analyzeDone {
-		// 等待缓存的待处理任务完成
-		time.Sleep(time.Duration(100) * time.Millisecond)
-	}
-	pool.additionClosed.Store(true)
-	// additionCh 可能仍有异步 producer（redirect/crawl/retry/append），
-	// 依赖 closeCh/ctx 停止消费循环，不直接关闭 channel
-	pool.Statistor.EndTime = time.Now().Unix()
+	pool.Cancel()
 	pool.reqPool.Release()
 	pool.scopePool.Release()
+	close(pool.processCh)
+	<-pool.handlerDone
+	pool.Statistor.EndTime = time.Now().Unix()
 }
 
 func (pool *BrutePool) safePath(u string) string {
@@ -714,9 +714,15 @@ func (pool *BrutePool) doCheck() {
 	}
 
 	if pool.Mod == HostSpray {
-		pool.checkCh <- struct{}{}
+		select {
+		case pool.checkCh <- struct{}{}:
+		case <-pool.ctx.Done():
+		}
 	} else if pool.Mod == PathSpray {
-		pool.checkCh <- struct{}{}
+		select {
+		case pool.checkCh <- struct{}{}:
+		case <-pool.ctx.Done():
+		}
 	}
 }
 
@@ -756,6 +762,7 @@ func (pool *BrutePool) doCrawl(bl *baseline.Baseline) {
 	pool.doScopeCrawl(bl)
 
 	go func() {
+		defer pool.wg.Done()
 		for _, u := range bl.URLs {
 			if u = pkg.FormatURL(bl.Url.Path, u); u == "" {
 				continue

core/pool/checkpool.go

@@ -34,6 +34,7 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) {
 			additionCh: make(chan *Unit, config.Thread*10),
 			closeCh:    make(chan struct{}),
 			processCh:  make(chan *baseline.Baseline, config.Thread*2),
+			handlerDone: make(chan struct{}),
 		},
 	}
 	pool.Request.Headers.Set("Connection", "close")
@@ -105,6 +106,9 @@ Loop:
 	pool.Close()
 }
 func (pool *CheckPool) Close() {
+	pool.Cancel()
+	close(pool.processCh)
+	<-pool.handlerDone
 	pool.Bar.Close()
 	pool.Pool.Release()
 }
@@ -139,7 +143,7 @@ func (pool *CheckPool) Invoke(v interface{}) {
 				ReqDepth:  unit.depth,
 			},
 		}
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 		return
 	}
 	start := time.Now()
@@ -172,10 +176,11 @@ func (pool *CheckPool) Invoke(v interface{}) {
 	if bl.RedirectURL != "" {
 		pool.doRedirect(bl, bl.ReqDepth)
 	}
-	pool.processCh <- bl
+	pool.sendProcess(bl)
 }
 
 func (pool *CheckPool) Handler() {
+	defer close(pool.handlerDone)
 	for bl := range pool.processCh {
 		if bl.IsValid {
 			params := map[string]interface{}{

core/pool/pool.go

@@ -21,13 +21,13 @@ type BasePool struct {
 	ctx       context.Context
 	processCh chan *baseline.Baseline // 待处理的baseline
 
-	reqCount       int
-	failedCount    int
-	additionCh     chan *Unit
-	additionClosed atomic.Bool
-	closeCh        chan struct{}
-	wg             *sync.WaitGroup
-	isFallback     atomic.Bool
+	reqCount    int
+	failedCount int
+	additionCh  chan *Unit
+	closeCh     chan struct{}
+	wg          *sync.WaitGroup
+	handlerDone chan struct{}
+	isFallback  atomic.Bool
 }
 
 func (pool *BasePool) doRetry(bl *baseline.Baseline) {
@@ -45,25 +45,21 @@ func (pool *BasePool) doRetry(bl *baseline.Baseline) {
 }
 
 func (pool *BasePool) addAddition(u *Unit) {
-	if pool.ctx.Err() != nil || pool.additionClosed.Load() {
+	if pool.ctx.Err() != nil {
 		return
 	}
-
 	pool.wg.Add(1)
 	select {
+	case pool.additionCh <- u:
 	case <-pool.ctx.Done():
 		pool.wg.Done()
-		return
-	case pool.additionCh <- u:
-		return
-	default:
-		go func() {
-			select {
-			case pool.additionCh <- u:
-			case <-pool.ctx.Done():
-				pool.wg.Done()
-			}
-		}()
+	}
+}
+
+func (pool *BasePool) sendProcess(bl *baseline.Baseline) {
+	select {
+	case pool.processCh <- bl:
+	case <-pool.ctx.Done():
 	}
 }
 
@@ -72,11 +68,19 @@ func (pool *BasePool) putToOutput(bl *baseline.Baseline) {
 		bl.Collect()
 	}
 	pool.Outwg.Add(1)
-	pool.OutputCh <- bl
+	select {
+	case pool.OutputCh <- bl:
+	case <-pool.ctx.Done():
+		pool.Outwg.Done()
+	}
 }
 
 func (pool *BasePool) putToFuzzy(bl *baseline.Baseline) {
 	pool.Outwg.Add(1)
 	bl.IsFuzzy = true
-	pool.FuzzyCh <- bl
+	select {
+	case pool.FuzzyCh <- bl:
+	case <-pool.ctx.Done():
+		pool.Outwg.Done()
+	}
 }