新增--retry插件

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, ":")

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 {
 		}

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,
 	}

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
 }

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 {

pkg/config.go

@@ -45,5 +45,6 @@ type Config struct {
 	Active          bool
 	Bak             bool
 	Common          bool
+	Retry           int
 	RandomUserAgent bool
 }