From 4687a7cf1bcccd1da598c595718c623817119c1c Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Fri, 23 Sep 2022 01:39:00 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E5=91=A8=E6=9C=9F=E6=80=A7wa?=
 =?UTF-8?q?f=E6=A3=80=E6=9F=A5=E5=92=8C=E8=BF=9E=E9=80=9A=E6=80=A7?=
 =?UTF-8?q?=E6=A3=80=E6=9F=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/pool.go | 56 +++++++++++++++++++++++++++++++++---------------
 1 file changed, 39 insertions(+), 17 deletions(-)

diff --git a/internal/pool.go b/internal/pool.go
index ce0c54f..65e5ce7 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -21,13 +21,15 @@ var (
 func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) {
 	pctx, cancel := context.WithCancel(ctx)
 	pool := &Pool{
-		Config:   config,
-		ctx:      pctx,
-		client:   pkg.NewClient(config.Thread, 2),
-		worder:   words.NewWorder(config.Wordlist),
-		outputCh: outputCh,
-		tempCh:   make(chan *baseline, config.Thread),
-		wg:       &sync.WaitGroup{},
+		Config:      config,
+		ctx:         pctx,
+		client:      pkg.NewClient(config.Thread, 2),
+		worder:      words.NewWorder(config.Wordlist),
+		outputCh:    outputCh,
+		tempCh:      make(chan *baseline, config.Thread),
+		wg:          &sync.WaitGroup{},
+		checkPeriod: 100,
+		errPeriod:   10,
 	}
 
 	switch config.Mod {
@@ -71,7 +73,15 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 		switch unit.source {
 		case CheckSource:
-			pool.baseline = bl
+			if pool.baseline == nil {
+				//初次check覆盖baseline
+				pool.baseline = bl
+			} else if bl.Err != nil {
+				logs.Log.Warn("maybe ip banned by waf")
+			} else if !pool.baseline.Equal(bl) {
+				logs.Log.Warn("maybe trigger risk control")
+			}
+
 		case WordSource:
 			// 异步进行性能消耗较大的深度对比
 			pool.tempCh <- bl
@@ -95,23 +105,28 @@ type Pool struct {
 	bar    *pkg.Bar
 	ctx    context.Context
 	//baseReq      *http.Request
-	baseline   *baseline
-	outputCh   chan *baseline
-	tempCh     chan *baseline
-	totalCount int
-	errorCount int
-	genReq     func(s string) (*fasthttp.Request, error)
+	baseline    *baseline
+	outputCh    chan *baseline
+	tempCh      chan *baseline
+	reqCount    int
+	errorCount  int
+	checkPeriod int
+	errPeriod   int
+	genReq      func(s string) (*fasthttp.Request, error)
 	//wordlist     []string
 	worder *words.Worder
 	wg     *sync.WaitGroup
 }
 
-func (p *Pool) Init() error {
-	//for i := 0; i < p.baseReqCount; i++ {
+func (p *Pool) check() {
 	p.wg.Add(1)
 	_ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
 	//}
 	p.wg.Wait()
+}
+
+func (p *Pool) Init() error {
+	p.check()
 	// todo 分析baseline
 	// 检测基本访问能力
 
@@ -119,6 +134,8 @@ func (p *Pool) Init() error {
 		return p.baseline.Err
 	}
 
+	p.baseline.Collect()
+
 	if p.baseline.RedirectURL != "" {
 		CheckRedirect = func(redirectURL string) bool {
 			if redirectURL == p.baseline.RedirectURL {
@@ -143,8 +160,13 @@ Loop:
 			if !ok {
 				break Loop
 			}
-			p.totalCount++
+			p.reqCount++
 			p.wg.Add(1)
+			if p.reqCount%p.checkPeriod == 0 {
+				go p.check()
+			} else if p.reqCount%p.errPeriod == 0 {
+				go p.check()
+			}
 			_ = p.pool.Invoke(newUnit(u, WordSource))
 		case <-time.NewTimer(time.Duration(p.DeadlineTime) * time.Second).C:
 			break Loop