internal/websocket/broadcast.go

@@ -24,8 +24,7 @@ func SendPermanentBroadCastMessage() {
 		mu.Lock()
 		for _, client := range Clients {
 			// Send message to client
-			client.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
-			err := client.Conn.WriteMessage(websocket.TextMessage, jsonMessage)
+			err := client.SendMessage(websocket.TextMessage, jsonMessage)
 			if err != nil {
 				helper.Logger.Info("Error broadcasting to client:", "msg", err, "host", client.RealIP, "clients", len(Clients))
 			}

internal/websocket/handle_connection.go

@@ -15,8 +15,6 @@ import (
 var Clients = make(map[net.Addr]*models.WebsocketClient)
 var mu sync.Mutex // Mutex to protect access to the Clients map
 
-const TIMEOUT = 10
-
 // Upgrade HTTP requests to WebSocket connections
 var upgrader = websocket.Upgrader{
 	CheckOrigin: func(r *http.Request) bool { return true },

internal/websocket/staleClients.go

@@ -2,27 +2,8 @@ package websocket
 
 import (
 	"time"
-
-	"github.com/gorilla/websocket"
-
-	"git.smsvc.net/pomodoro/GoTomato/pkg/models"
 )
 
-// Checks if the websockets last Pong is recent
-func isStale(client *models.WebsocketClient) bool {
-	return time.Since(client.LastPong) > (TIMEOUT * time.Second)
-}
-
-// send a Ping to the websocket client
-func sendPing(client *models.WebsocketClient) bool {
-	client.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
-	err := client.Conn.WriteMessage(websocket.PingMessage, nil)
-	if err != nil {
-		return false
-	}
-	return true
-}
-
 // Check and remove stale clients
 func RemoveStaleClients() {
 	ticker := time.NewTicker(30 * time.Second)
@@ -31,8 +12,7 @@ func RemoveStaleClients() {
 	for range ticker.C {
 		mu.Lock()
 		for _, client := range Clients {
-			// send ping and check for last pong
-			if !sendPing(client) || isStale(client) {
+			if client.IsStale() {
 				client.Conn.Close()
 				delete(Clients, client.Conn.LocalAddr())
 			}

pkg/models/client.go

@@ -6,6 +6,8 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+const TIMEOUT = 10
+
 // Represents a command from the client (start/stop)
 type ClientCommand struct {
 	Command  string         `json:"command"`  // Command send to the server
@@ -19,3 +21,25 @@ type WebsocketClient struct {
 	LastPong time.Time
 	RealIP   string
 }
+
+// Sends a message to the websocket.
+func (c *WebsocketClient) SendMessage(messageType int, data []byte) error {
+	c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
+	err := c.Conn.WriteMessage(websocket.PingMessage, nil)
+	if err != nil {
+		return err
+	}
+
+	c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
+	err = c.Conn.WriteMessage(messageType, data)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Checks if the websockets last Pong is recent
+func (c *WebsocketClient) IsStale() bool {
+	return time.Since(c.LastPong) > (TIMEOUT * time.Second)
+}