add client.IsStale()

internal/websocket/broadcast.go

@@ -22,11 +22,12 @@ func SendPermanentBroadCastMessage() {
 
 		// Iterate over all connected clients and broadcast the message
 		for _, client := range Clients {
-			if time.Now().Sub(client.LastPong) > (10 * time.Second) {
 			// 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))

internal/websocket/handle_connection.go

@@ -36,6 +36,12 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 		LastPong: time.Now(),
 		RealIP:   r.RemoteAddr,
 	}
+
+	client.Conn.SetPongHandler(func(appData string) error {
+		client.LastPong = time.Now()
+		return nil
+	})
+
 	mu.Lock()
 	Clients[ws.LocalAddr()] = &client
 	mu.Unlock()

pkg/models/client.go

@@ -24,12 +24,22 @@ type WebsocketClient struct {
 
 // Sends a message to the websocket.
 func (c *WebsocketClient) SendMessage(messageType int, data []byte) error {
-	c.Conn.SetPongHandler(func(s string) error {
-		c.LastPong = time.Now()
-		return nil
-	})
+	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))
-	c.Conn.WriteMessage(websocket.PingMessage, nil)
-	return c.Conn.WriteMessage(messageType, data)
+	err = c.Conn.WriteMessage(messageType, data)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Check if websockets last Pong is recent
+func (c *WebsocketClient) IsStale() bool {
+	return time.Since(c.LastPong) > (TIMEOUT * time.Second)
 }