优化check逻辑, 减少check发包.

添加相关接口, 为后续的通过dsl自定义过滤规则做准备

internal/pool.go

@@ -30,8 +30,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 		cancel:      cancel,
 		client:      ihttp.NewClient(config.Thread, 2, config.ClientType),
 		worder:      words.NewWorder(config.Wordlist),
-		outputCh:    config.OutputCh,
-		fuzzyCh:     config.FuzzyCh,
 		baselines:   make(map[int]*pkg.Baseline),
 		tempCh:      make(chan *pkg.Baseline, config.Thread),
 		wg:          sync.WaitGroup{},
@@ -131,8 +129,10 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			pool.tempCh <- bl
 
 			if pool.reqCount%pool.checkPeriod == 0 {
+				pool.reqCount++
 				go pool.check()
 			} else if pool.failedCount%pool.errPeriod == 0 {
+				pool.failedCount++
 				go pool.check()
 			}
 			pool.bar.Done()
@@ -142,7 +142,19 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 	})
 
 	pool.pool = p
-	go pool.comparing()
+	go func() {
+		for bl := range pool.tempCh {
+			if pool.customCompare != nil {
+				if pool.customCompare(bl) {
+					pool.OutputCh <- bl
+				}
+			} else {
+				pool.BaseCompare(bl)
+			}
+		}
+
+		pool.analyzeDone = true
+	}()
 	return pool, nil
 }
 
@@ -153,8 +165,6 @@ type Pool struct {
 	bar             *pkg.Bar
 	ctx             context.Context
 	cancel          context.CancelFunc
-	outputCh        chan *pkg.Baseline // 输出的chan, 全局统一
-	fuzzyCh         chan *pkg.Baseline
 	tempCh          chan *pkg.Baseline // 待处理的baseline
 	reqCount        int
 	failedCount     int
@@ -166,6 +176,7 @@ type Pool struct {
 	analyzeDone     bool
 	genReq          func(s string) (*ihttp.Request, error)
 	check           func()
+	customCompare   func(*pkg.Baseline) bool
 	worder          *words.Worder
 	wg              sync.WaitGroup
 	initwg          sync.WaitGroup // 初始化用, 之后改成锁
@@ -259,42 +270,37 @@ func (p *Pool) PreCompare(resp *ihttp.Response) error {
 	return nil
 }
 
-func (p *Pool) comparing() {
-Loop:
-	for bl := range p.tempCh {
-		if !bl.IsValid {
-			// precompare 确认无效数据直接送入管道
-			p.outputCh <- bl
-			continue
+func (p *Pool) BaseCompare(bl *pkg.Baseline) {
+	if !bl.IsValid {
+		// precompare 确认无效数据直接送入管道
+		p.OutputCh <- bl
+		return
+	}
+	var status int
+	base, ok := p.baselines[bl.Status]
+	if ok {
+		// 挑选对应状态码的baseline进行compare
+		if status = base.Compare(bl); status == 1 {
+			p.PutToInvalid(bl, "compare failed")
+			return
 		}
-		var status int
-		base, ok := p.baselines[bl.Status]
-		if ok {
-			// 挑选对应状态码的baseline进行compare
-			if status = base.Compare(bl); status == 1 {
-				p.PutToInvalid(bl, "compare failed")
-				continue
-			}
-		}
-
-		bl.Collect()
-		for _, f := range bl.Frameworks {
-			if f.Tag == "waf/cdn" {
-				p.PutToInvalid(bl, "waf")
-				continue Loop
-			}
-		}
-
-		if status == 0 && ok && base.FuzzyCompare(bl) {
-			p.PutToInvalid(bl, "fuzzy compare failed")
-			p.PutToFuzzy(bl)
-			continue
-		}
-
-		p.outputCh <- bl
 	}
 
-	p.analyzeDone = true
+	bl.Collect()
+	for _, f := range bl.Frameworks {
+		if f.Tag == "waf/cdn" {
+			p.PutToInvalid(bl, "waf")
+			return
+		}
+	}
+
+	if status == 0 && ok && base.FuzzyCompare(bl) {
+		p.PutToInvalid(bl, "fuzzy compare failed")
+		p.PutToFuzzy(bl)
+		return
+	}
+
+	p.OutputCh <- bl
 }
 
 func (p *Pool) addFuzzyBaseline(bl *pkg.Baseline) {
@@ -308,12 +314,12 @@ func (p *Pool) addFuzzyBaseline(bl *pkg.Baseline) {
 func (p *Pool) PutToInvalid(bl *pkg.Baseline, reason string) {
 	bl.IsValid = false
 	bl.Reason = reason
-	p.outputCh <- bl
+	p.OutputCh <- bl
 }
 
 func (p *Pool) PutToFuzzy(bl *pkg.Baseline) {
 	bl.IsFuzzy = true
-	p.fuzzyCh <- bl
+	p.FuzzyCh <- bl
 }
 
 func (p *Pool) resetFailed() {

pkg/baseline.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"github.com/chainreactors/parsers"
 	"github.com/chainreactors/spray/pkg/ihttp"
-	"math"
 	"net/url"
 	"strconv"
 	"strings"
@@ -101,7 +100,7 @@ func (bl *Baseline) Compare(other *Baseline) int {
 		return 1
 	}
 
-	if math.Abs(float64(bl.BodyLength-other.BodyLength)) < 16 {
+	if i := bl.BodyLength - other.BodyLength; i < 16 || i > -16 {
 		// 如果body length相等且md5相等, 则说明是同一个页面
 		if bl.BodyMd5 == parsers.Md5Hash(other.Body) {
 			// 如果length相等, md5也相等, 则判断为全同
@@ -121,10 +120,9 @@ func (bl *Baseline) Compare(other *Baseline) int {
 	return -1
 }
 
-var Distance uint8
+var Distance uint8 = 5
 
 func (bl *Baseline) FuzzyCompare(other *Baseline) bool {
-	// todo 模糊匹配
 	if parsers.SimhashCompare(other.BodySimhash, bl.BodySimhash) < Distance {
 		return true
 	}
@@ -241,3 +239,9 @@ func (bl *Baseline) Jsonify() string {
 	}
 	return string(bs)
 }
+
+func (bl *Baseline) ToMap() map[string]interface{} {
+	return map[string]interface{}{
+		"status": bl.Status,
+	}
+}

pkg/config.go

@@ -19,17 +19,15 @@ var ModMap = map[string]SprayMod{
 }
 
 type Config struct {
-	BaseURL     string
-	Wordlist    []string
-	Thread      int
-	Timeout     int
-	BaseReq     *http.Request
-	Method      string
-	Mod         SprayMod
-	Headers     http.Header
-	EnableFuzzy bool
-	ClientType  int
-	Fns         []func(string) string
-	OutputCh    chan *Baseline
-	FuzzyCh     chan *Baseline
+	BaseURL    string
+	Wordlist   []string
+	Thread     int
+	Timeout    int
+	Method     string
+	Mod        SprayMod
+	Headers    http.Header
+	ClientType int
+	Fns        []func(string) string
+	OutputCh   chan *Baseline
+	FuzzyCh    chan *Baseline
 }