From 05df91f053b1aa17d248885d3a5570082f31c38b Mon Sep 17 00:00:00 2001
From: shikong <919411476@qq.com>
Date: Sat, 20 Apr 2024 18:11:43 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=AD=E7=BA=BF=E9=87=8D=E8=BF=9E?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/client/client.go | 99 ++++++++++++++++++++++++++++++++------------
 1 file changed, 73 insertions(+), 26 deletions(-)

diff --git a/cmd/client/client.go b/cmd/client/client.go
index d8f1cf1..5baecd7 100644
--- a/cmd/client/client.go
+++ b/cmd/client/client.go
@@ -2,13 +2,16 @@ package client
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"github.com/gorilla/websocket"
 	"github.com/spf13/viper"
 	"gopkg.in/toast.v1"
 	"gotify-client/internal/client"
 	"log"
+	"net/http"
 	"net/url"
+	"time"
 
 	"gotify-client/constants"
 	"gotify-client/pkg/config"
@@ -43,11 +46,10 @@ func Main() {
 	}
 
 	if err := viper.ReadInConfig(); err != nil {
-		if _, ok := err.(viper.ConfigFileNotFoundError); ok {
+		var configFileNotFoundError viper.ConfigFileNotFoundError
+		if errors.As(err, &configFileNotFoundError) {
 			_ = toml.GenerateConfig()
 			logger.Log().Fatalf("未找到配置文件, 已生成示例配置文件于运行路径下")
-		} else {
-			logger.Log().Fatalf("配置解析失败 %s", err)
 		}
 	}
 
@@ -66,38 +68,83 @@ func Main() {
 	u.RawQuery = params.Encode()
 
 	logger.Log().Infof("url %s", u.String())
-	ws, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	dialer := &websocket.Dialer{
+		Proxy:            http.ProxyFromEnvironment,
+		HandshakeTimeout: 10 * time.Second,
+	}
+
+	wsBuilder := func() (*websocket.Conn, error) {
+		ws, _, err := dialer.Dial(u.String(), nil)
+		return ws, err
+	}
+	ws, err := wsBuilder()
 	if err != nil {
 		logger.Log().Fatalf("ws 连接失败: %s, 请检查配置是否有误", err)
 	}
+
 	defer func() {
-		_ = ws.Close()
+		err := recover()
+		if err != nil {
+			log.Println("panic:", err)
+			ws, _ = wsBuilder()
+		}
 	}()
 
-	done := make(chan struct{})
-
 	go func() {
-		defer close(done)
-		for {
-			_, rawMessage, err := ws.ReadMessage()
-			if err != nil {
-				log.Println("read:", err)
-				return
+		hasError := false
+
+		preHandler := func() {
+			if hasError {
+				for {
+					if ws != nil {
+						_ = ws.Close()
+					}
+
+					log.Println("尝试断线重连...")
+					ws, err = wsBuilder()
+					if err != nil {
+						log.Println("ws 重连异常:", err)
+						hasError = true
+						time.Sleep(time.Second * 3)
+					} else {
+						hasError = false
+						break
+					}
+				}
 			}
-
-			log.Printf("recv: %s", rawMessage)
-			data := &client.GotifyMessage{}
-			_ = json.Unmarshal(rawMessage, data)
-			message := fmt.Sprintf("%s", data.Message)
-
-			// windows 下 默认GBK中文编码转换
-			retTitle, _ := simplifiedchinese.GBK.NewEncoder().String(data.Title)
-			retMessage, _ := simplifiedchinese.GBK.NewEncoder().String(message)
-
-			notification.Title = retTitle
-			notification.Message = retMessage
-			_ = notification.Push()
 		}
+
+		handler := func() {
+			defer func() {
+				_ = ws.Close()
+			}()
+
+			preHandler()
+			for {
+				_, rawMessage, err := ws.ReadMessage()
+				if err != nil {
+					log.Println("ws 消息接收异常:", err)
+					hasError = true
+					preHandler()
+					continue
+				}
+
+				log.Printf("recv: %s", rawMessage)
+				data := &client.GotifyMessage{}
+				_ = json.Unmarshal(rawMessage, data)
+				message := fmt.Sprintf("%s", data.Message)
+
+				// windows 下 默认GBK中文编码转换
+				retTitle, _ := simplifiedchinese.GBK.NewEncoder().String(data.Title)
+				retMessage, _ := simplifiedchinese.GBK.NewEncoder().String(message)
+
+				notification.Title = retTitle
+				notification.Message = retMessage
+				_ = notification.Push()
+			}
+		}
+
+		handler()
 	}()
 
 	quit := make(chan os.Signal)