From 0bb74fcd8d9c4b5bca4752c0ba16be504f2d884d Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Thu, 3 Feb 2022 21:37:25 -0800
Subject: [PATCH] feat: client retry

---
 cmd/client.go | 33 +++++++++++++++++++++++++--------
 cmd/config.go | 10 ++++++----
 cmd/main.go   |  1 +
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/cmd/client.go b/cmd/client.go
index dffb732..4707225 100644
--- a/cmd/client.go
+++ b/cmd/client.go
@@ -110,13 +110,30 @@ func client(config *clientConfig) {
 		}
 	}
 	// Client
-	client, err := core.NewClient(config.Server, config.Protocol, auth, tlsConfig, quicConfig,
-		transport.DefaultClientTransport, uint64(config.UpMbps)*mbpsToBps, uint64(config.DownMbps)*mbpsToBps,
-		func(refBPS uint64) congestion.CongestionControl {
-			return hyCongestion.NewBrutalSender(congestion.ByteCount(refBPS))
-		}, obfuscator)
-	if err != nil {
-		logrus.WithField("error", err).Fatal("Failed to initialize client")
+	var client *core.Client
+	try := 0
+	for {
+		try += 1
+		c, err := core.NewClient(config.Server, config.Protocol, auth, tlsConfig, quicConfig,
+			transport.DefaultClientTransport, uint64(config.UpMbps)*mbpsToBps, uint64(config.DownMbps)*mbpsToBps,
+			func(refBPS uint64) congestion.CongestionControl {
+				return hyCongestion.NewBrutalSender(congestion.ByteCount(refBPS))
+			}, obfuscator)
+		if err != nil {
+			logrus.WithField("error", err).Error("Failed to initialize client")
+			if try <= config.Retry || config.Retry < 0 {
+				logrus.WithFields(logrus.Fields{
+					"retry":    try,
+					"interval": config.RetryInterval,
+				}).Info("Retrying...")
+				time.Sleep(time.Duration(config.RetryInterval) * time.Second)
+			} else {
+				logrus.Fatal("Out of retries, exiting...")
+			}
+		} else {
+			client = c
+			break
+		}
 	}
 	defer client.Close()
 	logrus.WithField("addr", config.Server).Info("Connected")
@@ -395,6 +412,6 @@ func client(config *clientConfig) {
 		}()
 	}
 
-	err = <-errChan
+	err := <-errChan
 	logrus.WithField("error", err).Fatal("Client shutdown")
 }
diff --git a/cmd/config.go b/cmd/config.go
index 2049599..13d31cf 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -97,10 +97,12 @@ func (r *Relay) Check() error {
 }
 
 type clientConfig struct {
-	Server   string `json:"server"`
-	Protocol string `json:"protocol"`
-	UpMbps   int    `json:"up_mbps"`
-	DownMbps int    `json:"down_mbps"`
+	Server        string `json:"server"`
+	Protocol      string `json:"protocol"`
+	UpMbps        int    `json:"up_mbps"`
+	DownMbps      int    `json:"down_mbps"`
+	Retry         int    `json:"retry"`
+	RetryInterval int    `json:"retry_interval"`
 	// Optional below
 	SOCKS5 struct {
 		Listen     string `json:"listen"`
diff --git a/cmd/main.go b/cmd/main.go
index 5439f8c..e6166c6 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -131,6 +131,7 @@ func initApp(c *cli.Context) error {
 				"version", "url",
 				"config", "file", "mode",
 				"addr", "src", "dst", "session", "action",
+				"retry", "interval",
 				"code", "msg", "error",
 			},
 			TimestampFormat: c.String("log-timestamp"),