From a5966355ae38b9f0241de5171e44f0a668c8ee80 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Thu, 17 Nov 2022 17:09:37 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E--white-status,=20=E7=99=BD?=
 =?UTF-8?q?=E5=90=8D=E5=8D=95=E7=8A=B6=E6=80=81=E7=A0=81=E5=B0=86=E8=B7=B3?=
 =?UTF-8?q?=E8=BF=87precompare?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/option.go | 11 +++++++++++
 internal/pool.go   | 20 +++++++-------------
 internal/runner.go | 11 ++---------
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/internal/option.go b/internal/option.go
index 7323295..b53c24b 100644
--- a/internal/option.go
+++ b/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
diff --git a/internal/pool.go b/internal/pool.go
index e70cea2..d8b6e44 100644
--- a/internal/pool.go
+++ b/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
 	}
 
diff --git a/internal/runner.go b/internal/runner.go
index eba69f0..9bed6f2 100644
--- a/internal/runner.go
+++ b/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)