From 1a656b26a15545258fdf97864e4986b5ddabe36b Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 4 May 2023 12:04:59 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E--retry=E6=8F=92=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/option.go | 45 +++++++++++++++++++-------------------------- internal/pool.go | 27 ++++++++++++++++++++------- internal/runner.go | 2 ++ internal/types.go | 2 ++ pkg/baseline.go | 1 + pkg/config.go | 1 + 6 files changed, 45 insertions(+), 33 deletions(-) diff --git a/internal/option.go b/internal/option.go index 0a5ae75..c79e559 100644 --- a/internal/option.go +++ b/internal/option.go @@ -99,23 +99,23 @@ type PluginOptions struct { } type ModeOptions struct { - RateLimit int `long:"rate-limit" default:"0" description:"Int, request rate limit (rate/s), e.g.: --rate-limit 100"` - Force bool `long:"force" description:"Bool, skip error break"` - CheckOnly bool `long:"check-only" description:"Bool, check only"` - NoScope bool `long:"no-scope" description:"Bool, no scope"` - Scope []string `long:"scope" description:"String, custom scope, e.g.: --scope *.example.com"` - Recursive string `long:"recursive" default:"current.IsDir()" description:"String,custom recursive rule, e.g.: --recursive current.IsDir()"` - Depth int `long:"depth" default:"0" description:"Int, recursive depth"` - CheckPeriod int `long:"check-period" default:"200" description:"Int, check period when request"` - ErrPeriod int `long:"error-period" default:"10" description:"Int, check period when error"` - BreakThreshold int `long:"error-threshold" default:"20" description:"Int, break when the error exceeds the threshold "` - BlackStatus string `long:"black-status" default:"400,410" description:"Strings (comma split),custom black status, "` - WhiteStatus string `long:"white-status" default:"200" description:"Strings (comma split), custom white status"` - FuzzyStatus string `long:"fuzzy-status" default:"404,403,500,501,502,503" description:"Strings (comma split), custom fuzzy status"` - UniqueStatus string `long:"unique-status" default:"403" description:"Strings (comma split), custom unique status"` - Unique bool `long:"unique" description:"Bool, unique response"` - - SimhashDistance int `long:"distance" default:"5"` + RateLimit int `long:"rate-limit" default:"0" description:"Int, request rate limit (rate/s), e.g.: --rate-limit 100"` + Force bool `long:"force" description:"Bool, skip error break"` + CheckOnly bool `long:"check-only" description:"Bool, check only"` + NoScope bool `long:"no-scope" description:"Bool, no scope"` + Scope []string `long:"scope" description:"String, custom scope, e.g.: --scope *.example.com"` + Recursive string `long:"recursive" default:"current.IsDir()" description:"String,custom recursive rule, e.g.: --recursive current.IsDir()"` + Depth int `long:"depth" default:"0" description:"Int, recursive depth"` + CheckPeriod int `long:"check-period" default:"200" description:"Int, check period when request"` + ErrPeriod int `long:"error-period" default:"10" description:"Int, check period when error"` + BreakThreshold int `long:"error-threshold" default:"20" description:"Int, break when the error exceeds the threshold "` + BlackStatus string `long:"black-status" default:"400,410" description:"Strings (comma split),custom black status, "` + WhiteStatus string `long:"white-status" default:"200" description:"Strings (comma split), custom white status"` + FuzzyStatus string `long:"fuzzy-status" default:"404,403,500,501,502,503" description:"Strings (comma split), custom fuzzy status"` + UniqueStatus string `long:"unique-status" default:"403" description:"Strings (comma split), custom unique status"` + Unique bool `long:"unique" description:"Bool, unique response"` + RetryCount int `long:"retry" default:"1" description:"Int, retry count"` + SimhashDistance int `long:"distance" default:"5"` } type MiscOptions struct { @@ -162,6 +162,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { Active: opt.Active, Bak: opt.Bak, Common: opt.Common, + RetryCount: opt.RetryCount, RandomUserAgent: opt.RandomUserAgent, } @@ -237,7 +238,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { if s.Len() > 0 { logs.Log.Important("Advance Mod: " + s.String()) } - + logs.Log.Important("Retry Count: " + strconv.Itoa(r.RetryCount)) if opt.NoScope { r.Scope = []string{"*"} } @@ -539,14 +540,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.RecursiveExpr = exp } - // 自定义scope - //var scopeexpr string - //if opt.NoScope { - // scopeexpr = "glob(current, '*')" - //} else if opt.Scope != "" { - // scopeexpr = fmt.Sprintf("glob(current, '*%s*')", opt.Scope) - //} - // prepare header for _, h := range opt.Headers { i := strings.Index(h, ":") diff --git a/internal/pool.go b/internal/pool.go index 37f53d4..b85f326 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -310,6 +310,9 @@ func (pool *Pool) Invoke(v interface{}) { }, } pool.failedBaselines = append(pool.failedBaselines, bl) + // 自动重放失败请求, 默认为一次 + pool.doRetry(bl) + } else { if unit.source <= 3 || unit.source == CrawlSource || unit.source == CommonFileSource { // 一些高优先级的source, 将跳过PreCompare @@ -617,9 +620,9 @@ func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) { return } reURL := FormatURL(bl.Url.Path, bl.RedirectURL) - pool.waiter.Add(1) go func() { + defer pool.waiter.Done() pool.addAddition(&Unit{ path: reURL, source: RedirectSource, @@ -649,7 +652,6 @@ func (pool *Pool) doCrawl(bl *pkg.Baseline) { if u = FormatURL(bl.Url.Path, u); u == "" { continue } - pool.waiter.Add(1) pool.addAddition(&Unit{ path: u, source: CrawlSource, @@ -698,7 +700,6 @@ func (pool *Pool) doRule(bl *pkg.Baseline) { go func() { defer pool.waiter.Done() for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) { - pool.waiter.Add(1) pool.addAddition(&Unit{ path: Dir(bl.Url.Path) + u, source: RuleSource, @@ -707,10 +708,24 @@ func (pool *Pool) doRule(bl *pkg.Baseline) { }() } +func (pool *Pool) doRetry(bl *pkg.Baseline) { + if bl.Retry >= pool.Retry { + return + } + pool.waiter.Add(1) + go func() { + defer pool.waiter.Done() + pool.addAddition(&Unit{ + path: bl.Path, + source: RetrySource, + retry: bl.Retry + 1, + }) + }() +} + func (pool *Pool) doActive() { defer pool.waiter.Done() for _, u := range pkg.ActivePath { - pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + u[1:], source: ActiveSource, @@ -726,7 +741,6 @@ func (pool *Pool) doBak() { } worder.Run() for w := range worder.C { - pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + w, source: BakSource, @@ -739,7 +753,6 @@ func (pool *Pool) doBak() { } worder.Run() for w := range worder.C { - pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + w, source: BakSource, @@ -750,7 +763,6 @@ func (pool *Pool) doBak() { func (pool *Pool) doCommonFile() { defer pool.waiter.Done() for _, u := range mask.SpecialWords["common_file"] { - pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + u, source: CommonFileSource, @@ -776,6 +788,7 @@ func (pool *Pool) doCheck() { func (pool *Pool) addAddition(u *Unit) { // 强行屏蔽报错, 防止goroutine泄露 + pool.waiter.Add(1) defer func() { if err := recover(); err != nil { } diff --git a/internal/runner.go b/internal/runner.go index 6c5030c..b3b0098 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -80,6 +80,7 @@ type Runner struct { Active bool Bak bool Common bool + RetryCount int RandomUserAgent bool } @@ -106,6 +107,7 @@ func (r *Runner) PrepareConfig() *pkg.Config { Active: r.Active, Bak: r.Bak, Common: r.Common, + Retry: r.RetryCount, ClientType: r.ClientType, RandomUserAgent: r.RandomUserAgent, } diff --git a/internal/types.go b/internal/types.go index c823082..6ff46a3 100644 --- a/internal/types.go +++ b/internal/types.go @@ -19,6 +19,7 @@ const ( BakSource CommonFileSource UpgradeSource + RetrySource ) func newUnit(path string, source int) *Unit { @@ -33,6 +34,7 @@ type Unit struct { number int path string source int + retry int frontUrl string depth int // redirect depth } diff --git a/pkg/baseline.go b/pkg/baseline.go index 201ad0c..ebcb357 100644 --- a/pkg/baseline.go +++ b/pkg/baseline.go @@ -112,6 +112,7 @@ type Baseline struct { RecuDepth int `json:"-"` URLs []string `json:"-"` Collected bool `json:"-"` + Retry int `json:"-"` } func (bl *Baseline) IsDir() bool { diff --git a/pkg/config.go b/pkg/config.go index 10d136f..cb6b788 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -45,5 +45,6 @@ type Config struct { Active bool Bak bool Common bool + Retry int RandomUserAgent bool }