优化--force的逻辑, 开启--force时将关闭check

internal/option.go

@@ -16,6 +16,7 @@ type Option struct {
 	InputOptions
 	OutputOptions
 	RequestOptions
+	ModeOptions
 	MiscOptions
 }
 
@@ -50,9 +51,15 @@ type RequestOptions struct {
 	Headers         []string `long:"header"`
 	Method          string   `long:"method"`
 	Cookie          string   `long:"cookie"`
+	SimhashDistance int      `long:"distance" default:"5"`
+}
+
+type ModeOptions struct {
 	Force          bool `long:"force"`
-	SimhashDistance int      `long:"distance"`
-	CheckOnly       bool     `long:"--check-only"`
+	CheckOnly      bool `long:"check-only"`
+	CheckPeriod    int  `long:"check-period" default:"100"`
+	ErrPeriod      int  `long:"error-period" default:"10"`
+	BreakThreshold int  `long:"error-threshold" default:"20"`
 }
 
 type MiscOptions struct {
@@ -87,6 +94,9 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		Fuzzy:          opt.Fuzzy,
 		Force:          opt.Force,
 		CheckOnly:      opt.CheckOnly,
+		CheckPeriod:    opt.CheckPeriod,
+		ErrPeriod:      opt.ErrPeriod,
+		BreakThreshold: opt.BreakThreshold,
 	}
 
 	err = pkg.LoadTemplates()
@@ -108,7 +118,10 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 	}
 
 	if opt.Force {
-		breakThreshold = 999999
+		// 如果开启了force模式, 将关闭check机制, err积累到一定数量自动退出机制
+		r.BreakThreshold = max
+		r.CheckPeriod = max
+		r.ErrPeriod = max
 	}
 
 	// prepare url

internal/pool.go

@@ -19,8 +19,7 @@ var (
 	CheckRedirect   func(string) bool
 	CheckWaf        func([]byte) bool
 )
-
-var breakThreshold int = 20
+var max = 2147483647
 
 func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 	pctx, cancel := context.WithCancel(ctx)
@@ -34,8 +33,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 		tempCh:      make(chan *pkg.Baseline, config.Thread),
 		wg:          sync.WaitGroup{},
 		initwg:      sync.WaitGroup{},
-		checkPeriod: 100,
-		errPeriod:   10,
 		reqCount:    1,
 		failedCount: 1,
 	}
@@ -49,7 +46,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			pool.wg.Add(1)
 			_ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
 
-			if pool.failedCount > breakThreshold {
+			if pool.failedCount > pool.BreakThreshold {
 				// 当报错次数超过上限是, 结束任务
 				pool.recover()
 				pool.cancel()
@@ -64,7 +61,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			pool.wg.Add(1)
 			_ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource))
 
-			if pool.failedCount > breakThreshold {
+			if pool.failedCount > pool.BreakThreshold {
 				// 当报错次数超过上限是, 结束任务
 				pool.recover()
 				pool.cancel()
@@ -109,7 +106,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			return
 		case CheckSource:
 			if bl.Err != "" {
-				logs.Log.Warnf("[check.error] maybe ip had banned by waf, break (%d/%d), error: %s", pool.failedCount, breakThreshold, bl.Err)
+				logs.Log.Warnf("[check.error] maybe ip had banned by waf, break (%d/%d), error: %s", pool.failedCount, pool.BreakThreshold, bl.Err)
 				pool.failedBaselines = append(pool.failedBaselines, bl)
 			} else if i := pool.base.Compare(bl); i < 1 {
 				if i == 0 {
@@ -128,10 +125,10 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			// 异步进行性能消耗较大的深度对比
 			pool.tempCh <- bl
 			pool.reqCount++
-			if pool.reqCount%pool.checkPeriod == 0 {
+			if pool.reqCount%pool.CheckPeriod == 0 {
 				pool.reqCount++
 				go pool.check()
-			} else if pool.failedCount%pool.errPeriod == 0 {
+			} else if pool.failedCount%pool.ErrPeriod == 0 {
 				pool.failedCount++
 				go pool.check()
 			}
@@ -168,8 +165,6 @@ type Pool struct {
 	tempCh          chan *pkg.Baseline // 待处理的baseline
 	reqCount        int
 	failedCount     int
-	checkPeriod     int
-	errPeriod       int
 	failedBaselines []*pkg.Baseline
 	base            *pkg.Baseline
 	baselines       map[int]*pkg.Baseline

internal/runner.go

@@ -38,6 +38,9 @@ type Runner struct {
 	Offset         int
 	Limit          int
 	Deadline       int
+	CheckPeriod    int
+	ErrPeriod      int
+	BreakThreshold int
 	CheckOnly      bool
 }
 
@@ -64,6 +67,9 @@ func (r *Runner) Prepare(ctx context.Context) error {
 			Fns:            r.Fns,
 			OutputCh:       r.OutputCh,
 			FuzzyCh:        r.FuzzyCh,
+			CheckPeriod:    r.CheckPeriod,
+			ErrPeriod:      r.ErrPeriod,
+			BreakThreshold: r.BreakThreshold,
 		}
 
 		if config.Mod == pkg.PathSpray {
@@ -84,6 +90,7 @@ func (r *Runner) Prepare(ctx context.Context) error {
 		if err != nil {
 			logs.Log.Error(err.Error())
 			if !r.Force {
+				// 如果没开启force, init失败将会关闭pool
 				pool.cancel()
 				r.poolwg.Done()
 				return

pkg/config.go

@@ -23,6 +23,9 @@ type Config struct {
 	Wordlist       []string
 	Thread         int
 	Timeout        int
+	CheckPeriod    int
+	ErrPeriod      int
+	BreakThreshold int
 	Method         string
 	Mod            SprayMod
 	Headers        http.Header