fix the bug of thread pool hanging

internal/option.go

@@ -349,7 +349,7 @@ func (opt *Option) NewRunner() (*Runner, error) {
 	}
 
 	opt.PrintPlugin()
-	if r.IsCheck == false {
+	if r.IsCheck == true {
 		logs.Log.Important("enabling brute mod, because of enabled brute plugin")
 	}
 

internal/pool/brutepool.go

@@ -51,7 +51,7 @@ func NewBrutePool(ctx context.Context, config *Config) (*BrutePool, error) {
 			additionCh: make(chan *Unit, config.Thread),
 			closeCh:    make(chan struct{}),
 			processCh:  make(chan *pkg.Baseline, config.Thread),
-			wg:         sync.WaitGroup{},
+			wg:         &sync.WaitGroup{},
 		},
 		base:  u.Scheme + "://" + u.Host,
 		isDir: strings.HasSuffix(u.Path, "/"),
@@ -196,7 +196,7 @@ func (pool *BrutePool) Upgrade(bl *pkg.Baseline) error {
 	return nil
 }
 
-func (pool *BrutePool) Run(offset, limit int) {
+func (pool *BrutePool) Run(ctx context.Context, offset, limit int) {
 	pool.Worder.Run()
 	if pool.Active {
 		pool.wg.Add(1)
@@ -279,7 +279,7 @@ Loop:
 			}
 		case <-pool.closeCh:
 			break Loop
-		case <-pool.ctx.Done():
+		case <-ctx.Done():
 			break Loop
 		case <-pool.ctx.Done():
 			break Loop

internal/pool/checkpool.go

@@ -29,7 +29,7 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) {
 				Timeout:   time.Duration(config.Timeout) * time.Second,
 				ProxyAddr: config.ProxyAddr,
 			}),
-			wg:         sync.WaitGroup{},
+			wg:         &sync.WaitGroup{},
 			additionCh: make(chan *Unit, 1024),
 			closeCh:    make(chan struct{}),
 			processCh:  make(chan *pkg.Baseline, config.Thread),
@@ -50,21 +50,35 @@ type CheckPool struct {
 func (pool *CheckPool) Run(ctx context.Context, offset, limit int) {
 	pool.Worder.Run()
 
+	var done bool
+	// 挂起一个监控goroutine, 每100ms判断一次done, 如果已经done, 则关闭closeCh, 然后通过Loop中的select case closeCh去break, 实现退出
+	go func() {
+		for {
+			if done {
+				pool.wg.Wait()
+				close(pool.closeCh)
+				return
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+	}()
+
 Loop:
 	for {
 		select {
 		case u, ok := <-pool.Worder.C:
 			if !ok {
-				break Loop
+				done = true
+				continue
 			}
 
 			if pool.reqCount < offset {
 				pool.reqCount++
-				break Loop
+				continue
 			}
 
 			if pool.reqCount > limit {
-				break Loop
+				continue
 			}
 
 			pool.wg.Add(1)
@@ -82,7 +96,7 @@ Loop:
 			break Loop
 		}
 	}
-	pool.wg.Wait()
+
 	pool.Close()
 }
 
@@ -128,6 +142,9 @@ func (pool *CheckPool) Invoke(v interface{}) {
 	} else {
 		bl = pkg.NewBaseline(req.URI(), req.Host(), resp)
 		bl.Collect()
+		if bl.Status == 400 {
+			pool.doUpgrade(bl)
+		}
 	}
 	bl.ReqDepth = unit.depth
 	bl.Source = unit.source
@@ -141,9 +158,6 @@ func (pool *CheckPool) Handler() {
 			if bl.RedirectURL != "" {
 				pool.doRedirect(bl, bl.ReqDepth)
 				pool.putToOutput(bl)
-			} else if bl.Status == 400 {
-				pool.doUpgrade(bl)
-				pool.putToOutput(bl)
 			} else {
 				params := map[string]interface{}{
 					"current": bl,

internal/pool/pool.go

@@ -29,7 +29,7 @@ type BasePool struct {
 	failedCount int
 	additionCh  chan *Unit
 	closeCh     chan struct{}
-	wg          sync.WaitGroup
+	wg          *sync.WaitGroup
 }
 
 func (pool *BasePool) doRedirect(bl *pkg.Baseline, depth int) {

internal/runner.go

@@ -207,7 +207,7 @@ func (r *Runner) Prepare(ctx context.Context) error {
 				}
 			}
 
-			brutePool.Run(brutePool.Statistor.Offset, limit)
+			brutePool.Run(ctx, brutePool.Statistor.Offset, limit)
 
 			if brutePool.IsFailed && len(brutePool.FailedBaselines) > 0 {
 				// 如果因为错误积累退出, end将指向第一个错误发生时, 防止resume时跳过大量目标