From 3c1f9a995397cb94a32f4e7f6bc42aa7b8f74d82 Mon Sep 17 00:00:00 2001
From: wzdnzd <wzdnzd@users.noreply.github.com>
Date: Sun, 4 Jun 2023 14:00:24 +0800
Subject: [PATCH] ProxyProvider health check also supports specifying expected
 status (#600)

Co-authored-by: wwqgtxx <wwqgtxx@gmail.com>
---
 adapter/outboundgroup/parser.go |  2 +-
 adapter/provider/healthcheck.go | 50 ++++++++++++++++++---------------
 adapter/provider/parser.go      | 17 +++++++----
 config/config.go                |  2 +-
 4 files changed, 42 insertions(+), 29 deletions(-)

diff --git a/adapter/outboundgroup/parser.go b/adapter/outboundgroup/parser.go
index fccf51fd..7ebaa6c0 100644
--- a/adapter/outboundgroup/parser.go
+++ b/adapter/outboundgroup/parser.go
@@ -81,7 +81,7 @@ func ParseProxyGroup(config map[string]any, proxyMap map[string]C.Proxy, provide
 			return nil, errDuplicateProvider
 		}
 
-		hc := provider.NewHealthCheck(ps, "", 0, true)
+		hc := provider.NewHealthCheck(ps, "", 0, true, nil)
 		pd, err := provider.NewCompatibleProvider(groupName, ps, hc)
 		if err != nil {
 			return nil, err
diff --git a/adapter/provider/healthcheck.go b/adapter/provider/healthcheck.go
index cc7056f1..330e306e 100644
--- a/adapter/provider/healthcheck.go
+++ b/adapter/provider/healthcheck.go
@@ -32,16 +32,17 @@ type extraOption struct {
 }
 
 type HealthCheck struct {
-	url       string
-	extra     map[string]*extraOption
-	mu        sync.Mutex
-	started   *atomic.Bool
-	proxies   []C.Proxy
-	interval  uint
-	lazy      bool
-	lastTouch *atomic.Int64
-	done      chan struct{}
-	singleDo  *singledo.Single[struct{}]
+	url            string
+	extra          map[string]*extraOption
+	mu             sync.Mutex
+	started        *atomic.Bool
+	proxies        []C.Proxy
+	interval       uint
+	lazy           bool
+	expectedStatus utils.IntRanges[uint16]
+	lastTouch      *atomic.Int64
+	done           chan struct{}
+	singleDo       *singledo.Single[struct{}]
 }
 
 func (hc *HealthCheck) process() {
@@ -153,7 +154,8 @@ func (hc *HealthCheck) check() {
 		b, _ := batch.New[bool](context.Background(), batch.WithConcurrencyNum[bool](10))
 
 		// execute default health check
-		hc.execute(b, hc.url, id, nil)
+		option := &extraOption{filters: nil, expectedStatus: hc.expectedStatus}
+		hc.execute(b, hc.url, id, option)
 
 		// execute extra health check
 		if len(hc.extra) != 0 {
@@ -178,7 +180,10 @@ func (hc *HealthCheck) execute(b *batch.Batch[bool], url, uid string, option *ex
 	var store = C.OriginalHistory
 	var expectedStatus utils.IntRanges[uint16]
 	if option != nil {
-		store = C.ExtraHistory
+		if url != hc.url {
+			store = C.ExtraHistory
+		}
+
 		expectedStatus = option.expectedStatus
 		if len(option.filters) != 0 {
 			filters := make([]string, 0, len(option.filters))
@@ -214,16 +219,17 @@ func (hc *HealthCheck) close() {
 	hc.done <- struct{}{}
 }
 
-func NewHealthCheck(proxies []C.Proxy, url string, interval uint, lazy bool) *HealthCheck {
+func NewHealthCheck(proxies []C.Proxy, url string, interval uint, lazy bool, expectedStatus utils.IntRanges[uint16]) *HealthCheck {
 	return &HealthCheck{
-		proxies:   proxies,
-		url:       url,
-		extra:     map[string]*extraOption{},
-		started:   atomic.NewBool(false),
-		interval:  interval,
-		lazy:      lazy,
-		lastTouch: atomic.NewInt64(0),
-		done:      make(chan struct{}, 1),
-		singleDo:  singledo.NewSingle[struct{}](time.Second),
+		proxies:        proxies,
+		url:            url,
+		extra:          map[string]*extraOption{},
+		started:        atomic.NewBool(false),
+		interval:       interval,
+		lazy:           lazy,
+		expectedStatus: expectedStatus,
+		lastTouch:      atomic.NewInt64(0),
+		done:           make(chan struct{}, 1),
+		singleDo:       singledo.NewSingle[struct{}](time.Second),
 	}
 }
diff --git a/adapter/provider/parser.go b/adapter/provider/parser.go
index 1df7f320..954db5cb 100644
--- a/adapter/provider/parser.go
+++ b/adapter/provider/parser.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/Dreamacro/clash/common/structure"
+	"github.com/Dreamacro/clash/common/utils"
 	"github.com/Dreamacro/clash/component/resource"
 	C "github.com/Dreamacro/clash/constant"
 	types "github.com/Dreamacro/clash/constant/provider"
@@ -14,10 +15,11 @@ import (
 var errVehicleType = errors.New("unsupport vehicle type")
 
 type healthCheckSchema struct {
-	Enable   bool   `provider:"enable"`
-	URL      string `provider:"url"`
-	Interval int    `provider:"interval"`
-	Lazy     bool   `provider:"lazy,omitempty"`
+	Enable         bool   `provider:"enable"`
+	URL            string `provider:"url"`
+	Interval       int    `provider:"interval"`
+	Lazy           bool   `provider:"lazy,omitempty"`
+	ExpectedStatus string `provider:"expected-status,omitempty"`
 }
 
 type proxyProviderSchema struct {
@@ -44,11 +46,16 @@ func ParseProxyProvider(name string, mapping map[string]any) (types.ProxyProvide
 		return nil, err
 	}
 
+	expectedStatus, err := utils.NewIntRanges[uint16](schema.HealthCheck.ExpectedStatus)
+	if err != nil {
+		return nil, err
+	}
+
 	var hcInterval uint
 	if schema.HealthCheck.Enable {
 		hcInterval = uint(schema.HealthCheck.Interval)
 	}
-	hc := NewHealthCheck([]C.Proxy{}, schema.HealthCheck.URL, hcInterval, schema.HealthCheck.Lazy)
+	hc := NewHealthCheck([]C.Proxy{}, schema.HealthCheck.URL, hcInterval, schema.HealthCheck.Lazy, expectedStatus)
 
 	path := C.Path.Resolve(schema.Path)
 
diff --git a/config/config.go b/config/config.go
index 2f563b30..191dbf13 100644
--- a/config/config.go
+++ b/config/config.go
@@ -654,7 +654,7 @@ func parseProxies(cfg *RawConfig) (proxies map[string]C.Proxy, providersMap map[
 		}
 		ps = append(ps, proxies[v])
 	}
-	hc := provider.NewHealthCheck(ps, "", 0, true)
+	hc := provider.NewHealthCheck(ps, "", 0, true, nil)
 	pd, _ := provider.NewCompatibleProvider(provider.ReservedName, ps, hc)
 	providersMap[provider.ReservedName] = pd