fix(engineio): respect pingTimeout and pingInterval to stop connection

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()),