fix(engineio): fix Noop packet emission when upgrading transport (#554)

Author: Totodore

crates/engineioxide/src/packet.rs

@@ -146,6 +146,11 @@ impl From<Packet> for tokio_tungstenite::tungstenite::Utf8Bytes {
         String::from(value).into()
     }
 }
+impl From<Packet> for Bytes {
+    fn from(value: Packet) -> Self {
+        String::from(value).into()
+    }
+}
 /// Deserialize a [Packet] from a [String] according to the Engine.IO protocol
 impl TryFrom<Str> for Packet {
     type Error = Error;

crates/engineioxide/src/transport/polling/mod.rs

@@ -102,6 +102,19 @@ where
         return Err(Error::TransportMismatch);
     }
 
+    if socket.is_upgrading() {
+        #[cfg(feature = "tracing")]
+        tracing::debug!(?sid, "socket is upgrading, sending NOOP packet");
+
+        #[cfg(feature = "v3")]
+        let data = payload::packet_encoder(Packet::Noop, socket.protocol, socket.supports_binary);
+        #[cfg(not(feature = "v3"))]
+        let data = payload::packet_encoder(Packet::Noop, socket.protocol);
+
+        let is_binary = false; // The noop packet is guaranteed to be serialized as text
+        return Ok(http_response(StatusCode::OK, data, is_binary)?);
+    }
+
     // If the socket is already locked, it means that the socket is being used by another request
     // In case of multiple http polling, session should be closed
     let rx = match socket.internal_rx.try_lock() {

crates/engineioxide/src/transport/polling/payload/encoder.rs

@@ -112,7 +112,7 @@ pub async fn v4_encoder(
 /// Encode one packet into a *binary* payload according to the
 /// [engine.io v3 protocol](https://github.com/socketio/engine.io-protocol/tree/v3#payload)
 #[cfg(feature = "v3")]
-pub fn v3_bin_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) -> Result<(), Error> {
+pub fn v3_bin_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) {
     use crate::transport::polling::payload::BINARY_PACKET_SEPARATOR_V3;
     use bytes::BufMut;
 
@@ -147,13 +147,12 @@ pub fn v3_bin_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) -> Resu
             data.extend_from_slice(packet.as_bytes()); // packet
         }
     };
-    Ok(())
 }
 
 /// Encode one packet into a *string* payload according to the
 /// [engine.io v3 protocol](https://github.com/socketio/engine.io-protocol/tree/v3#payload)
 #[cfg(feature = "v3")]
-pub fn v3_string_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) -> Result<(), Error> {
+pub fn v3_string_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) {
     use crate::transport::polling::payload::STRING_PACKET_SEPARATOR_V3;
     use bytes::BufMut;
     let packet: String = packet.into();
@@ -164,7 +163,6 @@ pub fn v3_string_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) -> R
         packet
     );
     data.put_slice(packet.as_bytes());
-    Ok(())
 }
 
 /// Encode multiple packet packet into a *string* payload if there is no binary packet or into a *binary* payload if there are binary packets
@@ -202,11 +200,11 @@ pub async fn v3_binary_encoder(
 
     if has_binary {
         for packet in packet_buffer {
-            v3_bin_packet_encoder(packet, &mut data)?;
+            v3_bin_packet_encoder(packet, &mut data);
         }
     } else {
         for packet in packet_buffer {
-            v3_string_packet_encoder(packet, &mut data)?;
+            v3_string_packet_encoder(packet, &mut data);
         }
     }
 
@@ -216,11 +214,11 @@ pub async fn v3_binary_encoder(
         for packet in packets {
             match packet {
                 Packet::BinaryV3(_) | Packet::Binary(_) => {
-                    v3_bin_packet_encoder(packet, &mut data)?;
+                    v3_bin_packet_encoder(packet, &mut data);
                     has_binary = true;
                 }
                 packet => {
-                    v3_string_packet_encoder(packet, &mut data)?;
+                    v3_string_packet_encoder(packet, &mut data);
                 }
             };
         }
@@ -250,15 +248,15 @@ pub async fn v3_string_encoder(
     let current_size = data.len() + PUNCTUATION_LEN + max_packet_size_len;
     while let Some(packets) = try_recv_packet(&mut rx, current_size, max_payload, true) {
         for packet in packets {
-            v3_string_packet_encoder(packet, &mut data)?;
+            v3_string_packet_encoder(packet, &mut data);
         }
     }
 
     // If there is no packet in the buffer, wait for the next packet
     if data.is_empty() {
         let packets = recv_packet(&mut rx).await?;
         for packet in packets {
-            v3_string_packet_encoder(packet, &mut data)?;
+            v3_string_packet_encoder(packet, &mut data);
         }
     }
 

crates/engineioxide/src/transport/polling/payload/mod.rs

@@ -87,3 +87,36 @@ pub async fn encoder(
         encoder::v4_encoder(rx, max_payload).await
     }
 }
+
+/// Encodes a single packet into a byte array.
+pub fn packet_encoder(
+    packet: Packet,
+    #[allow(unused_variables)] protocol: ProtocolVersion,
+    #[cfg(feature = "v3")] supports_binary: bool,
+) -> Bytes {
+    #[cfg(feature = "v3")]
+    {
+        match protocol {
+            ProtocolVersion::V4 => packet.into(),
+            ProtocolVersion::V3 if supports_binary => {
+                use bytes::BytesMut;
+
+                let size_hint = packet.get_size_hint(!supports_binary) + 2;
+                let mut bytes = BytesMut::with_capacity(size_hint);
+                encoder::v3_bin_packet_encoder(packet, &mut bytes);
+                bytes.freeze()
+            }
+            ProtocolVersion::V3 => {
+                use bytes::BytesMut;
+
+                let size_hint = packet.get_size_hint(!supports_binary) + 2;
+                let mut bytes = BytesMut::with_capacity(size_hint);
+                encoder::v3_string_packet_encoder(packet, &mut bytes);
+                bytes.freeze()
+            }
+        }
+    }
+
+    #[cfg(not(feature = "v3"))]
+    packet.into()
+}

crates/engineioxide/src/transport/ws.rs

@@ -320,7 +320,18 @@ where
     S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 {
     #[cfg(feature = "tracing")]
-    tracing::debug!("websocket connection upgrade");
+    tracing::debug!("starting websocket connection upgrade");
+
+    // The socket is upgrading, from now on all polling
+    // request will return a `Packet::Noop`
+    socket.start_upgrade();
+
+    // We send a last Noop request to close a potential waiting polling request.
+    socket.send(Packet::Noop)?;
+
+    // wait for any current polling connection to finish by waiting for the socket to be unlocked
+    // All other polling connection will be immediately closed with a NOOP packet.
+    let _lock = socket.internal_rx.lock().await;
 
     // Fetch the next packet from the ws stream, it should be a PingUpgrade packet
     let msg = match ws.next().await {
@@ -329,16 +340,15 @@ where
     };
     match Packet::try_from(msg)? {
         Packet::PingUpgrade => {
-            socket.start_upgrade();
+            #[cfg(feature = "tracing")]
+            tracing::debug!("received first ping upgrade");
+
             // Respond with a PongUpgrade packet
             ws.send(Message::Text(Packet::PongUpgrade.into())).await?;
         }
         p => Err(Error::BadPacket(p))?,
     };
 
-    // send a NOOP packet to any pending polling request so it closes gracefully
-    socket.send(Packet::Noop)?;
-
     // Fetch the next packet from the ws stream, it should be an Upgrade packet
     let msg = match ws.next().await {
         Some(Ok(Message::Text(d))) => d,
@@ -356,13 +366,11 @@ where
     match Packet::try_from(msg)? {
         Packet::Upgrade => {
             #[cfg(feature = "tracing")]
-            tracing::debug!("ws upgraded successful")
+            tracing::debug!("ws upgraded successfully")
         }
         p => Err(Error::BadPacket(p))?,
     };
 
-    // wait for any polling connection to finish by waiting for the socket to be unlocked
-    let _ = socket.internal_rx.lock().await;
     socket.upgrade_to_websocket();
     Ok(())
 }