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

@@ -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<Mutex<AsyncTransportType>>,
+    transport_raw: AsyncTransportType,
     on_close: OptionalCallback<()>,
     on_data: OptionalCallback<Bytes>,
     on_error: OptionalCallback<String>,
@@ -34,7 +33,7 @@ pub struct Socket {
     last_ping: Arc<Mutex<Instant>>,
     last_pong: Arc<Mutex<Instant>>,
     connection_data: Arc<HandshakePacket>,
-    generator: StreamGenerator<Packet>,
+    max_ping_timeout: u64,
 }
 
 impl Socket {
@@ -47,6 +46,8 @@ impl Socket {
         on_open: OptionalCallback<()>,
         on_packet: OptionalCallback<Packet>,
     ) -> 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<Packet>;
-
-    fn poll_next(
-        mut self: Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Option<Self::Item>> {
-        self.generator.poll_next_unpin(cx)
+    /// Returns the packet stream for the client.
+    pub(crate) fn as_stream<'a>(
+        &'a self,
+    ) -> Pin<Box<dyn Stream<Item = Result<Packet>> + 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()
     }
 }
 

engineio/src/asynchronous/client/async_client.rs

@@ -54,7 +54,8 @@ impl Client {
         socket: InnerSocket,
     ) -> Pin<Box<impl Stream<Item = Result<Packet>> + '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;

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[..]);

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<T> = std::result::Result<T, Error>;

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<HandshakePacket>,
     /// 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<RwLock<Option<crate::packet::IntoIter>>>,
+    max_ping_timeout: u64,
 }
 
 impl Socket {
@@ -40,6 +47,8 @@ impl Socket {
         on_open: OptionalCallback<()>,
         on_packet: OptionalCallback<Packet>,
     ) -> 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<u64> {
+        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));

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<Bytes>;
+    fn poll(&self, timeout: Duration) -> Result<Bytes>;
 
     /// Returns start of the url. ex. http://localhost:2998/engine.io/?EIO=4&transport=polling
     /// Must have EIO and transport already set.

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<Bytes> {
-        Ok(self.client.get(self.address()?).send()?.bytes()?)
+    fn poll(&self, timeout: Duration) -> Result<Bytes> {
+        Ok(self
+            .client
+            .get(self.address()?)
+            .timeout(timeout)
+            .send()?
+            .bytes()?)
     }
 
     fn base_url(&self) -> Result<Url> {

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<Bytes> {
+    fn poll(&self, timeout: Duration) -> Result<Bytes> {
         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<WebsocketTransport> {
         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(())
     }
 }

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<Bytes> {
+    fn poll(&self, timeout: Duration) -> Result<Bytes> {
         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()),