diff --git a/internal/pool/brutepool.go b/internal/pool/brutepool.go index 3b5ccee..16d3e11 100644 --- 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 index 56aeb00..89678f0 100644 --- 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 b/internal/pool/types.go similarity index 54% rename from internal/pool/unit.go rename to internal/pool/types.go index 06ff692..46cbea4 100644 --- a/internal/pool/unit.go +++ b/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 +}