From 55ab44fc90775b4526bf6b33149c6a766e55a77d Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Thu, 23 Dec 2021 00:17:28 +0800
Subject: [PATCH] proxy: support regex match in http(s) health checking

---
 README.md                  |  4 ++--
 config.go                  |  2 +-
 config/glider.conf.example |  5 +++--
 rule/check.go              | 23 +++++++++++++++++------
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 2a8c7d2..242adf2 100644
--- a/README.md
+++ b/README.md
@@ -103,8 +103,8 @@ glider -h
 glider 0.15.2 usage:
   -check string
     	check=tcp[://HOST:PORT]: tcp port connect check
-    	check=http://HOST[:PORT][/URI][#expect=STRING_IN_RESP_LINE]
-      check=https://HOST[:PORT][/URI][#expect=STRING_IN_RESP_LINE]
+    	check=http://HOST[:PORT][/URI][#expect=REGEX_MATCH_IN_RESP_LINE]
+    	check=https://HOST[:PORT][/URI][#expect=REGEX_MATCH_IN_RESP_LINE]
     	check=file://SCRIPT_PATH: run a check script, healthy when exitcode=0, environment variables: FORWARDER_ADDR
     	check=disable: disable health check (default "http://www.msftconnecttest.com/connecttest.txt#expect=200")
   -checkdisabledonly
diff --git a/config.go b/config.go
index 9be2fd4..138591e 100644
--- a/config.go
+++ b/config.go
@@ -51,7 +51,7 @@ func parseConfig() *Config {
 
 	flag.StringSliceVar(&conf.Forwards, "forward", nil, "forward url, format: SCHEME://[USER|METHOD:PASSWORD@][HOST]:PORT?PARAMS[,SCHEME://[USER|METHOD:PASSWORD@][HOST]:PORT?PARAMS]")
 	flag.StringVar(&conf.Strategy.Strategy, "strategy", "rr", "forward strategy, default: rr")
-	flag.StringVar(&conf.Strategy.Check, "check", "http://www.msftconnecttest.com/connecttest.txt#expect=200", "check=tcp[://HOST:PORT]: tcp port connect check\ncheck=http://HOST[:PORT][/URI][#expect=STRING_IN_RESP_LINE]\ncheck=https://HOST[:PORT][/URI][#expect=STRING_IN_RESP_LINE]\ncheck=file://SCRIPT_PATH: run a check script, healthy when exitcode=0, environment variables: FORWARDER_ADDR\ncheck=disable: disable health check")
+	flag.StringVar(&conf.Strategy.Check, "check", "http://www.msftconnecttest.com/connecttest.txt#expect=200", "check=tcp[://HOST:PORT]: tcp port connect check\ncheck=http://HOST[:PORT][/URI][#expect=REGEX_MATCH_IN_RESP_LINE]\ncheck=https://HOST[:PORT][/URI][#expect=REGEX_MATCH_IN_RESP_LINE]\ncheck=file://SCRIPT_PATH: run a check script, healthy when exitcode=0, environment variables: FORWARDER_ADDR\ncheck=disable: disable health check")
 	flag.IntVar(&conf.Strategy.CheckInterval, "checkinterval", 30, "fowarder check interval(seconds)")
 	flag.IntVar(&conf.Strategy.CheckTimeout, "checktimeout", 10, "fowarder check timeout(seconds)")
 	flag.IntVar(&conf.Strategy.CheckTolerance, "checktolerance", 0, "fowarder check tolerance(ms), switch only when new_latency < old_latency - tolerance, only used in lha mode")
diff --git a/config/glider.conf.example b/config/glider.conf.example
index 557fefb..bf3e9df 100644
--- a/config/glider.conf.example
+++ b/config/glider.conf.example
@@ -192,8 +192,9 @@ maxfailures=3
 
 # Forwarder health check:
 # check=tcp[://HOST:PORT]: tcp port connect check
-# check=http://HOST[:PORT][/URI][#expect=STRING_IN_RESP_LINE]
-# check=https://HOST[:PORT][/URI][#expect=STRING_IN_RESP_LINE]
+# check=http://HOST[:PORT][/URI][#expect=REGEX_MATCH_IN_RESP_LINE]
+# check=https://HOST[:PORT][/URI][#expect=REGEX_MATCH_IN_RESP_LINE]
+# e.g. check=https://www.netflix.com/title/81215567#expect=301|404
 # check=file://SCRIPT_PATH: run a check script, healthy when exitcode=0, environment variables: FORWARDER_ADDR
 # check=disable: disable health check
 check=http://www.msftconnecttest.com/connecttest.txt#expect=200
diff --git a/rule/check.go b/rule/check.go
index 37b4a05..e4c9e24 100644
--- a/rule/check.go
+++ b/rule/check.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"os"
 	"os/exec"
+	"regexp"
 	"strings"
 	"time"
 
@@ -44,16 +45,26 @@ func (c *tcpChecker) Check(dialer proxy.Dialer) (time.Duration, error) {
 }
 
 type httpChecker struct {
-	addr       string
-	uri        string
-	expect     string
-	timeout    time.Duration
+	addr    string
+	uri     string
+	expect  string
+	timeout time.Duration
+
 	tlsConfig  *tls.Config
 	serverName string
+
+	regex *regexp.Regexp
 }
 
 func newHttpChecker(addr, uri, expect string, timeout time.Duration, withTLS bool) *httpChecker {
-	c := &httpChecker{addr: addr, uri: uri, expect: expect, timeout: timeout}
+	c := &httpChecker{
+		addr:    addr,
+		uri:     uri,
+		expect:  expect,
+		timeout: timeout,
+		regex:   regexp.MustCompile(expect),
+	}
+
 	if _, p, _ := net.SplitHostPort(addr); p == "" {
 		if withTLS {
 			c.addr = net.JoinHostPort(addr, "443")
@@ -103,7 +114,7 @@ func (c *httpChecker) Check(dialer proxy.Dialer) (time.Duration, error) {
 		return 0, err
 	}
 
-	if !strings.Contains(line, c.expect) {
+	if !c.regex.MatchString(line) {
 		return 0, fmt.Errorf("expect: %s, got: %s", c.expect, line)
 	}