From 17ed1a1551235ff6f712ce5091af2b788231b9dd Mon Sep 17 00:00:00 2001
From: Sebastian Mark <smark@posteo.net>
Date: Wed, 20 Nov 2024 09:32:04 +0100
Subject: [PATCH] feat: improve connection check and reconnect

- enhance `Connected()` method
  - send ping message to verify connection status
- refactor reconnect code for better understanding
---
 internal/websocket/connect.go | 21 +++++++++++++++++++--
 internal/websocket/receive.go | 22 +++++++++-------------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/internal/websocket/connect.go b/internal/websocket/connect.go
index 1f35b7b..9711c23 100644
--- a/internal/websocket/connect.go
+++ b/internal/websocket/connect.go
@@ -1,9 +1,10 @@
 package websocket
 
 import (
-	"github.com/gorilla/websocket"
 	"time"
 
+	"github.com/gorilla/websocket"
+
 	ChronoTomato "git.smsvc.net/pomodoro/ChronoTomato/pkg/models"
 )
 
@@ -17,7 +18,23 @@ func Connect(url string) Client {
 }
 
 func (c Client) Connected() bool {
-	return c.Conn != nil
+	if c.Conn == nil {
+		return false
+	}
+
+	// Set a pong handler to signal when a pong is received
+	c.Conn.SetPongHandler(func(s string) error {
+		c.Conn.SetReadDeadline(time.Now().Add(10 * time.Second))
+		return nil
+	})
+
+	// Send a ping
+	err := c.Conn.WriteMessage(websocket.PingMessage, nil)
+	if err != nil {
+		return false
+	}
+
+	return true
 }
 
 // Disconnects from websocket
diff --git a/internal/websocket/receive.go b/internal/websocket/receive.go
index 510ad34..ca70c40 100644
--- a/internal/websocket/receive.go
+++ b/internal/websocket/receive.go
@@ -22,26 +22,22 @@ func (c *Client) ProcessServerMessages(channel chan<- GoTomato.ServerMessage) {
 		c.Conn.SetReadDeadline(time.Now().Add(10 * time.Second))
 		_, message, err := c.Conn.ReadMessage()
 		if err != nil {
+			c.LastErr = err
+
 			// On normal closure exit gracefully
 			if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
 				return
 			}
 
-			c.LastErr = err
-			// Try to reconnect on unexpected disconnect
-			for {
-				channel <- serverMessage // send last known ServerMessage to update view
-
-				time.Sleep(time.Second)
-
-				// reconnect while preserving password
+			// check connection and reconnect if no response
+			if !c.Connected() {
 				pw := c.Password
-				*c = Connect(c.Server)
-				c.Password = pw
-
-				if c.Connected() {
-					break
+				for !c.Connected() {
+					channel <- serverMessage // send last known ServerMessage to update view
+					time.Sleep(time.Second)
+					*c = Connect(c.Server)
 				}
+				c.Password = pw
 			}
 			continue
 		}