From 9ffb509265ac722a5d1ba61668d2ad2891a976fc Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Tue, 6 Aug 2024 03:58:30 +0800
Subject: [PATCH] enhance runner structure

---
 cmd/cmd.go                 | 16 +++++-----------
 go.mod                     |  2 +-
 go.sum                     |  2 ++
 internal/pool/checkpool.go | 22 ++++------------------
 internal/runner.go         |  6 ++++--
 5 files changed, 16 insertions(+), 32 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 3edd2bc..94f25e0 100644
--- a/cmd/cmd.go
+++ b/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 {
-		runner.RunWithCheck(ctx)
-	} else {
-		runner.Run(ctx)
+	err = runner.Prepare(ctx)
+	if err != nil {
+		logs.Log.Errorf(err.Error())
+		return
 	}
+
 	time.Sleep(1 * time.Second)
 }
diff --git a/go.mod b/go.mod
index 120386e..98faacc 100644
--- a/go.mod
+++ b/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
diff --git a/go.sum b/go.sum
index 7df73e9..aa9c067 100644
--- a/go.sum
+++ b/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=
diff --git a/internal/pool/checkpool.go b/internal/pool/checkpool.go
index a351aa6..45bc087 100644
--- a/internal/pool/checkpool.go
+++ b/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
-				continue
+				break Loop
 			}
 
 			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()
 }
 
diff --git a/internal/runner.go b/internal/runner.go
index c412057..aeeae11 100644
--- a/internal/runner.go
+++ b/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
 }