From bb889412d219753ae58557f75eb89b1c42832a83 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Thu, 20 Jan 2022 23:50:56 +0800
Subject: [PATCH] improve cmd

---
 cmd/gost/cmd.go          | 62 ++++++++++++++++++++++++++++------------
 pkg/chain/router.go      |  5 +++-
 pkg/metadata/metadata.go |  4 +++
 3 files changed, 52 insertions(+), 19 deletions(-)

diff --git a/cmd/gost/cmd.go b/cmd/gost/cmd.go
index b1f7f36..cfc4016 100644
--- a/cmd/gost/cmd.go
+++ b/cmd/gost/cmd.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/url"
 	"strings"
+	"time"
 
 	"github.com/go-gost/gost/pkg/config"
 	"github.com/go-gost/gost/pkg/metadata"
@@ -72,28 +73,21 @@ func buildConfigFromCmd(services, nodes stringList) (*config.Config, error) {
 		}
 
 		var nodes []*config.NodeConfig
-		if v := metadata.GetString(md, "ip"); v != "" {
-			_, port, err := net.SplitHostPort(nodeConfig.Addr)
-			if err != nil {
-				return nil, err
+		for _, host := range strings.Split(nodeConfig.Addr, ",") {
+			if host == "" {
+				continue
 			}
-			ips := parseIP(v, port)
-			for i := range ips {
-				nodeCfg := &config.NodeConfig{}
-				*nodeCfg = *nodeConfig
-				nodeCfg.Name = fmt.Sprintf("node-%d", i)
-				nodeCfg.Addr = ips[i]
-				nodes = append(nodes, nodeCfg)
-			}
-			md.Del("ip")
-		}
-		if len(nodes) == 0 {
-			nodes = append(nodes, nodeConfig)
+			nodeCfg := &config.NodeConfig{}
+			*nodeCfg = *nodeConfig
+			nodeCfg.Name = fmt.Sprintf("node-%d", len(nodes))
+			nodeCfg.Addr = host
+			nodes = append(nodes, nodeCfg)
 		}
 
 		chain.Hops = append(chain.Hops, &config.HopConfig{
-			Name:  fmt.Sprintf("hop-%d", i),
-			Nodes: nodes,
+			Name:     fmt.Sprintf("hop-%d", i),
+			Selector: parseSelector(md),
+			Nodes:    nodes,
 		})
 	}
 
@@ -265,6 +259,10 @@ func buildServiceConfig(url *url.URL) (*config.ServiceConfig, error) {
 		md.Set("dns", strings.Split(v, ","))
 	}
 
+	if svc.Forwarder != nil {
+		svc.Forwarder.Selector = parseSelector(md)
+	}
+
 	svc.Handler = &config.HandlerConfig{
 		Type:     handler,
 		Auths:    auths,
@@ -399,6 +397,34 @@ func parseAuthFromCmd(sa string) (*config.AuthConfig, error) {
 	}, nil
 }
 
+func parseSelector(md metadata.MapMetadata) *config.SelectorConfig {
+	strategy := metadata.GetString(md, "strategy")
+	maxFails := metadata.GetInt(md, "maxFails")
+	failTimeout := metadata.GetDuration(md, "failTimeout")
+	if strategy == "" && maxFails <= 0 && failTimeout <= 0 {
+		return nil
+	}
+	if strategy == "" {
+		strategy = "round"
+	}
+	if maxFails <= 0 {
+		maxFails = 1
+	}
+	if failTimeout <= 0 {
+		failTimeout = time.Second
+	}
+
+	md.Del("strategy")
+	md.Del("maxFails")
+	md.Del("failTimeout")
+
+	return &config.SelectorConfig{
+		Strategy:    strategy,
+		MaxFails:    maxFails,
+		FailTimeout: failTimeout,
+	}
+}
+
 func parseIP(s, port string) (ips []string) {
 	if s == "" {
 		return nil
diff --git a/pkg/chain/router.go b/pkg/chain/router.go
index 3061bee..0c266bf 100644
--- a/pkg/chain/router.go
+++ b/pkg/chain/router.go
@@ -74,9 +74,12 @@ func (r *Router) resolve(ctx context.Context, addr string) (string, error) {
 	}
 
 	host, port, err := net.SplitHostPort(addr)
-	if err != nil || host == "" {
+	if err != nil {
 		return "", err
 	}
+	if host == "" {
+		return addr, nil
+	}
 
 	if r.Hosts != nil {
 		if ips, _ := r.Hosts.Lookup("ip", host); len(ips) > 0 {
diff --git a/pkg/metadata/metadata.go b/pkg/metadata/metadata.go
index 9ce835b..1b54336 100644
--- a/pkg/metadata/metadata.go
+++ b/pkg/metadata/metadata.go
@@ -93,6 +93,10 @@ func GetDuration(md Metadata, key string) (v time.Duration) {
 		return time.Duration(vv) * time.Second
 	case string:
 		v, _ = time.ParseDuration(vv)
+		if v == 0 {
+			n, _ := strconv.Atoi(vv)
+			v = time.Duration(n) * time.Second
+		}
 	}
 	return
 }