fix --check-only panic

chainreactors · Mar 7, 2024 · 006e1af · 006e1af
1 parent 3da923b
commit 006e1af
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/internal/pool/brutepool.go b/internal/pool/brutepool.go
@@ -764,16 +764,3 @@ func (pool *BrutePool) resetFailed() {
 	pool.failedCount = 1
 	pool.FailedBaselines = nil
 }
-
-func NewBaselines() *Baselines {
-	return &Baselines{
-		baselines: map[int]*pkg.Baseline{},
-	}
-}
-
-type Baselines struct {
-	FailedBaselines []*pkg.Baseline
-	random          *pkg.Baseline
-	index           *pkg.Baseline
-	baselines       map[int]*pkg.Baseline
-}
diff --git a/internal/pool/checkpool.go b/internal/pool/checkpool.go
@@ -2,6 +2,7 @@ package pool
 
 import (
 	"context"
+	"errors"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/parsers"
 	"github.com/chainreactors/spray/internal/ihttp"
@@ -19,9 +20,10 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) {
 	pctx, cancel := context.WithCancel(ctx)
 	pool := &CheckPool{
 		&This{
-			Config: config,
-			ctx:    pctx,
-			Cancel: cancel,
+			Config:    config,
+			Statistor: pkg.NewStatistor(""),
+			ctx:       pctx,
+			Cancel:    cancel,
 			client: ihttp.NewClient(&ihttp.ClientConfig{
 				Thread:    config.Thread,
 				Type:      config.ClientType,
@@ -112,7 +114,7 @@ func (pool *CheckPool) Invoke(v interface{}) {
 		defer fasthttp.ReleaseRequest(req.FastRequest)
 	}
 
-	if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge {
+	if reqerr != nil && !errors.Is(reqerr, fasthttp.ErrBodyTooLarge) {
 		pool.failedCount++
 		bl = &pkg.Baseline{
 			SprayResult: &parsers.SprayResult{
@@ -156,9 +158,11 @@ func (pool *CheckPool) Invoke(v interface{}) {
 		}
 	}
 
+	if bl.Source == parsers.CheckSource {
+		pool.Bar.Done()
+	}
 	pool.reqCount++
 	pool.wg.Done()
-	pool.Bar.Done()
 }
 
 func (pool *CheckPool) doRedirect(bl *pkg.Baseline, depth int) {

diff --git a/internal/pool/unit.go → internal/pool/types.go b/internal/pool/unit.go → internal/pool/types.go
@@ -1,6 +1,9 @@
 package pool
 
-import "github.com/chainreactors/parsers"
+import (
+	"github.com/chainreactors/parsers"
+	"github.com/chainreactors/spray/pkg"
+)
 
 func newUnit(path string, source parsers.SpraySource) *Unit {
 	return &Unit{path: path, source: source}
@@ -18,3 +21,16 @@ type Unit struct {
 	frontUrl string
 	depth    int // redirect depth
 }
+
+func NewBaselines() *Baselines {
+	return &Baselines{
+		baselines: map[int]*pkg.Baseline{},
+	}
+}
+
+type Baselines struct {
+	FailedBaselines []*pkg.Baseline
+	random          *pkg.Baseline
+	index           *pkg.Baseline
+	baselines       map[int]*pkg.Baseline
+}