internal/websocket/broadcast.go

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

@@ -20,22 +20,6 @@ var upgrader = websocket.Upgrader{
 	CheckOrigin: func(r *http.Request) bool { return true },
 }
 
-// initialize a new WebsocketClient
-func newWebsocketClient(ws *websocket.Conn, address string) models.WebsocketClient {
-	wsclient := models.WebsocketClient{
-		Conn:     ws,
-		LastPong: time.Now(),
-		RealIP:   address,
-	}
-
-	wsclient.Conn.SetPongHandler(func(appData string) error {
-		wsclient.LastPong = time.Now()
-		return nil
-	})
-
-	return wsclient
-}
-
 // Upgrades HTTP requests to WebSocket connections and manages the client lifecycle
 func HandleConnection(w http.ResponseWriter, r *http.Request) {
 	// Upgrade initial GET request to a WebSocket
@@ -47,8 +31,11 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 	defer ws.Close()
 
 	// Register the new client
-	client := newWebsocketClient(ws, r.RemoteAddr)
-
+	client := models.WebsocketClient{
+		Conn:     ws,
+		LastPong: time.Now(),
+		RealIP:   r.RemoteAddr,
+	}
 	mu.Lock()
 	Clients[ws.LocalAddr()] = &client
 	mu.Unlock()

pkg/models/client.go

@@ -24,22 +24,12 @@ type WebsocketClient struct {
 
 // 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.SetPongHandler(func(s string) error {
+		c.LastPong = time.Now()
+		return nil
+	})
 
 	c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
-	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)
+	c.Conn.WriteMessage(websocket.PingMessage, nil)
+	return c.Conn.WriteMessage(messageType, data)
 }