comment updates

cmd/server/main.go

@@ -54,6 +54,7 @@ func Start() {
 	r := http.NewServeMux()
 	r.HandleFunc("/", websocket.HandleConnection)
 	go websocket.SendPermanentBroadCastMessage()
+	go websocket.RemoveStaleClients()
 
 	helper.Logger.Info("GoTomato started", "version", metadata.GoTomatoVersion)
 	helper.Logger.Info("Websocket listening", "address", listen)

internal/websocket/broadcast.go

@@ -21,18 +21,15 @@ func SendPermanentBroadCastMessage() {
 		}
 
 		// Iterate over all connected clients and broadcast the message
+		mu.Lock()
 		for _, client := range Clients {
-			// Remove unresponsive client
-			if client.IsStale() {
-				client.Conn.Close()
-				continue
-			}
 			// Send message to client
 			err := client.SendMessage(websocket.TextMessage, jsonMessage)
 			if err != nil {
 				helper.Logger.Info("Error broadcasting to client:", "msg", err, "host", client.RealIP, "clients", len(Clients))
 			}
 		}
+		mu.Unlock()
 
 		<-tick.C
 	}

internal/websocket/client_commands.go

@@ -22,7 +22,9 @@ func handleClientCommands(c models.WebsocketClient) {
 		_, message, err := ws.ReadMessage()
 		if err != nil {
 			// remove client on error/disconnect
+			mu.Lock()
 			delete(Clients, ws.LocalAddr())
+			mu.Unlock()
 			helper.Logger.Info("Client disconnected:", "msg", err, "host", c.RealIP, "clients", len(Clients))
 			break
 		}

internal/websocket/handle_connection.go

@@ -36,8 +36,7 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 		LastPong: time.Now(),
 		RealIP:   r.RemoteAddr,
 	}
-
-	client.Conn.SetPongHandler(func(appData string) error {
+	client.Conn.SetPongHandler(func(s string) error {
 		client.LastPong = time.Now()
 		return nil
 	})

internal/websocket/staleClients.go

@@ -0,0 +1,21 @@
+package websocket
+
+import (
+	"time"
+)
+
+// Check and remove stale clients
+func RemoveStaleClients() {
+	ticker := time.NewTicker(30 * time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		mu.Lock()
+		for _, client := range Clients {
+			if client.IsStale() {
+				client.Conn.Close()
+			}
+		}
+		mu.Unlock()
+	}
+}

pkg/models/client.go

@@ -39,7 +39,7 @@ func (c *WebsocketClient) SendMessage(messageType int, data []byte) error {
 	return nil
 }
 
-// Check if websockets last Pong is recent
+// Checks if the websockets last Pong is recent
 func (c *WebsocketClient) IsStale() bool {
 	return time.Since(c.LastPong) > (TIMEOUT * time.Second)
 }