From a4124b50aad8ff718f118bd49ca43d008869ea6e Mon Sep 17 00:00:00 2001
From: sjlleo <i@leo.moe>
Date: Tue, 9 Aug 2022 03:25:35 -0400
Subject: [PATCH] =?UTF-8?q?refactor:=20WebSocket=20=E6=8F=A1=E6=89=8B?=
 =?UTF-8?q?=E9=80=BB=E8=BE=91=E6=94=B9=E8=BF=9B=E4=BB=A5=E5=8F=8A=E5=BF=83?=
 =?UTF-8?q?=E8=B7=B3=E5=8C=85=E6=A3=80=E6=B5=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 wshandle/client.go | 107 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 88 insertions(+), 19 deletions(-)

diff --git a/wshandle/client.go b/wshandle/client.go
index 44fcaf8..3fb4b9d 100644
--- a/wshandle/client.go
+++ b/wshandle/client.go
@@ -12,6 +12,8 @@ import (
 )
 
 type WsConn struct {
+	Connecting   bool
+	Connected    bool            // 连接状态
 	MsgSendCh    chan string     // 消息发送通道
 	MsgReceiveCh chan string     // 消息接收通道
 	Done         chan struct{}   // 发送结束通道
@@ -23,14 +25,43 @@ type WsConn struct {
 
 var wsconn *WsConn
 
-func (c *WsConn) messageReceiveHandler() {
-	defer close(c.Done)
-	for {
-		_, msg, err := c.Conn.ReadMessage()
-		if err != nil {
-			return
+func (c *WsConn) keepAlive() {
+	go func() {
+		// 开启一个定时器
+		for {
+			<-time.After(time.Second * 54)
+			if c.Connected {
+				c.Conn.WriteMessage(websocket.TextMessage, []byte("ping"))
+			}
+		}
+	}()
+	for {
+		if !c.Connected && !c.Connecting {
+			c.Connecting = true
+			c.recreateWsConn()
+			// log.Println("WebSocket 连接意外断开，正在尝试重连...")
+			// return
+		}
+		// 降低检测频率，优化 CPU 占用情况
+		<-time.After(200 * time.Millisecond)
+	}
+}
+
+func (c *WsConn) messageReceiveHandler() {
+	// defer close(c.Done)
+	for {
+		if c.Connected {
+			_, msg, err := c.Conn.ReadMessage()
+			if err != nil {
+				// 读取信息出错，连接已经意外断开
+				// log.Println(err)
+				c.Connected = false
+				return
+			}
+			if string(msg) != "pong" {
+				c.MsgReceiveCh <- string(msg)
+			}
 		}
-		c.MsgReceiveCh <- string(msg)
 	}
 }
 
@@ -39,20 +70,26 @@ func (c *WsConn) messageSendHandler() {
 		// 循环监听发送
 		select {
 		case <-c.Done:
+			log.Println("发送协程已经退出")
 			return
 		case t := <-c.MsgSendCh:
-			err := c.Conn.WriteMessage(websocket.TextMessage, []byte(t))
-			if err != nil {
-				log.Println("write:", err)
-				return
+			// log.Println(t)
+			if !c.Connected {
+				c.MsgReceiveCh <- `{"ip":"` + t + `", "asnumber":"API服务端异常"}`
+			} else {
+				err := c.Conn.WriteMessage(websocket.TextMessage, []byte(t))
+				if err != nil {
+					log.Println("write:", err)
+					return
+				}
 			}
 		// 来自终端的中断运行请求
 		case <-c.Interrupt:
 			// 向 websocket 发起关闭连接任务
 			err := c.Conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
 			if err != nil {
-				log.Println("write close:", err)
-				return
+				// log.Println("write close:", err)
+				os.Exit(1)
 			}
 			select {
 			// 等到了结果，直接退出
@@ -66,6 +103,27 @@ func (c *WsConn) messageSendHandler() {
 	}
 }
 
+func (c *WsConn) recreateWsConn() {
+	u := url.URL{Scheme: "wss", Host: "api.leo.moe", Path: "/v2/ipGeoWs"}
+	// log.Printf("connecting to %s", u.String())
+
+	ws, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	c.Conn = ws
+	if err != nil {
+		log.Println("dial:", err)
+		// <-time.After(time.Second * 1)
+		c.Connected = false
+		c.Connecting = false
+		return
+	} else {
+		c.Connected = true
+	}
+	c.Connecting = false
+
+	c.Done = make(chan struct{})
+	go c.messageReceiveHandler()
+}
+
 func createWsConn() *WsConn {
 	// 设置终端中断通道
 	interrupt := make(chan os.Signal, 1)
@@ -75,17 +133,28 @@ func createWsConn() *WsConn {
 	// log.Printf("connecting to %s", u.String())
 
 	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+
+	wsconn = &WsConn{
+		Conn:         c,
+		Connected:    true,
+		Connecting:   false,
+		MsgSendCh:    make(chan string, 10),
+		MsgReceiveCh: make(chan string, 10),
+	}
+
 	if err != nil {
-		log.Fatal("dial:", err)
+		log.Println("dial:", err)
+		// <-time.After(time.Second * 1)
+		wsconn.Connected = false
+		wsconn.Done = make(chan struct{})
+		go wsconn.keepAlive()
+		go wsconn.messageSendHandler()
+		return wsconn
 	}
 	// defer c.Close()
 	// 将连接写入WsConn，方便随时可取
-	wsconn = &WsConn{
-		Conn:         c,
-		MsgSendCh:    make(chan string),
-		MsgReceiveCh: make(chan string),
-	}
 	wsconn.Done = make(chan struct{})
+	go wsconn.keepAlive()
 	go wsconn.messageReceiveHandler()
 	go wsconn.messageSendHandler()
 	return wsconn