add newWebsocketClient()

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
+			// 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

@@ -20,6 +20,22 @@ 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
@@ -31,11 +47,8 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 	defer ws.Close()
 
 	// Register the new client
-	client := models.WebsocketClient{
-		Conn:     ws,
-		LastPong: time.Now(),
-		RealIP:   r.RemoteAddr,
-	}
+	client := newWebsocketClient(ws, r.RemoteAddr)
+
 	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)
 }