新增--white-status, 白名单状态码将跳过precompare

internal/option.go

@@ -61,6 +61,7 @@ type ModeOptions struct {
 	ErrPeriod      int    `long:"error-period" default:"10"`
 	BreakThreshold int    `long:"error-threshold" default:"20"`
 	BlackStatus    string `long:"black-status" default:"default"`
+	WhiteStatus    string `long:"black-status" `
 }
 
 type MiscOptions struct {
@@ -137,6 +138,16 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		BlackStatus = []int{400, 404, 410}
 	}
 
+	if opt.WhiteStatus != "" {
+		for _, s := range strings.Split(opt.WhiteStatus, ",") {
+			si, err := strconv.Atoi(s)
+			if err != nil {
+				return nil, err
+			}
+			WhiteStatus = append(WhiteStatus, si)
+		}
+	}
+
 	// prepare url
 	var urls []string
 	var file *os.File

internal/pool.go

@@ -15,19 +15,9 @@ import (
 )
 
 var (
-	CheckBadStatus func(int) bool
-	CheckRedirect  func(string) bool
+	CheckRedirect func(string) bool
 )
 
-func CheckWaf(status int) bool {
-	for _, s := range WAFStatus {
-		if status == s {
-			return true
-		}
-	}
-	return false
-}
-
 var max = 2147483647
 
 func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
@@ -264,15 +254,19 @@ Loop:
 
 func (p *Pool) PreCompare(resp *ihttp.Response) error {
 	status := resp.StatusCode()
+	if IntsContains(WhiteStatus, status) {
+		// 如果为白名单状态码则直接返回
+		return nil
+	}
 	if p.base != nil && p.base.Status != 200 && p.base.Status == status {
 		return ErrSameStatus
 	}
 
-	if CheckBadStatus(status) {
+	if IntsContains(BlackStatus, status) {
 		return ErrBadStatus
 	}
 
-	if CheckWaf(status) {
+	if IntsContains(WAFStatus, status) {
 		return ErrWaf
 	}
 

internal/runner.go

@@ -14,7 +14,8 @@ import (
 )
 
 var (
-	BlackStatus = []int{}
+	WhiteStatus []int
+	BlackStatus []int
 	FuzzyStatus = []int{403, 500, 501, 502, 503}
 	WAFStatus   = []int{493, 418}
 )
@@ -49,14 +50,6 @@ type Runner struct {
 
 func (r *Runner) Prepare(ctx context.Context) error {
 	var err error
-	CheckBadStatus = func(status int) bool {
-		for _, black := range BlackStatus {
-			if black == status {
-				return true
-			}
-		}
-		return false
-	}
 
 	r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
 		u := i.(string)