diff --git a/src/api/notifications.rs b/src/api/notifications.rs index 8bfcd518..32204916 100644 --- a/src/api/notifications.rs +++ b/src/api/notifications.rs @@ -1,7 +1,7 @@ use std::{ net::IpAddr, sync::{Arc, LazyLock}, - time::Duration, + time::{Duration, Instant}, }; use chrono::{NaiveDateTime, Utc}; @@ -41,6 +41,14 @@ pub static WS_ANONYMOUS_SUBSCRIPTIONS: LazyLock = LazyLock::new(|| !CONFIG.enable_websocket() && !CONFIG.push_enabled()); pub fn routes() -> Vec { @@ -156,12 +164,16 @@ fn websockets_hub<'r>( rocket_ws::Stream! { ws => { let mut ws = ws; let _guard = guard; - let mut interval = tokio::time::interval(Duration::from_secs(15)); + let mut interval = tokio::time::interval(WS_PING_INTERVAL); + let mut last_received = Instant::now(); loop { tokio::select! { res = ws.next() => { match res { Some(Ok(message)) => { + // Any message, including a Pong, means the client is still there + last_received = Instant::now(); + match message { // Respond to any pings Message::Ping(ping) => yield Message::Pong(ping), @@ -195,7 +207,13 @@ fn websockets_hub<'r>( } } - _ = interval.tick() => yield Message::Ping(create_ping()) + _ = interval.tick() => { + // The client stopped responding without closing the connection, drop it + if last_received.elapsed() > WS_CLIENT_TIMEOUT { + break; + } + yield Message::Ping(create_ping()); + } } } }} @@ -229,12 +247,16 @@ fn anonymous_websockets_hub<'r>(ws: WebSocket, token: String, ip: ClientIp) -> R rocket_ws::Stream! { ws => { let mut ws = ws; let _guard = guard; - let mut interval = tokio::time::interval(Duration::from_secs(15)); + let mut interval = tokio::time::interval(WS_PING_INTERVAL); + let mut last_received = Instant::now(); loop { tokio::select! { res = ws.next() => { match res { Some(Ok(message)) => { + // Any message, including a Pong, means the client is still there + last_received = Instant::now(); + match message { // Respond to any pings Message::Ping(ping) => yield Message::Pong(ping), @@ -268,7 +290,13 @@ fn anonymous_websockets_hub<'r>(ws: WebSocket, token: String, ip: ClientIp) -> R } } - _ = interval.tick() => yield Message::Ping(create_ping()) + _ = interval.tick() => { + // The client stopped responding without closing the connection, drop it + if last_received.elapsed() > WS_CLIENT_TIMEOUT { + break; + } + yield Message::Ping(create_ping()); + } } } }}