server healthcheck support

rkonfj · Apr 28, 2023 · 0736dbc · 0736dbc
1 parent 9e939f8
commit 0736dbc
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 66 deletions.
diff --git a/cmd/s5/cmd.go b/cmd/s5/cmd.go
@@ -124,10 +124,11 @@ func defaultOptions() *server.Config {
 		Geoip2: "country.mmdb",
 		Listen: "0.0.0.0:2080",
 		Servers: []server.TohServer{{
-			Name:    "us1",
-			Api:     "wss://us-l4-vultr.synf.in/ws",
-			Key:     "5868a941-3025-4c6d-ad3a-41e29bb42e5f",
-			Ruleset: []string{"https://raw.githubusercontent.com/rkonfj/toh/main/ruleset.txt"},
+			Name:        "us1",
+			Api:         "wss://us-l4-vultr.synf.in/ws",
+			Key:         "5868a941-3025-4c6d-ad3a-41e29bb42e5f",
+			Ruleset:     []string{"https://raw.githubusercontent.com/rkonfj/toh/main/ruleset.txt"},
+			Healthcheck: "https://www.google.com/generate_204",
 		}},
 	}
 }
diff --git a/cmd/s5/server/server.go b/cmd/s5/server/server.go
@@ -4,11 +4,11 @@ import (
 	"context"
 	"errors"
 	"io"
-	"math/rand"
 	"net"
 	"net/http"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"time"
 
@@ -29,10 +29,11 @@ type Config struct {
 }
 
 type TohServer struct {
-	Name    string   `yaml:"name"`
-	Api     string   `yaml:"api"`
-	Key     string   `yaml:"key"`
-	Ruleset []string `yaml:"ruleset"`
+	Name        string   `yaml:"name"`
+	Api         string   `yaml:"api"`
+	Key         string   `yaml:"key"`
+	Ruleset     []string `yaml:"ruleset"`
+	Healthcheck string   `yaml:"healthcheck"`
 }
 
 type ServerGroup struct {
@@ -50,9 +51,11 @@ type RulebasedSocks5Server struct {
 }
 
 type Server struct {
-	name    string
-	client  *client.TohClient
-	ruleset *ruleset.Ruleset
+	name       string
+	client     *client.TohClient
+	httpClient *http.Client
+	ruleset    *ruleset.Ruleset
+	latency    time.Duration
 }
 
 type Group struct {
@@ -81,13 +84,26 @@ func NewSocks5Server(dataPath string, cfg Config) (socks5Server *RulebasedSocks5
 		server := &Server{
 			name:   s.Name,
 			client: c,
+			httpClient: &http.Client{
+				Timeout: 120 * time.Second,
+				Transport: &http.Transport{
+					DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+						addr, err := spec.ResolveIP(ctx, c.DialTCP, addr)
+						if err != nil {
+							return nil, err
+						}
+						return c.DialTCP(ctx, addr)
+					},
+				},
+			},
 		}
 		if s.Ruleset != nil {
 			server.ruleset, err = ruleset.Parse(c, s.Name, s.Ruleset, dataPath)
 			if err != nil {
 				return
 			}
 		}
+		go healthcheck(server, s.Healthcheck)
 		socks5Server.servers = append(socks5Server.servers, server)
 	}
 
@@ -113,8 +129,7 @@ func NewSocks5Server(dataPath string, cfg Config) (socks5Server *RulebasedSocks5
 		socks5Server.groups = append(socks5Server.groups, group)
 	}
 
-	httpClient := securityHttpClient(socks5Server.servers)
-	socks5Server.geoip2db, err = openGeoip2(httpClient, dataPath, cfg.Geoip2)
+	socks5Server.geoip2db, err = openGeoip2(selectServer(socks5Server.servers).httpClient, dataPath, cfg.Geoip2)
 	if err != nil {
 		return
 	}
@@ -141,23 +156,10 @@ func (s *RulebasedSocks5Server) dialUDP(ctx context.Context, addr string) (diale
 }
 
 func selectServer(servers []*Server) *Server {
-	return servers[rand.Intn(len(servers))]
-}
-
-func securityHttpClient(servers []*Server) *http.Client {
-	return &http.Client{
-		Timeout: 120 * time.Second,
-		Transport: &http.Transport{
-			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
-				server := selectServer(servers)
-				addr, err := spec.ResolveIP(ctx, server.client.DialTCP, addr)
-				if err != nil {
-					return nil, err
-				}
-				return server.client.DialTCP(ctx, addr)
-			},
-		},
-	}
+	sort.Slice(servers, func(i, j int) bool {
+		return servers[i].latency < servers[j].latency
+	})
+	return servers[0]
 }
 
 func openGeoip2(httpClient *http.Client, dataPath, geoip2Path string) (*geoip2.Reader, error) {

diff --git a/cmd/s5/server/stats.go b/cmd/s5/server/stats.go
@@ -1,6 +1,9 @@
 package server
 
 import (
+	"strings"
+	"time"
+
 	"github.com/rkonfj/toh/socks5"
 	"github.com/sirupsen/logrus"
 )
@@ -17,3 +20,20 @@ func logTrafficEvent(e *socks5.TrafficEvent) {
 		WithField("stats_out", e.RemoteAddr).
 		Info()
 }
+
+func healthcheck(server *Server, url string) {
+	if strings.TrimSpace(url) == "" {
+		server.latency = time.Duration(0)
+		return
+	}
+	for {
+		t1 := time.Now()
+		_, err := server.httpClient.Get(url)
+		if err != nil {
+			server.latency = server.httpClient.Timeout
+		} else {
+			server.latency = time.Since(t1)
+		}
+		time.Sleep(15 * time.Second)
+	}
+}
diff --git a/cmd/s5/server/tcp.go b/cmd/s5/server/tcp.go
@@ -9,7 +9,6 @@ import (
 )
 
 func (s *RulebasedSocks5Server) dialTCP(ctx context.Context, addr string) (dialerName string, conn net.Conn, err error) {
-	log := logrus.WithField(spec.AppAddr.String(), ctx.Value(spec.AppAddr))
 	host, _, err := net.SplitHostPort(addr)
 	if err != nil {
 		return
@@ -30,21 +29,14 @@ func (s *RulebasedSocks5Server) dialTCP(ctx context.Context, addr string) (diale
 		if len(s.groups) > 0 {
 			for _, g := range s.groups {
 				if g.ruleset.CountryMatch(c.Country.IsoCode) {
-					server := selectServer(g.servers)
-					dialerName = server.name
-					log.Infof("%s using %s.%s", addr, g.name, dialerName)
-					conn, err = server.client.DialTCP(ctx, addr)
-					return
+					return dialTCPUseServer(ctx, addr, g.name, selectServer(g.servers))
 				}
 			}
 		}
 
-		for _, toh := range s.servers {
-			if toh.ruleset.CountryMatch(c.Country.IsoCode) {
-				log.Infof("%s using %s", addr, toh.name)
-				dialerName = toh.name
-				conn, err = toh.client.DialTCP(ctx, addr)
-				return
+		for _, s := range s.servers {
+			if s.ruleset.CountryMatch(c.Country.IsoCode) {
+				return dialTCPUseServer(ctx, addr, "", s)
 			}
 		}
 
@@ -54,46 +46,44 @@ func (s *RulebasedSocks5Server) dialTCP(ctx context.Context, addr string) (diale
 	if len(s.groups) > 0 {
 		for _, g := range s.groups {
 			if g.ruleset.SpecialMatch(host) {
-				server := selectServer(g.servers)
-				dialerName = server.name
-				log.Infof("%s using %s.%s", addr, g.name, dialerName)
-				conn, err = server.client.DialTCP(ctx, addr)
-				return
+				return dialTCPUseServer(ctx, addr, g.name, selectServer(g.servers))
 			}
 		}
 
 		for _, g := range s.groups {
 			if g.ruleset.WildcardMatch(host) {
-				server := selectServer(g.servers)
-				dialerName = server.name
-				log.Infof("%s using %s.%s", addr, g.name, dialerName)
-				conn, err = server.client.DialTCP(ctx, addr)
-				return
+				return dialTCPUseServer(ctx, addr, g.name, selectServer(g.servers))
 			}
 		}
 	}
 
-	for _, toh := range s.servers {
-		if toh.ruleset.SpecialMatch(host) {
-			log.Infof("%s using %s", addr, toh.name)
-			dialerName = toh.name
-			conn, err = toh.client.DialTCP(ctx, addr)
-			return
+	for _, s := range s.servers {
+		if s.ruleset.SpecialMatch(host) {
+			return dialTCPUseServer(ctx, addr, "", s)
 		}
 	}
 
-	for _, toh := range s.servers {
-		if toh.ruleset.WildcardMatch(host) {
-			log.Infof("%s using %s", addr, toh.name)
-			dialerName = toh.name
-			conn, err = toh.client.DialTCP(ctx, addr)
-			return
+	for _, s := range s.servers {
+		if s.ruleset.WildcardMatch(host) {
+			return dialTCPUseServer(ctx, addr, "", s)
 		}
 	}
 
 direct:
-	log.Infof("%s using direct", addr)
+	logrus.WithField(spec.AppAddr.String(), ctx.Value(spec.AppAddr)).Infof("%s using direct", addr)
 	dialerName = "direct"
 	conn, err = s.defaultDialer.DialContext(ctx, "tcp", addr)
 	return
 }
+
+func dialTCPUseServer(ctx context.Context, addr, groupName string, server *Server) (dialerName string, conn net.Conn, err error) {
+	log := logrus.WithField(spec.AppAddr.String(), ctx.Value(spec.AppAddr))
+	dialerName = server.name
+	if groupName == "" {
+		log.Infof("%s using %s latency %s", addr, dialerName, server.latency)
+	} else {
+		log.Infof("%s using %s.%s latency %s", addr, groupName, dialerName, server.latency)
+	}
+	conn, err = server.client.DialTCP(ctx, addr)
+	return
+}