From 7d67e81cb0ae27b2939cda311cc1fe30520a34ab Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Thu, 17 Nov 2022 05:48:46 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96basecompare=E9=80=BB=E8=BE=91?=
 =?UTF-8?q?,=20=E5=8A=A0=E5=85=A5index=20baseline=E7=9A=84=E5=AF=B9?=
 =?UTF-8?q?=E6=AF=94?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/pool.go  | 39 ++++++++++++++++++++++++++++++---------
 internal/types.go |  3 ++-
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/internal/pool.go b/internal/pool.go
index d21fa77..a1f4c45 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -89,7 +89,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			bl = &pkg.Baseline{Url: pool.BaseURL + unit.path, IsValid: false, Err: reqerr.Error(), Reason: ErrRequestFailed.Error()}
 			pool.failedBaselines = append(pool.failedBaselines, bl)
 		} else {
-			if err = pool.PreCompare(resp); unit.source == CheckSource || unit.source == InitSource || err == nil {
+			if err = pool.PreCompare(resp); unit.source != WordSource || err == nil {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
 				bl = pkg.NewBaseline(req.URI(), req.Host(), resp)
 				pool.addFuzzyBaseline(bl)
@@ -99,11 +99,15 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 		}
 
 		switch unit.source {
-		case InitSource:
+		case InitRandomSource:
 			pool.base = bl
 			pool.addFuzzyBaseline(bl)
 			pool.initwg.Done()
 			return
+		case InitIndexSource:
+			pool.index = bl
+			pool.initwg.Done()
+			return
 		case CheckSource:
 			if bl.Err != "" {
 				logs.Log.Warnf("[check.error] maybe ip had banned by waf, break (%d/%d), error: %s", pool.failedCount, pool.BreakThreshold, bl.Err)
@@ -167,6 +171,7 @@ type Pool struct {
 	failedCount     int
 	failedBaselines []*pkg.Baseline
 	base            *pkg.Baseline
+	index           *pkg.Baseline
 	baselines       map[int]*pkg.Baseline
 	analyzeDone     bool
 	genReq          func(s string) (*ihttp.Request, error)
@@ -178,8 +183,9 @@ type Pool struct {
 }
 
 func (p *Pool) Init() error {
-	p.initwg.Add(1)
-	p.pool.Invoke(newUnit(pkg.RandPath(), InitSource))
+	p.initwg.Add(2)
+	p.pool.Invoke(newUnit(pkg.RandPath(), InitRandomSource))
+	p.pool.Invoke(newUnit("/", InitIndexSource))
 	p.initwg.Wait()
 	// todo 分析baseline
 	// 检测基本访问能力
@@ -188,8 +194,16 @@ func (p *Pool) Init() error {
 		return fmt.Errorf(p.base.String())
 	}
 
+	if p.index.Err != "" {
+		return fmt.Errorf(p.index.String())
+	}
+
 	p.base.Collect()
-	logs.Log.Important("[baseline.init] " + p.base.String())
+	p.index.Collect()
+
+	logs.Log.Important("[baseline.random] " + p.base.String())
+	logs.Log.Important("[baseline.index] " + p.base.String())
+
 	if p.base.RedirectURL != "" {
 		CheckRedirect = func(redirectURL string) bool {
 			if redirectURL == p.base.RedirectURL {
@@ -270,10 +284,17 @@ func (p *Pool) BaseCompare(bl *pkg.Baseline) {
 	}
 	var status = -1
 	base, ok := p.baselines[bl.Status] // 挑选对应状态码的baseline进行compare
-	if !ok && p.base.Status == bl.Status {
-		// 当other的状态码与base相同时, 会使用base
-		ok = true
-		base = p.base
+	if !ok {
+		if p.base.Status == bl.Status {
+			// 当other的状态码与base相同时, 会使用base
+			ok = true
+			base = p.base
+		} else if p.index.Status == bl.Status {
+			// 当other的状态码与index相同时, 会使用index
+			ok = true
+			base = p.index
+		}
+
 	}
 
 	if ok {
diff --git a/internal/types.go b/internal/types.go
index 0001589..123429b 100644
--- a/internal/types.go
+++ b/internal/types.go
@@ -31,7 +31,8 @@ type sourceType int
 
 const (
 	CheckSource sourceType = iota + 1
-	InitSource
+	InitRandomSource
+	InitIndexSource
 	WordSource
 	WafSource
 )