From cbbe483a8ca7e156fb007ce4c34cacfe44b1268a Mon Sep 17 00:00:00 2001 From: Tobias Kantusch Date: Mon, 4 Sep 2023 16:17:58 +0200 Subject: [PATCH] fix(engineio): respect pingTimeout and pingInterval to stop connection The Engine.IO protocol features a heartbeat mechanism which ensures that the connection between server and client is alive. During that heartbeat, the server sends PINGs to which the client responds with PONGs. Both parties can therefore detect whether the connection is still alive, based on the pingInterval and pingTimeout received in the initial handshake. However, we previously didn't make use of that in the Engine.IO implementation, which lead to disconnects not being properly recognized. We now respect these settings and return an error and disconnect the connection, once the pingInterval+pingTimeout time has passed. See also https://socket.io/docs/v4/how-it-works/#disconnection-detection --- engineio/src/asynchronous/async_socket.rs | 66 +++++++++++++++---- .../src/asynchronous/client/async_client.rs | 3 +- engineio/src/client/client.rs | 4 +- engineio/src/error.rs | 2 + engineio/src/socket.rs | 37 ++++++++++- engineio/src/transport.rs | 4 +- engineio/src/transports/polling.rs | 10 ++- engineio/src/transports/websocket.rs | 15 +++-- engineio/src/transports/websocket_secure.rs | 9 ++- 9 files changed, 119 insertions(+), 31 deletions(-) diff --git a/engineio/src/asynchronous/async_socket.rs b/engineio/src/asynchronous/async_socket.rs index e213d2a8..8718e0ea 100644 --- a/engineio/src/asynchronous/async_socket.rs +++ b/engineio/src/asynchronous/async_socket.rs @@ -9,7 +9,7 @@ use std::{ use async_stream::try_stream; use bytes::Bytes; -use futures_util::{Stream, StreamExt}; +use futures_util::{stream, Stream, StreamExt}; use tokio::{runtime::Handle, sync::Mutex, time::Instant}; use crate::{ @@ -19,12 +19,11 @@ use crate::{ Error, Packet, PacketId, }; -use super::generator::StreamGenerator; - #[derive(Clone)] pub struct Socket { handle: Handle, transport: Arc>, + transport_raw: AsyncTransportType, on_close: OptionalCallback<()>, on_data: OptionalCallback, on_error: OptionalCallback, @@ -34,7 +33,7 @@ pub struct Socket { last_ping: Arc>, last_pong: Arc>, connection_data: Arc, - generator: StreamGenerator, + max_ping_timeout: u64, } impl Socket { @@ -47,6 +46,8 @@ impl Socket { on_open: OptionalCallback<()>, on_packet: OptionalCallback, ) -> Self { + let max_ping_timeout = handshake.ping_interval + handshake.ping_timeout; + Socket { handle: Handle::current(), on_close, @@ -55,11 +56,12 @@ impl Socket { on_open, on_packet, transport: Arc::new(Mutex::new(transport.clone())), + transport_raw: transport, connected: Arc::new(AtomicBool::default()), last_ping: Arc::new(Mutex::new(Instant::now())), last_pong: Arc::new(Mutex::new(Instant::now())), connection_data: Arc::new(handshake), - generator: StreamGenerator::new(Self::stream(transport)), + max_ping_timeout, } } @@ -202,6 +204,23 @@ impl Socket { *self.last_ping.lock().await = Instant::now(); } + /// Returns the time in milliseconds that is left until a new ping must be received. + /// This is used to detect whether we have been disconnected from the server. + /// See https://socket.io/docs/v4/how-it-works/#disconnection-detection + async fn time_to_next_ping(&self) -> u64 { + match Instant::now().checked_duration_since(*self.last_ping.lock().await) { + Some(since_last_ping) => { + let since_last_ping = since_last_ping.as_millis() as u64; + if since_last_ping > self.max_ping_timeout { + 0 + } else { + self.max_ping_timeout - since_last_ping + } + } + None => 0, + } + } + pub(crate) fn handle_packet(&self, packet: Packet) { if let Some(on_packet) = self.on_packet.as_ref() { let on_packet = on_packet.clone(); @@ -224,16 +243,35 @@ impl Socket { self.connected.store(false, Ordering::Release); } -} -impl Stream for Socket { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - self.generator.poll_next_unpin(cx) + /// Returns the packet stream for the client. + pub(crate) fn as_stream<'a>( + &'a self, + ) -> Pin> + Send + 'a>> { + stream::unfold( + Self::stream(self.transport_raw.clone()), + |mut stream| async { + // Wait for the next payload or until we should have received the next ping. + match tokio::time::timeout( + std::time::Duration::from_millis(self.time_to_next_ping().await), + stream.next(), + ) + .await + { + Ok(result) => result.map(|result| (result, stream)), + // We didn't receive a ping in time and now consider the connection as closed. + Err(_) => { + // Be nice and disconnect properly. + if let Err(e) = self.disconnect().await { + Some((Err(e), stream)) + } else { + Some((Err(Error::PingTimeout()), stream)) + } + } + } + }, + ) + .boxed() } } diff --git a/engineio/src/asynchronous/client/async_client.rs b/engineio/src/asynchronous/client/async_client.rs index 100c5cbd..99b0d0dd 100644 --- a/engineio/src/asynchronous/client/async_client.rs +++ b/engineio/src/asynchronous/client/async_client.rs @@ -54,7 +54,8 @@ impl Client { socket: InnerSocket, ) -> Pin> + 'static + Send>> { Box::pin(try_stream! { - for await item in socket.clone() { + let socket = socket.clone(); + for await item in socket.as_stream() { let packet = item?; socket.handle_incoming_packet(packet.clone()).await?; yield packet; diff --git a/engineio/src/client/client.rs b/engineio/src/client/client.rs index 0af38243..dc22ff77 100644 --- a/engineio/src/client/client.rs +++ b/engineio/src/client/client.rs @@ -1,5 +1,6 @@ use super::super::socket::Socket as InnerSocket; use crate::callback::OptionalCallback; +use crate::socket::DEFAULT_MAX_POLL_TIMEOUT; use crate::transport::Transport; use crate::error::{Error, Result}; @@ -128,7 +129,8 @@ impl ClientBuilder { let mut url = self.url.clone(); - let handshake: HandshakePacket = Packet::try_from(transport.poll()?)?.try_into()?; + let handshake: HandshakePacket = + Packet::try_from(transport.poll(DEFAULT_MAX_POLL_TIMEOUT)?)?.try_into()?; // update the base_url with the new sid url.query_pairs_mut().append_pair("sid", &handshake.sid[..]); diff --git a/engineio/src/error.rs b/engineio/src/error.rs index cc21401e..440fe10a 100644 --- a/engineio/src/error.rs +++ b/engineio/src/error.rs @@ -52,6 +52,8 @@ pub enum Error { InvalidHeaderNameFromReqwest(#[from] reqwest::header::InvalidHeaderName), #[error("Invalid header value")] InvalidHeaderValueFromReqwest(#[from] reqwest::header::InvalidHeaderValue), + #[error("The server did not send a PING packet in time")] + PingTimeout(), } pub(crate) type Result = std::result::Result; diff --git a/engineio/src/socket.rs b/engineio/src/socket.rs index 0d4127d0..b5231ab6 100644 --- a/engineio/src/socket.rs +++ b/engineio/src/socket.rs @@ -6,12 +6,18 @@ use crate::packet::{HandshakePacket, Packet, PacketId, Payload}; use bytes::Bytes; use std::convert::TryFrom; use std::sync::RwLock; +use std::time::Duration; use std::{fmt::Debug, sync::atomic::Ordering}; use std::{ sync::{atomic::AtomicBool, Arc, Mutex}, time::Instant, }; +/// The default maximum ping timeout as calculated from the pingInterval and pingTimeout. +/// See https://socket.io/docs/v4/server-options/#pinginterval and +/// https://socket.io/docs/v4/server-options/#pingtimeout +pub const DEFAULT_MAX_POLL_TIMEOUT: Duration = Duration::from_secs(45); + /// An `engine.io` socket which manages a connection with the server and allows /// it to register common callbacks. #[derive(Clone)] @@ -28,6 +34,7 @@ pub struct Socket { connection_data: Arc, /// Since we get packets in payloads it's possible to have a state where only some of the packets have been consumed. remaining_packets: Arc>>, + max_ping_timeout: u64, } impl Socket { @@ -40,6 +47,8 @@ impl Socket { on_open: OptionalCallback<()>, on_packet: OptionalCallback, ) -> Self { + let max_ping_timeout = handshake.ping_interval + handshake.ping_timeout; + Socket { on_close, on_data, @@ -52,6 +61,7 @@ impl Socket { last_pong: Arc::new(Mutex::new(Instant::now())), connection_data: Arc::new(handshake), remaining_packets: Arc::new(RwLock::new(None)), + max_ping_timeout, } } @@ -126,9 +136,13 @@ impl Socket { } } - // Iterator has run out of packets, get a new payload - // TODO: 0.3.X timeout? - let data = self.transport.as_transport().poll()?; + // Iterator has run out of packets, get a new payload. + // Make sure that payload is received within time_to_next_ping, as otherwise the heart + // stopped beating and we disconnect. + let data = self + .transport + .as_transport() + .poll(Duration::from_millis(self.time_to_next_ping()?))?; if data.is_empty() { continue; @@ -165,6 +179,23 @@ impl Socket { Ok(()) } + /// Returns the time in milliseconds that is left until a new ping must be received. + /// This is used to detect whether we have been disconnected from the server. + /// See https://socket.io/docs/v4/how-it-works/#disconnection-detection + fn time_to_next_ping(&self) -> Result { + match Instant::now().checked_duration_since(*self.last_ping.lock()?) { + Some(since_last_ping) => { + let since_last_ping = since_last_ping.as_millis() as u64; + if since_last_ping > self.max_ping_timeout { + Ok(0) + } else { + Ok(self.max_ping_timeout - since_last_ping) + } + } + None => Ok(0), + } + } + pub(crate) fn handle_packet(&self, packet: Packet) { if let Some(on_packet) = self.on_packet.as_ref() { spawn_scoped!(on_packet(packet)); diff --git a/engineio/src/transport.rs b/engineio/src/transport.rs index 39ff5c91..a0945c50 100644 --- a/engineio/src/transport.rs +++ b/engineio/src/transport.rs @@ -2,7 +2,7 @@ use super::transports::{PollingTransport, WebsocketSecureTransport, WebsocketTra use crate::error::Result; use adler32::adler32; use bytes::Bytes; -use std::time::SystemTime; +use std::time::{Duration, SystemTime}; use url::Url; pub trait Transport { @@ -13,7 +13,7 @@ pub trait Transport { /// Performs the server long polling procedure as long as the client is /// connected. This should run separately at all time to ensure proper /// response handling from the server. - fn poll(&self) -> Result; + fn poll(&self, timeout: Duration) -> Result; /// Returns start of the url. ex. http://localhost:2998/engine.io/?EIO=4&transport=polling /// Must have EIO and transport already set. diff --git a/engineio/src/transports/polling.rs b/engineio/src/transports/polling.rs index e1ca1b5c..26aee87b 100644 --- a/engineio/src/transports/polling.rs +++ b/engineio/src/transports/polling.rs @@ -8,6 +8,7 @@ use reqwest::{ header::HeaderMap, }; use std::sync::{Arc, RwLock}; +use std::time::Duration; use url::Url; #[derive(Debug, Clone)] @@ -77,8 +78,13 @@ impl Transport for PollingTransport { Ok(()) } - fn poll(&self) -> Result { - Ok(self.client.get(self.address()?).send()?.bytes()?) + fn poll(&self, timeout: Duration) -> Result { + Ok(self + .client + .get(self.address()?) + .timeout(timeout) + .send()? + .bytes()?) } fn base_url(&self) -> Result { diff --git a/engineio/src/transports/websocket.rs b/engineio/src/transports/websocket.rs index a9d9f4c2..7281cfa7 100644 --- a/engineio/src/transports/websocket.rs +++ b/engineio/src/transports/websocket.rs @@ -8,7 +8,7 @@ use crate::{ }; use bytes::Bytes; use http::HeaderMap; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; use tokio::runtime::Runtime; use url::Url; @@ -46,9 +46,12 @@ impl Transport for WebsocketTransport { .block_on(async { self.inner.emit(data, is_binary_att).await }) } - fn poll(&self) -> Result { + fn poll(&self, timeout: Duration) -> Result { self.runtime.block_on(async { - let r = self.inner.poll_next().await; + let r = match tokio::time::timeout(timeout, self.inner.poll_next()).await { + Ok(r) => r, + Err(_) => return Err(Error::PingTimeout()), + }; match r { Ok(b) => b.ok_or(Error::IncompletePacket()), Err(_) => Err(Error::IncompletePacket()), @@ -81,6 +84,8 @@ mod test { use crate::ENGINE_IO_VERSION; use std::str::FromStr; + const TIMEOUT_DURATION: Duration = Duration::from_secs(45); + fn new() -> Result { let url = crate::test::engine_io_server()?.to_string() + "engine.io/?EIO=" @@ -123,8 +128,8 @@ mod test { format!("{:?}", transport), format!("WebsocketTransport(base_url: {:?})", transport.base_url()) ); - println!("{:?}", transport.poll().unwrap()); - println!("{:?}", transport.poll().unwrap()); + println!("{:?}", transport.poll(TIMEOUT_DURATION).unwrap()); + println!("{:?}", transport.poll(TIMEOUT_DURATION).unwrap()); Ok(()) } } diff --git a/engineio/src/transports/websocket_secure.rs b/engineio/src/transports/websocket_secure.rs index 98f64da0..345e0880 100644 --- a/engineio/src/transports/websocket_secure.rs +++ b/engineio/src/transports/websocket_secure.rs @@ -10,7 +10,7 @@ use crate::{ use bytes::Bytes; use http::HeaderMap; use native_tls::TlsConnector; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; use tokio::runtime::Runtime; use url::Url; @@ -54,9 +54,12 @@ impl Transport for WebsocketSecureTransport { .block_on(async { self.inner.emit(data, is_binary_att).await }) } - fn poll(&self) -> Result { + fn poll(&self, timeout: Duration) -> Result { self.runtime.block_on(async { - let r = self.inner.poll_next().await; + let r = match tokio::time::timeout(timeout, self.inner.poll_next()).await { + Ok(r) => r, + Err(_) => return Err(Error::PingTimeout()), + }; match r { Ok(b) => b.ok_or(Error::IncompletePacket()), Err(_) => Err(Error::IncompletePacket()),