Repository: 1c3t3a/rust-socketio (language: rust)

Issue:
EngineIO does not timeout on missing pings
## Problem Description

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.

See also https://socket.io/docs/v4/how-it-works/#disconnection-detection

Currently, this is not respected and the client does not recognize (and then terminate) a dead connection.

## Steps To Reproduce

1. Start a long running client and server (like the ones from #311) but start them in separate Docker containers
2. Start the server and then the client
3. Pause the server container with `docker pause`
4. Wait for 60 seconds and notice that the client continues to run and doesn't notice that the server isn't alive anymore.

## Possible Solution

Polling via HTTP or waiting for the next packet in WebSockets should be guarded with a timeout.


Discussion / hints:

Fix patch:
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<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()
     }
 }
 
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<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;
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<T> = std::result::Result<T, Error>;
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<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));
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<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.
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<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> {
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<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(())
     }
 }
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<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()),


Test patch: