From b530438c058efe7c705c5c55b42e102c6d6f5761 Mon Sep 17 00:00:00 2001
From: eiiko6 <myeiiko6@gmail.com>
Date: Tue, 13 Jan 2026 19:17:20 +0100
Subject: [PATCH] added pingpong to websockets

---
 src/main.rs      |  2 +-
 src/routes/ws.rs | 59 ++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 6d52bdb..e54be3c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -36,7 +36,7 @@ pub struct Cli {
     port: String,
 
     /// Database URL
-    #[arg(short, long, default_value = "localhost:5432")]
+    #[arg(short, long, default_value = "0.0.0.0:5432")]
     database: String,
 
     /// Data directory path
diff --git a/src/routes/ws.rs b/src/routes/ws.rs
index 213a536..62686f2 100644
--- a/src/routes/ws.rs
+++ b/src/routes/ws.rs
@@ -10,6 +10,8 @@ use axum::{
     response::IntoResponse,
 };
 use serde::Deserialize;
+use std::time::Duration;
+use tokio::select;
 use uuid::Uuid;
 
 use crate::auth::{create_jwt, verify_jwt};
@@ -115,13 +117,56 @@ async fn handle_socket(
     mut socket: WebSocket,
     mut receiver: tokio::sync::broadcast::Receiver<crate::routes::messages::Message>,
 ) {
-    while let Ok(msg) = receiver.recv().await {
-        if socket
-            .send(WsMessage::Text(serde_json::to_string(&msg).unwrap().into()))
-            .await
-            .is_err()
-        {
-            break;
+    let mut ping_interval = tokio::time::interval(Duration::from_secs(30));
+
+    loop {
+        select! {
+            // Receive broadcast messages and send to client
+            msg = receiver.recv() => {
+                if let Ok(msg) = msg {
+                    if let Ok(json) = serde_json::to_string(&msg) {
+                        if socket.send(WsMessage::Text(json.into())).await.is_err() {
+                            tracing::warn!("WebSocket send failed, closing socket");
+                            break;
+                        }
+                    }
+                } else {
+                    break;
+                }
+            }
+
+            // Send Ping
+            _ = ping_interval.tick() => {
+                if socket.send(WsMessage::Ping(vec![].into())).await.is_err() {
+                    tracing::error!("Failed to send ping, closing connection");
+                    break;
+                }
+
+                tracing::info!("Ping sent successfully");
+            }
+
+            // Get incoming messages from client
+            client_msg = socket.recv() => {
+                match client_msg {
+                    Some(Ok(msg)) => {
+                        match msg {
+                            WsMessage::Pong(_) => {
+                                tracing::info!("Received Pong");
+                            }
+                            WsMessage::Ping(_) => {
+                                tracing::info!("Received Ping from client");
+                            }
+                            // WsMessage::Text(_) => {}
+                            WsMessage::Close(_) => {
+                                tracing::info!("Client disconnected");
+                                break;
+                            }
+                            _ => {}
+                        }
+                    },
+                    _ => break,
+                }
+            }
         }
     }
 }