From 637eb21377635771343ab8ac5ebeabc7477e75ae Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Fri, 23 Sep 2022 01:47:24 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E6=8A=A5=E9=94=99=E9=98=88?=
 =?UTF-8?q?=E5=80=BC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/baseline.go |  4 ++--
 internal/pool.go     | 20 +++++++++++++++-----
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/internal/baseline.go b/internal/baseline.go
index 6414ca7..bf67a09 100644
--- a/internal/baseline.go
+++ b/internal/baseline.go
@@ -9,7 +9,7 @@ import (
 	"strings"
 )
 
-func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) (*baseline, error) {
+func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
 	bl := &baseline{
 		Url:       u,
 		UrlString: u.String(),
@@ -24,7 +24,7 @@ func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) (*baseline, error) {
 	bl.HeaderLength = resp.Header.Len()
 	bl.RedirectURL = string(resp.Header.Peek("Location"))
 	bl.Raw = append(bl.Header, bl.Body...)
-	return bl, nil
+	return bl
 }
 
 func NewInvalidBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
diff --git a/internal/pool.go b/internal/pool.go
index 65e5ce7..2be75e6 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -18,11 +18,14 @@ var (
 	CheckWaf        func(*http.Response) bool
 )
 
+var breakThreshold int = 10
+
 func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) {
 	pctx, cancel := context.WithCancel(ctx)
 	pool := &Pool{
 		Config:      config,
 		ctx:         pctx,
+		cancel:      cancel,
 		client:      pkg.NewClient(config.Thread, 2),
 		worder:      words.NewWorder(config.Wordlist),
 		outputCh:    outputCh,
@@ -57,15 +60,14 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			//logs.Log.Debugf("%s request error, %s", strurl, err.Error())
 			pool.errorCount++
 			bl = &baseline{Err: err}
+
 		} else {
 			defer fasthttp.ReleaseResponse(resp)
 			defer fasthttp.ReleaseRequest(req)
 			//defer resp.Body.Close() // 必须要关闭body ,否则keep-alive无法生效
 			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
-				bl, err = NewBaseline(req.URI(), resp)
-			} else if err == ErrWaf {
-				cancel()
+				bl = NewBaseline(req.URI(), resp)
 			} else {
 				bl = NewInvalidBaseline(req.URI(), resp)
 			}
@@ -104,12 +106,14 @@ type Pool struct {
 	pool   *ants.PoolWithFunc
 	bar    *pkg.Bar
 	ctx    context.Context
+	cancel context.CancelFunc
 	//baseReq      *http.Request
 	baseline    *baseline
 	outputCh    chan *baseline
 	tempCh      chan *baseline
 	reqCount    int
 	errorCount  int
+	failedCount int
 	checkPeriod int
 	errPeriod   int
 	genReq      func(s string) (*fasthttp.Request, error)
@@ -119,10 +123,15 @@ type Pool struct {
 }
 
 func (p *Pool) check() {
-	p.wg.Add(1)
+	var wg sync.WaitGroup
+	wg.Add(1)
 	_ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
 	//}
-	p.wg.Wait()
+	wg.Wait()
+
+	if p.failedCount > breakThreshold {
+		p.cancel()
+	}
 }
 
 func (p *Pool) Init() error {
@@ -131,6 +140,7 @@ func (p *Pool) Init() error {
 	// 检测基本访问能力
 
 	if p.baseline != nil && p.baseline.Err != nil {
+		p.cancel()
 		return p.baseline.Err
 	}