enhance runner structure

cmd/cmd.go

@@ -127,13 +127,6 @@ func Spray() {
 	}
 
 	ctx, canceler := context.WithTimeout(context.Background(), time.Duration(runner.Deadline)*time.Second)
-
-	err = runner.Prepare(ctx)
-	if err != nil {
-		logs.Log.Errorf(err.Error())
-		return
-	}
-
 	go func() {
 		exitChan := make(chan os.Signal, 2)
 		signal.Notify(exitChan, os.Interrupt, syscall.SIGTERM)
@@ -154,10 +147,11 @@ func Spray() {
 		}()
 	}()
 
-	if runner.IsCheck {
+	err = runner.Prepare(ctx)
-		runner.RunWithCheck(ctx)
+	if err != nil {
-	} else {
+		logs.Log.Errorf(err.Error())
-		runner.Run(ctx)
+		return
 	}
+
 	time.Sleep(1 * time.Second)
 }

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/chainreactors/fingers v0.0.0-20240716172449-2fc3147b9c2a
 	github.com/chainreactors/logs v0.0.0-20240207121836-c946f072f81f
 	github.com/chainreactors/parsers v0.0.0-20240708072709-07deeece7ce2
-	github.com/chainreactors/utils v0.0.0-20240716182459-e85f2b01ee16
+	github.com/chainreactors/utils v0.0.0-20240805193040-ff3b97aa3c3f
 	github.com/chainreactors/words v0.4.1-0.20240510105042-5ba5c2edc508
 	github.com/expr-lang/expr v1.16.9
 	github.com/gookit/config/v2 v2.2.5

go.sum

@@ -98,6 +98,8 @@ github.com/chainreactors/utils v0.0.0-20240704062557-662d623b74f4/go.mod h1:JA4e
 github.com/chainreactors/utils v0.0.0-20240715080349-d2d0484c95ed/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU=
 github.com/chainreactors/utils v0.0.0-20240716182459-e85f2b01ee16 h1:TCOshCp7PrWqhP/HSAM5kT3VxoOe7EoJbRseyoSX3RM=
 github.com/chainreactors/utils v0.0.0-20240716182459-e85f2b01ee16/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU=
+github.com/chainreactors/utils v0.0.0-20240805193040-ff3b97aa3c3f h1:2NKmadFYP9vCwC0YrazgttFACleOhxScTPzg0i76YAY=
+github.com/chainreactors/utils v0.0.0-20240805193040-ff3b97aa3c3f/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU=
 github.com/chainreactors/words v0.4.1-0.20240510105042-5ba5c2edc508 h1:iT4HWkoZzUAfQYcQMRH8XyrMau9tCVE0zSuFQnkhrqw=
 github.com/chainreactors/words v0.4.1-0.20240510105042-5ba5c2edc508/go.mod h1:DUDx7PdsMEm5PvVhzkFyppzpiUhQb8dOJaWjVc1SMVk=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=

internal/pool/checkpool.go

@@ -50,35 +50,21 @@ 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 {
-				done = true
+				break Loop
-				continue
 			}
 
 			if pool.reqCount < offset {
 				pool.reqCount++
-				continue
+				break Loop
 			}
 
 			if pool.reqCount > limit {
-				continue
+				break Loop
 			}
 
 			pool.wg.Add(1)
@@ -96,7 +82,7 @@ Loop:
 			break Loop
 		}
 	}
-
+	pool.wg.Wait()
 	pool.Close()
 }
 

internal/runner.go

@@ -111,6 +111,7 @@ func (r *Runner) AppendFunction(fn func(string) []string) {
 }
 
 func (r *Runner) Prepare(ctx context.Context) error {
+	r.OutputHandler()
 	var err error
 	if r.IsCheck {
 		// 仅check, 类似httpx
@@ -138,6 +139,7 @@ func (r *Runner) Prepare(ctx context.Context) error {
 			checkPool.Run(ctx, r.Offset, r.Count)
 			r.poolwg.Done()
 		})
+		r.RunWithCheck(ctx)
 	} else {
 		// 完整探测模式
 		go func() {
@@ -214,12 +216,12 @@ func (r *Runner) Prepare(ctx context.Context) error {
 			r.PrintStat(brutePool)
 			r.Done()
 		})
+		r.Run(ctx)
 	}
-
 	if err != nil {
 		return err
 	}
-	r.OutputHandler()
+
 	return nil
 }