From c94a4c4d107a5a3041a6c7d437ebcb56a1b0071a Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Thu, 10 Nov 2022 17:19:05 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=A4=9A=E4=B8=AAbug=20?=
 =?UTF-8?q?=E6=96=B0=E5=A2=9E--force=E5=8F=82=E6=95=B0,=20=E5=B0=86?=
 =?UTF-8?q?=E5=BF=BD=E7=95=A5=E6=8A=A5=E9=94=99=E6=95=B0=E9=87=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/baseline.go |  2 ++
 internal/option.go   |  9 ++++++++-
 internal/pool.go     |  9 ++++++---
 internal/runner.go   | 17 ++++++++++++-----
 4 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/internal/baseline.go b/internal/baseline.go
index d382ea4..86aabf4 100644
--- a/internal/baseline.go
+++ b/internal/baseline.go
@@ -49,6 +49,8 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
 		bl.Host = host
 	}
 
+	bl.Body = resp.Body()
+	bl.BodyLength = resp.ContentLength()
 	bl.RedirectURL = string(resp.GetHeader("Location"))
 
 	return bl
diff --git a/internal/option.go b/internal/option.go
index 33cecd3..f515b22 100644
--- a/internal/option.go
+++ b/internal/option.go
@@ -48,6 +48,7 @@ type RequestOptions struct {
 	Headers []string `long:"header"`
 	Method  string   `long:"method"`
 	Cookie  string   `long:"cookie"`
+	Force   bool     `long:"force"`
 }
 
 type MiscOptions struct {
@@ -73,7 +74,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		PoolSize: opt.PoolSize,
 		Mod:      opt.Mod,
 		Timeout:  opt.Timeout,
-		Probes:   strings.Split(opt.OutputProbe, ","),
 		Deadline: opt.Deadline,
 		Offset:   opt.Offset,
 		Limit:    opt.Limit,
@@ -94,6 +94,10 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		logs.Log.Writer = r.Progress.Bypass()
 	}
 
+	if opt.Force {
+		breakThreshold = 999999
+	}
+
 	// prepare url
 	var urls []string
 	var file *os.File
@@ -217,6 +221,9 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		}
 	}
 
+	if opt.OutputProbe != "" {
+		r.Probes = strings.Split(opt.OutputProbe, ",")
+	}
 	return r, nil
 }
 
diff --git a/internal/pool.go b/internal/pool.go
index 96b5ba2..5c72dd4 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -19,7 +19,7 @@ var (
 	CheckWaf        func([]byte) bool
 )
 
-var breakThreshold int = 10
+var breakThreshold int = 20
 
 func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) {
 	pctx, cancel := context.WithCancel(ctx)
@@ -35,6 +35,8 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 		initwg:      sync.WaitGroup{},
 		checkPeriod: 100,
 		errPeriod:   10,
+		reqCount:    1,
+		failedCount: 1,
 	}
 
 	switch config.Mod {
@@ -87,6 +89,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 		if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge {
 			pool.failedCount++
 			bl = &baseline{Url: pool.BaseURL + unit.path, Err: reqerr}
+			pool.failedBaselines = append(pool.failedBaselines, bl)
 		} else {
 			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource || unit.source == InitSource {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
@@ -271,8 +274,8 @@ func (p *Pool) ResetFailed() {
 func (p *Pool) Recover() {
 	logs.Log.Errorf("failed request exceeds the threshold , task will exit. Breakpoint %d", p.reqCount)
 	logs.Log.Error("collecting failed check")
-	for _, bl := range p.failedBaselines {
-		logs.Log.Error(bl.String())
+	for i, bl := range p.failedBaselines {
+		logs.Log.Errorf("[failed.%d] %s", i, bl.String())
 	}
 }
 
diff --git a/internal/runner.go b/internal/runner.go
index f99a4de..73e1bfd 100644
--- a/internal/runner.go
+++ b/internal/runner.go
@@ -113,15 +113,22 @@ Loop:
 }
 
 func (r *Runner) Outputting() {
+	var outFunc func(baseline2 *baseline)
+	if len(r.Probes) > 0 {
+		outFunc = func(bl *baseline) {
+			logs.Log.Console("[+] " + bl.Format(r.Probes) + "\n")
+		}
+	} else {
+		outFunc = func(bl *baseline) {
+			logs.Log.Console("[+] " + bl.String() + "\n")
+		}
+	}
+
 	for {
 		select {
 		case bl := <-r.OutputCh:
 			if bl.IsValid {
-				if len(r.Probes) > 0 {
-					logs.Log.Console("[+]" + bl.Format(r.Probes) + "\n")
-				} else {
-					logs.Log.Console("[+] " + bl.String() + "\n")
-				}
+				outFunc(bl)
 			} else {
 				logs.Log.Debug(bl.String())
 			}