diff --git a/Cargo.lock b/Cargo.lock index 5b599e6b..1782324e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -815,7 +815,6 @@ name = "engineioxide" version = "0.17.5" dependencies = [ "axum", - "base64", "bytes", "codspeed-criterion-compat", "engineioxide-core", @@ -826,8 +825,6 @@ dependencies = [ "http-body-util", "hyper", "hyper-util", - "itoa", - "memchr", "pin-project-lite", "serde", "serde_json", @@ -849,8 +846,20 @@ version = "0.2.1" dependencies = [ "base64", "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "itoa", + "memchr", + "pin-project-lite", "rand 0.10.1", "serde", + "serde_json", + "smallvec", + "tokio", + "tokio-stream", + "tracing", ] [[package]] diff --git a/crates/engineioxide-core/Cargo.toml b/crates/engineioxide-core/Cargo.toml index 485e57e3..cd4b642d 100644 --- a/crates/engineioxide-core/Cargo.toml +++ b/crates/engineioxide-core/Cargo.toml @@ -17,6 +17,29 @@ rand = "0.10" base64 = "0.22" serde.workspace = true bytes.workspace = true +serde_json.workspace = true +http-body.workspace = true +http-body-util.workspace = true +http.workspace = true +futures-util.workspace = true +smallvec.workspace = true +pin-project-lite.workspace = true + +# Engine.io V3 payload +itoa = { workspace = true, optional = true } +memchr = { version = "2.7", optional = true } + +# Tracing +tracing = { workspace = true, optional = true } + + +[features] +v3 = ["dep:memchr", "dep:itoa"] +tracing = ["dep:tracing"] + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt", "time", "sync"] } +tokio-stream.workspace = true [lints] workspace = true diff --git a/crates/engineioxide-core/src/lib.rs b/crates/engineioxide-core/src/lib.rs index 9b5fff8f..fe94cc6b 100644 --- a/crates/engineioxide-core/src/lib.rs +++ b/crates/engineioxide-core/src/lib.rs @@ -1,8 +1,13 @@ #![warn(missing_docs)] #![doc = include_str!("../README.md")] +mod packet; +mod protocol; mod sid; mod str; +pub use packet::{OpenPacket, Packet, PacketBuf, PacketParseError}; +pub use protocol::{ProtocolVersion, TransportType, UnknownTransportError}; pub use sid::Sid; pub use str::Str; +pub mod payload; diff --git a/crates/engineioxide/src/packet.rs b/crates/engineioxide-core/src/packet.rs similarity index 68% rename from crates/engineioxide/src/packet.rs rename to crates/engineioxide-core/src/packet.rs index ead18cae..0d16416d 100644 --- a/crates/engineioxide/src/packet.rs +++ b/crates/engineioxide-core/src/packet.rs @@ -1,16 +1,14 @@ +use std::{fmt, time::Duration}; + use base64::{Engine, engine::general_purpose}; use bytes::Bytes; -use engineioxide_core::{Sid, Str}; use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use smallvec::{SmallVec, smallvec}; -use std::time::Duration; +use smallvec::SmallVec; -use crate::config::EngineIoConfig; -use crate::errors::Error; -use crate::{ProtocolVersion, TransportType}; +use crate::{ProtocolVersion, Sid, Str, TransportType}; /// A Packet type to use when receiving and sending data from the client -#[derive(Debug, Clone, PartialEq, PartialOrd)] +#[derive(Debug, Clone, PartialEq)] pub enum Packet { /// Open packet used to initiate a connection Open(OpenPacket), @@ -53,6 +51,67 @@ pub enum Packet { BinaryV3(Bytes), // Not part of the protocol, used internally } +/// An error that occurs when parsing a packet. +#[derive(Debug)] +pub enum PacketParseError { + /// Invalid connect packet + InvalidConnectPacket(serde_json::Error), + /// The packet type is invalid. + InvalidPacketType(Option), + /// The packet payload is invalid. + InvalidPacketPayload, + /// The packet length is invalid. + InvalidPacketLen, + /// The packet chunk is invalid + InvalidUtf8Boundary(std::str::Utf8Error), + /// The base64 decoding failed. + Base64Decode(base64::DecodeError), + /// The payload is too large. + PayloadTooLarge { + /// The maximum allowed payload size. + max: u64, + }, +} +impl fmt::Display for PacketParseError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + PacketParseError::InvalidConnectPacket(e) => write!(f, "invalid connect packet: {e}"), + PacketParseError::InvalidPacketType(c) => write!(f, "invalid packet type: {c:?}"), + PacketParseError::InvalidPacketPayload => write!(f, "invalid packet payload"), + PacketParseError::InvalidPacketLen => write!(f, "invalid packet length"), + PacketParseError::InvalidUtf8Boundary(err) => write!( + f, + "invalid utf8 boundary when parsing payload into packet chunks: {err}" + ), + PacketParseError::Base64Decode(err) => write!(f, "base64 decode error: {err}"), + PacketParseError::PayloadTooLarge { max } => { + write!(f, "payload too large: max {max}") + } + } + } +} +impl From for PacketParseError { + fn from(err: base64::DecodeError) -> Self { + PacketParseError::Base64Decode(err) + } +} +impl From for PacketParseError { + fn from(err: std::string::FromUtf8Error) -> Self { + PacketParseError::InvalidUtf8Boundary(err.utf8_error()) + } +} +impl From for PacketParseError { + fn from(err: std::str::Utf8Error) -> Self { + PacketParseError::InvalidUtf8Boundary(err) + } +} +impl From for PacketParseError { + fn from(err: serde_json::Error) -> Self { + PacketParseError::InvalidConnectPacket(err) + } +} +impl std::error::Error for PacketParseError {} + impl Packet { /// Check if the packet is a binary packet pub fn is_binary(&self) -> bool { @@ -60,7 +119,7 @@ impl Packet { } /// If the packet is a message packet (text), it returns the message - pub(crate) fn into_message(self) -> Str { + pub fn into_message(self) -> Str { match self { Packet::Message(msg) => msg, _ => panic!("Packet is not a message"), @@ -68,7 +127,7 @@ impl Packet { } /// If the packet is a binary packet, it returns the binary data - pub(crate) fn into_binary(self) -> Bytes { + pub fn into_binary(self) -> Bytes { match self { Packet::Binary(data) => data, Packet::BinaryV3(data) => data, @@ -81,7 +140,7 @@ impl Packet { /// If b64 is true, it returns the max size when serialized to base64 /// /// The base64 max size factor is `ceil(n / 3) * 4` - pub(crate) fn get_size_hint(&self, b64: bool) -> usize { + pub fn get_size_hint(&self, b64: bool) -> usize { match self { Packet::Open(_) => 156, // max possible size for the open packet serialized Packet::Close => 1, @@ -110,6 +169,12 @@ impl Packet { } } +impl From for Bytes { + fn from(value: Packet) -> Self { + String::from(value).into() + } +} + /// Serialize a [Packet] to a [String] according to the Engine.IO protocol impl From for String { fn from(packet: Packet) -> String { @@ -143,25 +208,19 @@ impl From for String { buffer } } -impl From for tokio_tungstenite::tungstenite::Utf8Bytes { - fn from(value: Packet) -> Self { - String::from(value).into() - } -} -impl From for Bytes { - fn from(value: Packet) -> Self { - String::from(value).into() - } -} +/// Deserialize a [Packet] from a [String] according to the Engine.IO protocol impl Packet { /// Parses a packet from a string value using the specified protocol version. - pub fn parse(protocol: ProtocolVersion, value: impl Into) -> Result { + pub fn parse( + protocol: ProtocolVersion, + value: impl Into, + ) -> Result { let value = value.into(); let packet_type = value .as_bytes() .first() - .ok_or(Error::InvalidPacketType(None))?; + .ok_or(PacketParseError::InvalidPacketType(None))?; let is_upgrade = value.len() == 6 && &value[1..6] == "probe"; let res = match packet_type { b'1' => Packet::Close, @@ -182,29 +241,39 @@ impl Packet { .decode(value.slice(1..).as_bytes())? .into(), ), - c => Err(Error::InvalidPacketType(Some(*c as char)))?, + c => Err(PacketParseError::InvalidPacketType(Some(*c as char)))?, }; Ok(res) } } /// An OpenPacket is used to initiate a connection -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, PartialOrd)] +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "camelCase")] pub struct OpenPacket { - sid: Sid, - upgrades: SmallVec<[TransportType; 1]>, + /// The session ID. + pub sid: Sid, + + /// The list of available transport upgrades. + pub upgrades: SmallVec<[TransportType; 1]>, + + /// The ping interval, used in the heartbeat mechanism. #[serde( serialize_with = "serialize_duration_millis", deserialize_with = "deserialize_duration_from_millis" )] - ping_interval: Duration, + pub ping_interval: Duration, + + /// The ping timeout, used in the heartbeat mechanism. #[serde( serialize_with = "serialize_duration_millis", deserialize_with = "deserialize_duration_from_millis" )] - ping_timeout: Duration, - max_payload: u64, + pub ping_timeout: Duration, + + /// The maximum number of bytes per chunk, used by the client to + /// aggregate packets into payloads. + pub max_payload: u64, } /// Helper to serialize a duration as milliseconds @@ -225,28 +294,28 @@ where Ok(Duration::from_millis(millis)) } -impl OpenPacket { - /// Create a new [OpenPacket] - /// If the current transport is polling, the server will always allow the client to upgrade to websocket - pub fn new(transport: TransportType, sid: Sid, config: &EngineIoConfig) -> Self { - let upgrades = if transport == TransportType::Polling { - smallvec![TransportType::Websocket] - } else { - smallvec![] - }; - OpenPacket { - sid, - upgrades, - ping_interval: config.ping_interval, - ping_timeout: config.ping_timeout, - max_payload: config.max_payload, +/// This default implementation should only be used for testing purposes. +impl Default for OpenPacket { + fn default() -> Self { + Self { + sid: Sid::ZERO, + upgrades: smallvec::smallvec![TransportType::Websocket], + ping_interval: Duration::from_millis(25000), + ping_timeout: Duration::from_millis(20000), + max_payload: 100000, } } } +/// Buffered packets to send to the client. +/// It is used to ensure atomicity when sending multiple packets to the client. +/// +/// The [`PacketBuf`] stack size will impact the dynamically allocated buffer +/// of the internal mpsc channel. +pub type PacketBuf = SmallVec<[Packet; 2]>; + #[cfg(test)] mod tests { - use crate::config::EngineIoConfig; use super::*; use std::time::Duration; @@ -254,11 +323,13 @@ mod tests { #[test] fn test_open_packet() { let sid = Sid::new(); - let packet = Packet::Open(OpenPacket::new( - TransportType::Polling, + let packet = Packet::Open(OpenPacket { sid, - &EngineIoConfig::default(), - )); + upgrades: smallvec::smallvec![TransportType::Websocket], + ping_interval: Duration::from_millis(25000), + ping_timeout: Duration::from_millis(20000), + max_payload: 100000, + }); let packet_str: String = packet.into(); assert_eq!( packet_str, @@ -275,12 +346,6 @@ mod tests { assert_eq!(packet_str, "4hello"); } - #[test] - fn test_message_packet_deserialize() { - let packet = Packet::parse(ProtocolVersion::V4, "4hello").unwrap(); - assert_eq!(packet, Packet::Message("hello".into())); - } - #[test] fn test_binary_packet() { let packet = Packet::Binary(vec![1, 2, 3].into()); @@ -288,12 +353,6 @@ mod tests { assert_eq!(packet_str, "bAQID"); } - #[test] - fn test_binary_packet_deserialize() { - let packet = Packet::parse(ProtocolVersion::V4, "bAQID").unwrap(); - assert_eq!(packet, Packet::Binary(vec![1, 2, 3].into())); - } - #[test] fn test_binary_packet_v4_deserialize_payload_starting_with_4() { let data = vec![0xE0, 0xE1, 0xE2]; @@ -312,27 +371,16 @@ mod tests { assert_eq!(packet_str, "b4AQID"); } - #[test] - fn test_binary_packet_v3_deserialize() { - let packet = Packet::parse(ProtocolVersion::V3, "b4AQID").unwrap(); - assert_eq!(packet, Packet::BinaryV3(vec![1, 2, 3].into())); - } - #[test] fn test_packet_get_size_hint() { // Max serialized packet - let open = OpenPacket::new( - TransportType::Polling, - Sid::new(), - &EngineIoConfig { - max_buffer_size: usize::MAX, - max_payload: u64::MAX, - ping_interval: Duration::MAX, - ping_timeout: Duration::MAX, - transports: TransportType::Polling as u8 | TransportType::Websocket as u8, - ..Default::default() - }, - ); + let open = OpenPacket { + sid: Sid::new(), + ping_interval: Duration::MAX, + ping_timeout: Duration::MAX, + max_payload: u64::MAX, + upgrades: smallvec::smallvec![TransportType::Websocket], + }; let size = serde_json::to_string(&open).unwrap().len(); let packet = Packet::Open(open); assert_eq!(packet.get_size_hint(false), size); diff --git a/crates/engineioxide/src/transport/polling/payload/buf.rs b/crates/engineioxide-core/src/payload/buf.rs similarity index 100% rename from crates/engineioxide/src/transport/polling/payload/buf.rs rename to crates/engineioxide-core/src/payload/buf.rs diff --git a/crates/engineioxide/src/transport/polling/payload/decoder.rs b/crates/engineioxide-core/src/payload/decoder.rs similarity index 88% rename from crates/engineioxide/src/transport/polling/payload/decoder.rs rename to crates/engineioxide-core/src/payload/decoder.rs index a748fd81..cc152308 100644 --- a/crates/engineioxide/src/transport/polling/payload/decoder.rs +++ b/crates/engineioxide-core/src/payload/decoder.rs @@ -5,11 +5,9 @@ //! - v3_decoder: Decodes the payload stream according to the [engine.io v3 protocol](https://github.com/socketio/engine.io-protocol/tree/v3#payload) //! -use futures_core::Stream; -use futures_util::StreamExt; -use http::StatusCode; +use crate::{Packet, PacketParseError, ProtocolVersion}; +use futures_util::{Stream, StreamExt}; -use crate::{ProtocolVersion, errors::Error, packet::Packet}; use bytes::Buf; use http_body::Body; use http_body_util::BodyStream; @@ -22,9 +20,6 @@ struct Payload { buffer: BufList, end_of_stream: bool, current_payload_size: u64, - - #[cfg(feature = "v3")] - yield_packets: u32, } impl Payload { @@ -34,15 +29,13 @@ impl Payload { buffer: BufList::new(), end_of_stream: false, current_payload_size: 0, - #[cfg(feature = "v3")] - yield_packets: 0, } } } /// Polls the body stream for data and adds it to the chunk list in the state /// Returns an error if the packet length exceeds the maximum allowed payload size -async fn poll_body(state: &mut Payload, max_payload: u64) -> Result<(), Error> +async fn poll_body(state: &mut Payload, max_payload: u64) -> Result<(), PacketParseError> where B: Body + Unpin, E: std::fmt::Debug, @@ -59,7 +52,7 @@ where Err(_e) => { #[cfg(feature = "tracing")] tracing::debug!("error reading body stream: {:?}", _e); - Err(Error::HttpErrorResponse(StatusCode::BAD_REQUEST)) + Err(PacketParseError::InvalidPacketPayload) } }?; if state.current_payload_size + (data.remaining() as u64) <= max_payload { @@ -67,11 +60,14 @@ where state.buffer.push(data); Ok(()) } else { - Err(Error::PayloadTooLarge) + Err(PacketParseError::PayloadTooLarge { max: max_payload }) } } -pub fn v4_decoder(body: B, max_payload: u64) -> impl Stream> +pub fn v4_decoder( + body: B, + max_payload: u64, +) -> impl Stream> where B: Body + Unpin, E: std::fmt::Debug, @@ -93,11 +89,14 @@ where } // Read from the buffer until the packet separator is found - if let Err(e) = (&mut state.buffer) + if let Err(_err) = (&mut state.buffer) .reader() .read_until(PACKET_SEPARATOR_V4, &mut packet_buf) { - break Some((Err(Error::Io(e)), state)); + #[cfg(feature = "tracing")] + tracing::debug!("failed to read packet payload: {_err}"); + + break Some((Err(PacketParseError::InvalidPacketPayload), state)); } let separator_found = packet_buf.ends_with(&[PACKET_SEPARATOR_V4]); @@ -111,7 +110,7 @@ where || (state.end_of_stream && state.buffer.remaining() == 0 && !packet_buf.is_empty()) { let packet = String::from_utf8(packet_buf) - .map_err(|_| Error::InvalidPacketLength) + .map_err(PacketParseError::from) .and_then(|v| Packet::parse(ProtocolVersion::V4, v)); // Convert the packet buffer to a Packet object break Some((packet, state)); // Emit the packet and the updated state } else if state.end_of_stream && state.buffer.remaining() == 0 { @@ -125,14 +124,14 @@ where pub fn v3_binary_decoder( body: B, max_payload: u64, -) -> impl Stream> +) -> impl Stream> where B: Body + Unpin, E: std::fmt::Debug, { use std::io::Read; - use crate::transport::polling::payload::{ + use crate::payload::{ BINARY_PACKET_IDENTIFIER_V3, BINARY_PACKET_SEPARATOR_V3, STRING_PACKET_IDENTIFIER_V3, }; @@ -156,11 +155,14 @@ where // If there is no packet_type found if packet_type.is_none() && state.buffer.remaining() > 0 { // Read from the buffer until the packet separator is found - if let Err(e) = (&mut state.buffer) + if let Err(_err) = (&mut state.buffer) .reader() .read_until(BINARY_PACKET_SEPARATOR_V3, &mut packet_buf) { - break Some((Err(Error::Io(e)), state)); + #[cfg(feature = "tracing")] + tracing::debug!("failed to read packet payload: {_err}"); + + break Some((Err(PacketParseError::InvalidPacketPayload), state)); } // Extract packet_type and packet_size @@ -173,11 +175,11 @@ where Some(&STRING_PACKET_IDENTIFIER_V3) => { packet_type = Some(STRING_PACKET_IDENTIFIER_V3) } - _ => break Some((Err(Error::InvalidPacketLength), state)), + _ => break Some((Err(PacketParseError::InvalidPacketLen), state)), } if packet_buf.len() > 9 { - break Some((Err(Error::InvalidPacketLength), state)); + break Some((Err(PacketParseError::InvalidPacketLen), state)); } let size_str = &packet_buf[1..] @@ -187,7 +189,7 @@ where if let Ok(size) = size_str.parse() { packet_size = size; } else { - break Some((Err(Error::InvalidPacketLength), state)); + break Some((Err(PacketParseError::InvalidPacketLen), state)); } packet_buf.clear(); } @@ -204,10 +206,10 @@ where // Read the packet data let packet = match packet_type.unwrap() { STRING_PACKET_IDENTIFIER_V3 => String::from_utf8(packet_buf) - .map_err(|_| Error::InvalidPacketLength) + .map_err(PacketParseError::from) .and_then(|v| Packet::parse(ProtocolVersion::V3, v)), // Convert the packet buffer to a Packet object BINARY_PACKET_IDENTIFIER_V3 => Ok(Packet::BinaryV3(packet_buf.into())), - _ => Err(Error::InvalidPacketLength), + _ => Err(PacketParseError::InvalidPacketLen), }; break Some((packet, state)); @@ -216,7 +218,7 @@ where } else if state.end_of_stream { // EOS reached with leftover bytes that cannot form a complete // packet (truncated header or truncated body). - break Some((Err(Error::InvalidPacketLength), state)); + break Some((Err(PacketParseError::InvalidPacketLen), state)); } } }) @@ -250,10 +252,10 @@ fn utf16_len(s: &str) -> usize { pub fn v3_string_decoder( body: impl Body + Unpin, max_payload: u64, -) -> impl Stream> { +) -> impl Stream> { use std::io::ErrorKind; - use crate::transport::polling::payload::STRING_PACKET_SEPARATOR_V3; + use crate::payload::STRING_PACKET_SEPARATOR_V3; #[cfg(feature = "tracing")] tracing::debug!("decoding payload with v3 string decoder"); @@ -269,10 +271,14 @@ pub fn v3_string_decoder( { break Some((Err(e), state)); } - if state.end_of_stream && state.buffer.remaining() == 0 && state.yield_packets > 0 { + if state.end_of_stream && state.buffer.remaining() == 0 { + if !packet_buf.is_empty() { + // Leftover unparsed bytes at end of stream: either a length + // token that never found its `:` separator or a + // truncated packet body. + break Some((Err(PacketParseError::InvalidPacketLen), state)); + } break None; // Reached end of stream with no more data, end the stream - } else if state.end_of_stream && state.buffer.remaining() == 0 { - return Some((Err(Error::InvalidPacketLength), state)); } let mut reader = (&mut state.buffer).reader(); @@ -285,7 +291,12 @@ pub fn v3_string_decoder( let available = match reader.fill_buf() { Ok(n) => n, Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, - Err(e) => return Some((Err(Error::Io(e)), state)), + Err(_err) => { + #[cfg(feature = "tracing")] + tracing::debug!("failed to read packet payload: {_err}"); + + return Some((Err(PacketParseError::InvalidPacketPayload), state)); + } }; let old_len = packet_buf.len(); packet_buf.extend_from_slice(available); @@ -294,9 +305,10 @@ pub fn v3_string_decoder( Some(i) => { // Extract the packet length from the available data packet_utf16_len = match std::str::from_utf8(&packet_buf[..i]) - .map_err(|_| Error::InvalidPacketLength) + .map_err(PacketParseError::from) .and_then(|s| { - s.parse::().map_err(|_| Error::InvalidPacketLength) + s.parse::() + .map_err(|_| PacketParseError::InvalidPacketLen) }) { Ok(size) => size, Err(e) => return Some((Err(e), state)), @@ -306,7 +318,7 @@ pub fn v3_string_decoder( (true, i + 1 - old_len) // Mark as done and set the used bytes count } None if state.end_of_stream && remaining - available.len() == 0 => { - return Some((Err(Error::InvalidPacketLength), state)); + return Some((Err(PacketParseError::InvalidPacketLen), state)); } // Reached end of stream and end of bufferered chunks without finding the separator None => (false, available.len()), // Continue reading more data } @@ -365,8 +377,7 @@ pub fn v3_string_decoder( // SAFETY: packet_buf is a valid utf8 string checkd above let packet = unsafe { String::from_utf8_unchecked(packet_buf) }; let packet = Packet::parse(ProtocolVersion::V3, packet) - .map_err(|_| Error::InvalidPacketLength); - state.yield_packets += 1; + .map_err(|_| PacketParseError::InvalidPacketLen); break Some((packet, state)); // Emit the packet and the updated state } } else if state.end_of_stream && state.buffer.remaining() == 0 { @@ -384,8 +395,6 @@ mod tests { use http_body::Frame; use http_body_util::{Full, StreamBody}; - use crate::packet::Packet; - use super::*; const MAX_PAYLOAD: u64 = 100_000; @@ -450,7 +459,10 @@ mod tests { let payload = v4_decoder(stream, MAX_PAYLOAD); futures_util::pin_mut!(payload); let packet = payload.next().await.unwrap(); - assert!(matches!(packet, Err(Error::PayloadTooLarge))); + assert!(matches!( + packet, + Err(PacketParseError::PayloadTooLarge { max: MAX_PAYLOAD }) + )); } } @@ -545,6 +557,18 @@ mod tests { } } + #[cfg(feature = "v3")] + #[tokio::test] + async fn string_invalid_packet_format_v3() { + let data = Full::new(Bytes::from_static(b"abc")); + let payload = v3_string_decoder(data, MAX_PAYLOAD); + futures_util::pin_mut!(payload); + assert!(matches!( + payload.next().await, + Some(Err(crate::PacketParseError::InvalidPacketLen)) + )); + } + #[cfg(feature = "v3")] #[tokio::test] async fn binary_payload_stream_v3() { @@ -607,7 +631,10 @@ mod tests { let payload = v3_binary_decoder(stream, MAX_PAYLOAD); futures_util::pin_mut!(payload); let packet = payload.next().await.unwrap(); - assert!(matches!(packet, Err(Error::PayloadTooLarge))); + assert!(matches!( + packet, + Err(PacketParseError::PayloadTooLarge { max: MAX_PAYLOAD }) + )); } for i in 1..DATA.len() { let stream = StreamBody::new(futures_util::stream::iter( @@ -618,7 +645,10 @@ mod tests { let payload = v3_string_decoder(stream, MAX_PAYLOAD); futures_util::pin_mut!(payload); let packet = payload.next().await.unwrap(); - assert!(matches!(packet, Err(Error::PayloadTooLarge))); + assert!(matches!( + packet, + Err(PacketParseError::PayloadTooLarge { max: MAX_PAYLOAD }) + )); } } diff --git a/crates/engineioxide/src/transport/polling/payload/encoder.rs b/crates/engineioxide-core/src/payload/encoder.rs similarity index 61% rename from crates/engineioxide/src/transport/polling/payload/encoder.rs rename to crates/engineioxide-core/src/payload/encoder.rs index a56a9c9a..af0ac536 100644 --- a/crates/engineioxide/src/transport/polling/payload/encoder.rs +++ b/crates/engineioxide-core/src/payload/encoder.rs @@ -7,30 +7,28 @@ //! * binary encoder (used when there are binary packets and the client supports binary) //! -use tokio::sync::MutexGuard; +use std::pin::Pin; + +use futures_util::{FutureExt, Stream, StreamExt}; +use smallvec::smallvec; use crate::{ - errors::Error, packet::Packet, peekable::PeekableReceiver, socket::PacketBuf, - transport::polling::payload::Payload, + packet::PacketBuf, + payload::{Payload, peekable::Peekable}, }; -/// Try to immediately poll a new packet buf from the rx channel and check that the new packet can be added to the payload -/// -/// Manually close the channel if the packet is a close packet -/// It will allow to notify the [`Socket`](crate::socket::Socket) that the session is closed -/// -/// ## Arguments -/// * `rx` - The channel to poll -/// * `payload_len` - The current payload length -/// * `max_payload` - The maximum payload length -/// * `b64` - If binary packets should be encoded in base64 -fn try_recv_packet( - rx: &mut MutexGuard<'_, PeekableReceiver>, +#[cfg(feature = "v3")] +use crate::Packet; + +/// Try to immediately poll a new packet buf from the rx channel and check +/// that the new packet can be added to the payload without exceeding the max payload size +fn try_poll_packet( + mut rx: Pin<&mut Peekable>>, payload_len: usize, max_payload: u64, b64: bool, ) -> Option { - if let Some(packets) = rx.peek() { + if let Some(packets) = rx.as_mut().peek().now_or_never().flatten() { let size = packets.iter().map(|p| p.get_size_hint(b64)).sum::(); if (payload_len + size) as u64 > max_payload { #[cfg(feature = "tracing")] @@ -39,44 +37,30 @@ fn try_recv_packet( } } - let packets = rx.try_recv().ok(); - - if Some(&Packet::Close) == packets.as_ref().and_then(|p| p.first()) { - #[cfg(feature = "tracing")] - tracing::debug!("Received close packet, closing channel"); - rx.try_recv().ok(); - rx.close(); - } + let packets = rx.next().now_or_never().flatten(); #[cfg(feature = "tracing")] tracing::debug!("sending packet: {:?}", packets); packets } -/// Same as [`try_recv_packet`] +/// Same as [`try_poll_packet`] /// but wait for a new packet if there is no packet in the buffer -async fn recv_packet( - rx: &mut MutexGuard<'_, PeekableReceiver>, -) -> Result { - let packet = rx.recv().await.ok_or(Error::Aborted)?; - if Some(&Packet::Close) == packet.first() { - #[cfg(feature = "tracing")] - tracing::debug!("Received close packet, closing channel"); - rx.close(); - } +async fn poll_packet(mut rx: Pin<&mut Peekable>>) -> PacketBuf { + let packet = rx.next().await.unwrap_or(smallvec![]); // if the channel is closed yield an empty packet #[cfg(feature = "tracing")] tracing::debug!("sending packet: {:?}", packet); - Ok(packet) + packet } /// Encode multiple packets into a string payload according to the /// [engine.io v4 protocol](https://socket.io/fr/docs/v4/engine-io-protocol/#http-long-polling-1) pub async fn v4_encoder( - mut rx: MutexGuard<'_, PeekableReceiver>, + mut rx: Pin<&mut Peekable>>, max_payload: u64, -) -> Result { - use crate::transport::polling::payload::PACKET_SEPARATOR_V4; +) -> Payload { + use crate::payload::PACKET_SEPARATOR_V4; #[cfg(feature = "tracing")] tracing::debug!("encoding payload with v4 encoder"); @@ -85,7 +69,7 @@ pub async fn v4_encoder( // Send all packets in the buffer const PUNCTUATION_LEN: usize = 1; while let Some(packets) = - try_recv_packet(&mut rx, data.len() + PUNCTUATION_LEN, max_payload, true) + try_poll_packet(rx.as_mut(), data.len() + PUNCTUATION_LEN, max_payload, true) { for packet in packets { let packet: String = packet.into(); @@ -99,7 +83,7 @@ pub async fn v4_encoder( // If there is no packet in the buffer, wait for the next packet if data.is_empty() { - let packets = recv_packet(&mut rx).await?; + let packets = poll_packet(rx.as_mut()).await; for packet in packets { if !data.is_empty() { data.push(std::char::from_u32(PACKET_SEPARATOR_V4 as u32).unwrap()); @@ -110,14 +94,14 @@ pub async fn v4_encoder( } } - Ok(Payload::new(data.into(), false)) + Payload::new(data.into(), false) } /// 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) { - use crate::transport::polling::payload::BINARY_PACKET_SEPARATOR_V3; + use crate::payload::BINARY_PACKET_SEPARATOR_V3; use bytes::BufMut; let mut itoa = itoa::Buffer::new(); @@ -157,7 +141,7 @@ pub fn v3_bin_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) { /// [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) { - use crate::transport::polling::payload::STRING_PACKET_SEPARATOR_V3; + use crate::payload::STRING_PACKET_SEPARATOR_V3; use bytes::BufMut; let packet: String = packet.into(); let packet = format!( @@ -173,9 +157,9 @@ pub fn v3_string_packet_encoder(packet: Packet, data: &mut bytes::BytesMut) { /// according to the [engine.io v3 protocol](https://github.com/socketio/engine.io-protocol/tree/v3#payload) #[cfg(feature = "v3")] pub async fn v3_binary_encoder( - mut rx: MutexGuard<'_, PeekableReceiver>, + mut rx: Pin<&mut Peekable>>, max_payload: u64, -) -> Result { +) -> Payload { let mut data = bytes::BytesMut::new(); let mut packet_buffer: Vec = Vec::new(); @@ -189,7 +173,7 @@ pub async fn v3_binary_encoder( // buffer all packets to find if there is binary packets let mut has_binary = false; - while let Some(packets) = try_recv_packet(&mut rx, estimated_size, max_payload, false) { + while let Some(packets) = try_poll_packet(rx.as_mut(), estimated_size, max_payload, false) { for packet in packets { if packet.is_binary() { has_binary = true; @@ -214,7 +198,7 @@ pub async fn v3_binary_encoder( // If there is no packet in the buffer, wait for the next packet if data.is_empty() { - let packets = recv_packet(&mut rx).await?; + let packets = poll_packet(rx.as_mut()).await; has_binary = packets.iter().any(|p| p.is_binary()); for packet in packets { if has_binary { @@ -227,16 +211,16 @@ pub async fn v3_binary_encoder( #[cfg(feature = "tracing")] tracing::debug!("sending packet: {:?}", &data); - Ok(Payload::new(data.freeze(), has_binary)) + Payload::new(data.freeze(), has_binary) } /// Encode multiple packet 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 async fn v3_string_encoder( - mut rx: MutexGuard<'_, PeekableReceiver>, + mut rx: Pin<&mut Peekable>>, max_payload: u64, -) -> Result { +) -> Payload { let mut data = bytes::BytesMut::new(); #[cfg(feature = "tracing")] @@ -245,8 +229,8 @@ pub async fn v3_string_encoder( const PUNCTUATION_LEN: usize = 2; // number of digits of the max packet size, used to approximate the payload size let max_packet_size_len = max_payload.checked_ilog10().unwrap_or(0) as usize + 1; - while let Some(packets) = try_recv_packet( - &mut rx, + while let Some(packets) = try_poll_packet( + rx.as_mut(), data.len() + PUNCTUATION_LEN + max_packet_size_len, max_payload, true, @@ -258,21 +242,19 @@ pub async fn v3_string_encoder( // If there is no packet in the buffer, wait for the next packet if data.is_empty() { - let packets = recv_packet(&mut rx).await?; + let packets = poll_packet(rx.as_mut()).await; for packet in packets { v3_string_packet_encoder(packet, &mut data); } } - Ok(Payload::new(data.freeze(), false)) + Payload::new(data.freeze(), false) } #[cfg(test)] mod tests { use bytes::Bytes; - use tokio::sync::Mutex; - - use PacketBuf; + use futures_util::stream; use super::*; const MAX_PAYLOAD: u64 = 100_000; @@ -280,18 +262,15 @@ mod tests { #[tokio::test] async fn encode_v4_payload() { const PAYLOAD: &str = "4hello€\x1ebAQIDBA==\x1e4hello€"; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let rx = Mutex::new(PeekableReceiver::new(rx)); - let rx = rx.lock().await; - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Binary(Bytes::from_static(&[ - 1, 2, 3, 4 - ]))]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD).await.unwrap(); + + let rx = stream::iter([ + smallvec![Packet::Message("hello€".into())], + smallvec![Packet::Binary(Bytes::from_static(&[1, 2, 3, 4]))], + smallvec![Packet::Message("hello€".into())], + ]); + let rx = std::pin::pin!(Peekable::new(rx)); + + let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD).await; assert_eq!(data, PAYLOAD.as_bytes()); } @@ -299,8 +278,8 @@ mod tests { async fn encode_v4_payload_parked_poll_multi_packet_batch() { const PAYLOAD: &str = "4hello€\x1ebAQIDBA=="; let (tx, rx) = tokio::sync::mpsc::channel::(10); - let rx = Mutex::new(PeekableReceiver::new(rx)); - let rx = rx.lock().await; + let rx = tokio_stream::wrappers::ReceiverStream::new(rx); + let rx = std::pin::pin!(Peekable::new(rx)); tokio::spawn(async move { tokio::time::sleep(std::time::Duration::from_millis(50)).await; tx.try_send(smallvec::smallvec![ @@ -309,38 +288,33 @@ mod tests { ]) .unwrap(); }); - let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD).await.unwrap(); - assert_eq!(data, PAYLOAD); + let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD).await; + assert_eq!(data, PAYLOAD.as_bytes()); } #[tokio::test] async fn max_payload_v4() { const MAX_PAYLOAD: u64 = 10; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Binary(Bytes::from_static(&[ - 1, 2, 3, 4 - ]))]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); + + let rx = stream::iter([ + smallvec![Packet::Message("hello€".into())], + smallvec![Packet::Binary(Bytes::from_static(&[1, 2, 3, 4]))], + smallvec![Packet::Message("hello€".into())], + smallvec![Packet::Message("hello€".into())], + ]); + + let mut rx = std::pin::pin!(Peekable::new(rx)); + { - let rx = mutex.lock().await; - let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD).await.unwrap(); + let Payload { data, .. } = v4_encoder(rx.as_mut(), MAX_PAYLOAD).await; assert_eq!(data, "4hello€".as_bytes()); } { - let rx = mutex.lock().await; - let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD + 10).await.unwrap(); + let Payload { data, .. } = v4_encoder(rx.as_mut(), MAX_PAYLOAD + 10).await; assert_eq!(data, "bAQIDBA==\x1e4hello€".as_bytes()); } { - let rx = mutex.lock().await; - let Payload { data, .. } = v4_encoder(rx, MAX_PAYLOAD + 10).await.unwrap(); + let Payload { data, .. } = v4_encoder(rx.as_mut(), MAX_PAYLOAD + 10).await; assert_eq!(data, "4hello€".as_bytes()); } } @@ -353,12 +327,9 @@ mod tests { // (packet type '4' + non-BMP codepoint U+1D54A) has 2 codepoints // but 3 UTF-16 code units. const PAYLOAD: &str = "3:4𝕊"; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - let rx = mutex.lock().await; - tx.try_send(smallvec::smallvec![Packet::Message("𝕊".into())]) - .unwrap(); - let Payload { data, .. } = v3_string_encoder(rx, MAX_PAYLOAD).await.unwrap(); + let rx = stream::iter([smallvec![Packet::Message("𝕊".into())]]); + let rx = std::pin::pin!(Peekable::new(rx)); + let Payload { data, .. } = v3_string_encoder(rx, MAX_PAYLOAD).await; assert_eq!(data, PAYLOAD.as_bytes()); } @@ -366,21 +337,16 @@ mod tests { #[tokio::test] async fn encode_v3b64_payload() { const PAYLOAD: &str = "7:4hello€10:b4AQIDBA==7:4hello€"; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - let rx = mutex.lock().await; + let rx = stream::iter([ + smallvec![Packet::Message("hello€".into())], + smallvec![Packet::BinaryV3(Bytes::from_static(&[1, 2, 3, 4]))], + smallvec![Packet::Message("hello€".into())], + ]); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::BinaryV3(Bytes::from_static( - &[1, 2, 3, 4] - ))]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); + let rx = std::pin::pin!(Peekable::new(rx)); let Payload { data, has_binary, .. - } = v3_string_encoder(rx, MAX_PAYLOAD).await.unwrap(); + } = v3_string_encoder(rx, MAX_PAYLOAD).await; assert_eq!(data, PAYLOAD.as_bytes()); assert!(!has_binary); } @@ -390,32 +356,25 @@ mod tests { async fn max_payload_v3_b64() { const MAX_PAYLOAD: u64 = 10; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::BinaryV3(Bytes::from_static( - &[1, 2, 3, 4] - ))]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); + let rx = stream::iter(vec![ + smallvec::smallvec![Packet::Message("hello€".into())], + smallvec::smallvec![Packet::BinaryV3(Bytes::from_static(&[1, 2, 3, 4]))], + smallvec::smallvec![Packet::Message("hello€".into())], + smallvec::smallvec![Packet::Message("hello€".into())], + ]); + let mut rx = std::pin::pin!(Peekable::new(rx)); + { - let rx = mutex.lock().await; - let Payload { data, .. } = v3_string_encoder(rx, MAX_PAYLOAD).await.unwrap(); + let Payload { data, .. } = v3_string_encoder(rx.as_mut(), MAX_PAYLOAD).await; assert_eq!(data, "7:4hello€".as_bytes()); } { - let rx = mutex.lock().await; - let Payload { data, .. } = v3_string_encoder(rx, MAX_PAYLOAD + 10).await.unwrap(); + let Payload { data, .. } = v3_string_encoder(rx.as_mut(), MAX_PAYLOAD + 10).await; assert_eq!(data, "10:b4AQIDBA==".as_bytes()); } { // Next call drains one of the remaining Message packets. - let rx = mutex.lock().await; - let Payload { data, .. } = v3_string_encoder(rx, MAX_PAYLOAD + 10).await.unwrap(); + let Payload { data, .. } = v3_string_encoder(rx.as_mut(), MAX_PAYLOAD + 10).await; assert_eq!(data, "7:4hello€".as_bytes()); } } @@ -426,19 +385,16 @@ mod tests { const PAYLOAD: [u8; 20] = [ 0, 9, 255, 52, 104, 101, 108, 108, 111, 226, 130, 172, 1, 5, 255, 4, 1, 2, 3, 4, ]; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - let rx = mutex.lock().await; - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::BinaryV3(Bytes::from_static( - &[1, 2, 3, 4] - ))]) - .unwrap(); + let rx = stream::iter([ + smallvec![Packet::Message("hello€".into())], + smallvec![Packet::BinaryV3(Bytes::from_static(&[1, 2, 3, 4]))], + ]); + let rx = std::pin::pin!(Peekable::new(rx)); + let Payload { data, has_binary, .. - } = v3_binary_encoder(rx, MAX_PAYLOAD).await.unwrap(); + } = v3_binary_encoder(rx, MAX_PAYLOAD).await; assert_eq!(*data, PAYLOAD); assert!(has_binary); } @@ -455,8 +411,8 @@ mod tests { 0, 9, 255, 52, 104, 101, 108, 108, 111, 226, 130, 172, 1, 5, 255, 4, 1, 2, 3, 4, ]; let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - let rx = mutex.lock().await; + let rx = tokio_stream::wrappers::ReceiverStream::new(rx); + let rx = std::pin::pin!(Peekable::new(rx)); tokio::spawn(async move { tokio::time::sleep(std::time::Duration::from_millis(50)).await; tx.try_send(smallvec::smallvec![ @@ -467,7 +423,7 @@ mod tests { }); let Payload { data, has_binary, .. - } = v3_binary_encoder(rx, MAX_PAYLOAD).await.unwrap(); + } = v3_binary_encoder(rx, MAX_PAYLOAD).await; assert_eq!(*data, PAYLOAD); assert!(has_binary); } @@ -481,26 +437,21 @@ mod tests { 0, 1, 1, 255, 52, 104, 101, 108, 108, 111, 111, 111, 226, 130, 172, 1, 5, 255, 4, 1, 2, 3, 4, ]; - let (tx, rx) = tokio::sync::mpsc::channel::(10); - let mutex = Mutex::new(PeekableReceiver::new(rx)); - tx.try_send(smallvec::smallvec![Packet::Message("hellooo€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::BinaryV3(Bytes::from_static( - &[1, 2, 3, 4] - ))]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); - tx.try_send(smallvec::smallvec![Packet::Message("hello€".into())]) - .unwrap(); + + let rx = stream::iter([ + smallvec![Packet::Message("hellooo€".into())], + smallvec![Packet::BinaryV3(Bytes::from_static(&[1, 2, 3, 4]))], + smallvec![Packet::Message("hello€".into())], + smallvec![Packet::Message("hello€".into())], + ]); + let mut rx = std::pin::pin!(Peekable::new(rx)); + { - let rx = mutex.lock().await; - let Payload { data, .. } = v3_binary_encoder(rx, MAX_PAYLOAD).await.unwrap(); + let Payload { data, .. } = v3_binary_encoder(rx.as_mut(), MAX_PAYLOAD).await; assert_eq!(*data, PAYLOAD); } { - let rx = mutex.lock().await; - let Payload { data, .. } = v3_binary_encoder(rx, MAX_PAYLOAD).await.unwrap(); + let Payload { data, .. } = v3_binary_encoder(rx.as_mut(), MAX_PAYLOAD).await; assert_eq!(data, "7:4hello€7:4hello€".as_bytes()); } } diff --git a/crates/engineioxide-core/src/payload/mod.rs b/crates/engineioxide-core/src/payload/mod.rs new file mode 100644 index 00000000..23f153b7 --- /dev/null +++ b/crates/engineioxide-core/src/payload/mod.rs @@ -0,0 +1,252 @@ +//! Payload encoder and decoder for polling transport. + +use crate::{Packet, PacketParseError, ProtocolVersion, packet::PacketBuf}; + +use bytes::Bytes; +use futures_util::Stream; +use http::HeaderValue; + +mod buf; +mod decoder; +mod encoder; + +const PACKET_SEPARATOR_V4: u8 = b'\x1e'; +#[cfg(feature = "v3")] +const STRING_PACKET_SEPARATOR_V3: u8 = b':'; +#[cfg(feature = "v3")] +const BINARY_PACKET_SEPARATOR_V3: u8 = 0xff; +#[cfg(feature = "v3")] +const STRING_PACKET_IDENTIFIER_V3: u8 = 0x00; +#[cfg(feature = "v3")] +const BINARY_PACKET_IDENTIFIER_V3: u8 = 0x01; +#[cfg(feature = "v3")] +const BINARY_CONTENT_TYPE: HeaderValue = HeaderValue::from_static("application/octet-stream"); + +/// Decode a payload into a stream of packets. +pub fn decoder( + body: impl http_body::Body + Unpin, + #[allow(unused)] content_type: Option<&HeaderValue>, + #[allow(unused)] protocol: ProtocolVersion, + max_payload: u64, +) -> impl Stream> { + #[cfg(feature = "tracing")] + tracing::debug!(?content_type, %protocol, "decoding payload"); + + #[cfg(feature = "v3")] + { + use futures_util::future::Either; + + let is_binary = content_type == Some(&BINARY_CONTENT_TYPE); + match protocol { + ProtocolVersion::V4 => Either::Left(decoder::v4_decoder(body, max_payload)), + ProtocolVersion::V3 if is_binary => { + Either::Right(Either::Left(decoder::v3_binary_decoder(body, max_payload))) + } + ProtocolVersion::V3 => { + Either::Right(Either::Right(decoder::v3_string_decoder(body, max_payload))) + } + } + } + + #[cfg(not(feature = "v3"))] + { + decoder::v4_decoder(body, max_payload) + } +} + +/// A payload to transmit to the client through http polling +pub struct Payload { + /// The data of the payload. + pub data: Bytes, + /// Whether the payload contains binary data. + pub has_binary: bool, + /// A peeked packet buffer that could not be sent due to the max payload size. + pub peeked: Option, +} +impl Payload { + /// Creates a new payload with the given data and binary flag. + pub fn new(data: Bytes, has_binary: bool) -> Self { + Self { + data, + has_binary, + peeked: None, + } + } +} + +/// Encodes a payload into a byte stream. +pub async fn encoder( + rx: impl Stream, + #[allow(unused)] protocol: ProtocolVersion, + #[allow(unused)] supports_binary: bool, + max_payload: u64, +) -> Payload { + let mut rx = std::pin::pin!(peekable::Peekable::new(rx)); + + #[cfg(feature = "v3")] + let mut payload = match protocol { + ProtocolVersion::V4 => encoder::v4_encoder(rx.as_mut(), max_payload).await, + ProtocolVersion::V3 if supports_binary => { + encoder::v3_binary_encoder(rx.as_mut(), max_payload).await + } + ProtocolVersion::V3 => encoder::v3_string_encoder(rx.as_mut(), max_payload).await, + }; + + #[cfg(not(feature = "v3"))] + let mut payload = encoder::v4_encoder(rx.as_mut(), max_payload).await; + + // Recover any packet that was peeked but rejected (max payload exceeded) so + // the caller can hand it back to the next encoding pass instead of losing it. + payload.peeked = rx.take_peeked(); + + payload +} + +/// Encodes a single packet into a byte array. +pub fn packet_encoder( + packet: Packet, + #[allow(unused)] protocol: ProtocolVersion, + #[allow(unused)] supports_binary: bool, +) -> Bytes { + #[cfg(feature = "v3")] + match protocol { + ProtocolVersion::V4 => packet.into(), + ProtocolVersion::V3 if supports_binary && packet.is_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() +} + +mod peekable { + //! [`Peekable`] implementation taken from [`futures_util::stream::Peekable`] but + //! with an [`Peekable::take_peeked`] method to get back a potentially peeked item. + use std::{ + pin::Pin, + task::{Context, Poll}, + }; + + use futures_util::{Stream, StreamExt as _, stream::Fuse}; + use pin_project_lite::pin_project; + + pin_project! { + /// A `Stream` that implements a `peek` method. + /// + /// The `peek` method can be used to retrieve a reference + /// to the next `Stream::Item` if available. A subsequent + /// call to `poll` will return the owned item. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Peekable { + #[pin] + stream: Fuse, + peeked: Option, + } + } + + impl Peekable { + pub(super) fn new(stream: St) -> Self { + Self { + stream: stream.fuse(), + peeked: None, + } + } + + /// Produces a future which retrieves a reference to the next item + /// in the stream, or `None` if the underlying stream terminates. + pub fn peek(self: Pin<&mut Self>) -> Peek<'_, St> { + Peek { inner: Some(self) } + } + + /// Takes the peeked item out of the buffer, if any. + /// + /// Only the `peeked` field is moved out, so this is safe to call + /// through a pinned reference even when the underlying stream is `!Unpin`. + pub fn take_peeked(self: Pin<&mut Self>) -> Option { + self.project().peeked.take() + } + + /// Peek retrieves a reference to the next item in the stream. + /// + /// This method polls the underlying stream and return either a reference + /// to the next item if the stream is ready or passes through any errors. + pub fn poll_peek(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + + Poll::Ready(loop { + if this.peeked.is_some() { + break this.peeked.as_ref(); + } else if let Some(item) = futures_util::ready!(this.stream.as_mut().poll_next(cx)) + { + *this.peeked = Some(item); + } else { + break None; + } + }) + } + } + + impl Stream for Peekable { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + if let Some(item) = this.peeked.take() { + return Poll::Ready(Some(item)); + } + this.stream.poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option) { + let peek_len = usize::from(self.peeked.is_some()); + let (lower, upper) = self.stream.size_hint(); + let lower = lower.saturating_add(peek_len); + let upper = match upper { + Some(x) => x.checked_add(peek_len), + None => None, + }; + (lower, upper) + } + } + + pin_project! { + /// Future for the [`Peekable::peek`](self::Peekable::peek) method. + #[must_use = "futures do nothing unless polled"] + pub struct Peek<'a, St: Stream> { + inner: Option>>, + } + } + + impl<'a, St> Future for Peek<'a, St> + where + St: Stream, + { + type Output = Option<&'a St::Item>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let inner = self.project().inner; + if let Some(peekable) = inner { + futures_util::ready!(peekable.as_mut().poll_peek(cx)); + + inner.take().unwrap().poll_peek(cx) + } else { + panic!("Peek polled after completion") + } + } + } +} diff --git a/crates/engineioxide-core/src/protocol.rs b/crates/engineioxide-core/src/protocol.rs new file mode 100644 index 00000000..acc9714f --- /dev/null +++ b/crates/engineioxide-core/src/protocol.rs @@ -0,0 +1,125 @@ +use std::{fmt, str::FromStr}; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// The type of `transport` used to connect to the client/server. +#[derive(Debug, Copy, Clone, PartialEq)] +pub enum TransportType { + /// Polling transport + Polling = 0x01, + /// Websocket transport + Websocket = 0x02, +} + +impl From for TransportType { + fn from(t: u8) -> Self { + match t { + 0x01 => TransportType::Polling, + 0x02 => TransportType::Websocket, + _ => panic!("unknown transport type"), + } + } +} + +impl FromStr for TransportType { + type Err = UnknownTransportError; + + fn from_str(s: &str) -> Result { + match s { + "websocket" => Ok(TransportType::Websocket), + "polling" => Ok(TransportType::Polling), + _ => Err(UnknownTransportError), + } + } +} +impl From for &'static str { + fn from(t: TransportType) -> Self { + match t { + TransportType::Polling => "polling", + TransportType::Websocket => "websocket", + } + } +} +impl From for String { + fn from(t: TransportType) -> Self { + match t { + TransportType::Polling => "polling".into(), + TransportType::Websocket => "websocket".into(), + } + } +} + +impl Serialize for TransportType { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str((*self).into()) + } +} + +impl<'de> Deserialize<'de> for TransportType { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + Self::from_str(&s).map_err(serde::de::Error::custom) + } +} + +/// Cannot determine the transport type to connect to the client/server. +#[derive(Debug, Copy, Clone)] +pub struct UnknownTransportError; +impl std::fmt::Display for UnknownTransportError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "unknown transport type") + } +} +impl std::error::Error for UnknownTransportError {} + +/// == ProtocolVersion == + +#[derive(Debug)] +pub struct UnknownProtocolVersionError; +impl std::fmt::Display for UnknownProtocolVersionError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "unknown protocol version") + } +} +impl std::error::Error for UnknownProtocolVersionError {} + +/// The engine.io protocol version +#[derive(Debug, Copy, Clone, PartialEq)] +pub enum ProtocolVersion { + /// The protocol version 3 + V3 = 3, + /// The protocol version 4 + V4 = 4, +} + +impl fmt::Display for ProtocolVersion { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&(*self as u8), f) + } +} +impl FromStr for ProtocolVersion { + type Err = UnknownProtocolVersionError; + + #[cfg(feature = "v3")] + fn from_str(s: &str) -> Result { + match s { + "3" => Ok(ProtocolVersion::V3), + "4" => Ok(ProtocolVersion::V4), + _ => Err(UnknownProtocolVersionError), + } + } + + #[cfg(not(feature = "v3"))] + fn from_str(s: &str) -> Result { + match s { + "4" => Ok(ProtocolVersion::V4), + _ => Err(UnknownProtocolVersionError), + } + } +} diff --git a/crates/engineioxide/Cargo.toml b/crates/engineioxide/Cargo.toml index b604481d..030f98be 100644 --- a/crates/engineioxide/Cargo.toml +++ b/crates/engineioxide/Cargo.toml @@ -33,21 +33,15 @@ tokio-util.workspace = true tower-service.workspace = true tower-layer.workspace = true hyper.workspace = true -tokio-tungstenite.workspace = true +tokio-tungstenite = { workspace = true, features = ["handshake"] } http-body-util.workspace = true pin-project-lite.workspace = true smallvec.workspace = true hyper-util = { workspace = true, features = ["tokio"] } -base64 = "0.22" - # Tracing tracing = { workspace = true, optional = true } -# Engine.io V3 payload -itoa = { workspace = true, optional = true } -memchr = { version = "2.7", optional = true } - [dev-dependencies] tokio = { workspace = true, features = ["macros", "parking_lot"] } tracing-subscriber.workspace = true @@ -55,9 +49,11 @@ hyper = { workspace = true, features = ["server", "http1"] } criterion.workspace = true axum.workspace = true tokio-stream.workspace = true +tokio-util.workspace = true + [features] -v3 = ["memchr", "itoa"] -tracing = ["dep:tracing"] +v3 = ["engineioxide-core/v3"] +tracing = ["dep:tracing", "engineioxide-core/tracing"] __test_harness = [] [[bench]] diff --git a/crates/engineioxide/benches/packet_encode.rs b/crates/engineioxide/benches/packet_encode.rs index 0cfc8842..f1d8f483 100644 --- a/crates/engineioxide/benches/packet_encode.rs +++ b/crates/engineioxide/benches/packet_encode.rs @@ -1,15 +1,11 @@ use bytes::Bytes; use criterion::{BatchSize, Criterion, black_box, criterion_group, criterion_main}; -use engineioxide::{OpenPacket, Packet, TransportType, config::EngineIoConfig, socket::Sid}; +use engineioxide::{OpenPacket, Packet}; fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("engineio_packet/encode"); group.bench_function("Encode packet open", |b| { - let packet = Packet::Open(OpenPacket::new( - black_box(TransportType::Polling), - black_box(Sid::ZERO), - &EngineIoConfig::default(), - )); + let packet = Packet::Open(OpenPacket::default()); b.iter_batched( || packet.clone(), TryInto::::try_into, diff --git a/crates/engineioxide/src/config.rs b/crates/engineioxide/src/config.rs index aa74a583..2cd9f603 100644 --- a/crates/engineioxide/src/config.rs +++ b/crates/engineioxide/src/config.rs @@ -32,7 +32,7 @@ use std::{borrow::Cow, time::Duration}; -use crate::service::TransportType; +use engineioxide_core::TransportType; /// Configuration for the engine.io engine & transports #[derive(Debug, Clone)] diff --git a/crates/engineioxide/src/engine.rs b/crates/engineioxide/src/engine.rs index 99353b7c..c0fdabc1 100644 --- a/crates/engineioxide/src/engine.rs +++ b/crates/engineioxide/src/engine.rs @@ -3,13 +3,12 @@ use std::{ sync::{Arc, RwLock}, }; -use engineioxide_core::Sid; +use engineioxide_core::{ProtocolVersion, Sid, TransportType}; use http::request::Parts; use crate::{ config::EngineIoConfig, handler::EngineIoHandler, - service::{ProtocolVersion, TransportType}, socket::{DisconnectReason, Socket}, }; @@ -45,7 +44,7 @@ impl EngineIo { protocol: ProtocolVersion, transport: TransportType, req: Parts, - #[cfg(feature = "v3")] supports_binary: bool, + supports_binary: bool, ) -> Arc> { let engine = self.clone(); let close_fn = Box::new(move |sid, reason| engine.close_session(sid, reason)); @@ -56,7 +55,6 @@ impl EngineIo { &self.config, req, close_fn, - #[cfg(feature = "v3")] supports_binary, ); let socket = Arc::new(socket); @@ -90,11 +88,23 @@ impl EngineIo { return; }; - // Try to close the internal channel if it is available - // E.g. with polling transport the channel is not always locked so it is necessary to close it here - socket.internal_rx.try_lock().map(|mut rx| rx.close()).ok(); + // Notify every transport task listening on this token to stop. Cancellation-aware + // transports (the ws `forward_to_socket` task and pending polling requests) will release + // the `internal_rx` lock as soon as they observe the cancellation. socket.cancellation_token.cancel(); - self.handler.on_disconnect(socket, reason); + + let handler = self.handler.clone(); + tokio::spawn(async move { + // Wait for the cancelled transport task to yield the `internal_rx` lock, then close the + // channel so that `Socket::closed` resolves and no more packets can be buffered. + // When no transport holds the lock (e.g. polling between requests) this resolves + // immediately. + socket.internal_rx.lock().await.rx.close(); + + #[cfg(feature = "tracing")] + tracing::debug!(sid = %socket.id, "session closed, notifying handler"); + handler.on_disconnect(socket, reason); + }); #[cfg(feature = "tracing")] tracing::debug!( @@ -149,7 +159,6 @@ mod tests { ProtocolVersion::V4, TransportType::Polling, Request::<()>::default().into_parts().0, - #[cfg(feature = "v3")] true, ); assert_eq!(engine.sockets.read().unwrap().len(), 1); @@ -164,7 +173,6 @@ mod tests { ProtocolVersion::V4, TransportType::Polling, Request::<()>::default().into_parts().0, - #[cfg(feature = "v3")] true, ); assert_eq!(engine.sockets.read().unwrap().len(), 1); @@ -179,7 +187,6 @@ mod tests { ProtocolVersion::V4, TransportType::Polling, Request::<()>::default().into_parts().0, - #[cfg(feature = "v3")] true, ); assert_eq!(engine.sockets.read().unwrap().len(), 1); diff --git a/crates/engineioxide/src/errors.rs b/crates/engineioxide/src/errors.rs index 29b335a0..393be172 100644 --- a/crates/engineioxide/src/errors.rs +++ b/crates/engineioxide/src/errors.rs @@ -1,52 +1,35 @@ use http::{Response, StatusCode}; -use tokio::sync::mpsc; use tokio_tungstenite::tungstenite; use crate::body::ResponseBody; -use crate::packet::Packet; -use engineioxide_core::Sid; +use engineioxide_core::{Packet, Sid}; + +pub use engineioxide_core::PacketParseError; #[derive(thiserror::Error, Debug)] pub enum Error { - #[error("error decoding binary packet from polling request: {0:?}")] - Base64(#[from] base64::DecodeError), - #[error("error decoding packet: {0:?}")] - StrUtf8(#[from] std::str::Utf8Error), - #[error("io error: {0:?}")] - Io(#[from] std::io::Error), - #[error("bad packet received")] + #[error("error decoding packet from request: {0}")] + PacketParse(#[from] PacketParseError), + #[error("invalid packet received: {0:?}")] BadPacket(Packet), - #[error("ws transport error: {0:?}")] + #[error("ws transport error: {0}")] WsTransport(#[from] Box), - #[error("http error: {0:?}")] - Http(#[from] http::Error), - #[error("internal channel error: {0:?}")] - SendChannel(#[from] mpsc::error::TrySendError), - #[error("internal channel error: {0:?}")] - RecvChannel(#[from] mpsc::error::TryRecvError), #[error("heartbeat timeout")] HeartbeatTimeout, #[error("upgrade error")] Upgrade, #[error("multiple ws upgrade requests")] MultipleWebsocketRequests, - #[error("aborted connection")] - Aborted, - #[error("http error response: {0:?}")] - HttpErrorResponse(StatusCode), + #[error("multiple http polling error")] + MultipleHttpPolling, + #[error("invalid websocket Sec-WebSocket-Key http header")] + InvalidWebSocketKey, #[error("unknown session id")] UnknownSessionID(Sid), #[error("transport mismatch")] TransportMismatch, - #[error("payload too large")] - PayloadTooLarge, - - #[error("Invalid packet length")] - InvalidPacketLength, - #[error("Invalid packet type")] - InvalidPacketType(Option), } /// Convert an error into an http response @@ -62,18 +45,15 @@ impl From for Response> { .unwrap() }; match err { - Error::HttpErrorResponse(code) => Response::builder() - .status(code) + Error::PacketParse(PacketParseError::PayloadTooLarge { .. }) => Response::builder() + .status(413) .body(ResponseBody::empty_response()) .unwrap(), - Error::BadPacket(_) | Error::InvalidPacketLength | Error::InvalidPacketType(_) => { - Response::builder() - .status(400) - .body(ResponseBody::empty_response()) - .unwrap() - } - Error::PayloadTooLarge => Response::builder() - .status(413) + Error::BadPacket(_) + | Error::PacketParse(_) + | Error::MultipleHttpPolling + | Error::InvalidWebSocketKey => Response::builder() + .status(400) .body(ResponseBody::empty_response()) .unwrap(), diff --git a/crates/engineioxide/src/lib.rs b/crates/engineioxide/src/lib.rs index b8f34b3b..94cb84d7 100644 --- a/crates/engineioxide/src/lib.rs +++ b/crates/engineioxide/src/lib.rs @@ -2,13 +2,12 @@ #![cfg_attr(docsrs, feature(doc_cfg))] #![doc = include_str!("../README.md")] -pub use engineioxide_core::Str; -pub use service::{ProtocolVersion, TransportType}; +pub use engineioxide_core::{ProtocolVersion, Str, TransportType}; pub use socket::{DisconnectReason, Socket}; #[doc(hidden)] #[cfg(feature = "__test_harness")] -pub use packet::*; +pub use engineioxide_core::{OpenPacket, Packet, PacketParseError}; pub mod config; pub mod handler; @@ -25,6 +24,4 @@ pub mod sid { mod body; mod engine; mod errors; -mod packet; -mod peekable; mod transport; diff --git a/crates/engineioxide/src/peekable.rs b/crates/engineioxide/src/peekable.rs deleted file mode 100644 index bc8ddb6b..00000000 --- a/crates/engineioxide/src/peekable.rs +++ /dev/null @@ -1,79 +0,0 @@ -use tokio::sync::mpsc::{Receiver, error::TryRecvError}; - -/// Peekable receiver for polling transport -/// It is a thin wrapper around a [`Receiver`](tokio::sync::mpsc::Receiver) that allows to peek the next packet without consuming it -/// -/// Its main goal is to be able to peek the next packet without consuming it to calculate the -/// packet length when using polling transport to check if it fits according to the max_payload setting -#[derive(Debug)] -pub struct PeekableReceiver { - rx: Receiver, - next: Option, -} -impl PeekableReceiver { - pub fn new(rx: Receiver) -> Self { - Self { rx, next: None } - } - pub fn peek(&mut self) -> Option<&T> { - if self.next.is_none() { - self.next = self.rx.try_recv().ok(); - } - self.next.as_ref() - } - pub async fn recv(&mut self) -> Option { - if self.next.is_none() { - self.rx.recv().await - } else { - self.next.take() - } - } - pub fn try_recv(&mut self) -> Result { - if self.next.is_none() { - self.rx.try_recv() - } else { - Ok(self.next.take().unwrap()) - } - } - - pub fn close(&mut self) { - self.rx.close() - } -} - -#[cfg(test)] -mod tests { - use tokio::sync::Mutex; - - #[tokio::test] - async fn peek() { - use super::PeekableReceiver; - use crate::packet::Packet; - use tokio::sync::mpsc::channel; - - let (tx, rx) = channel(1); - let rx = Mutex::new(PeekableReceiver::new(rx)); - let mut rx = rx.lock().await; - - assert!(rx.peek().is_none()); - - tx.send(Packet::Ping).await.unwrap(); - assert_eq!(rx.peek(), Some(&Packet::Ping)); - assert_eq!(rx.recv().await, Some(Packet::Ping)); - assert!(rx.peek().is_none()); - - tx.send(Packet::Pong).await.unwrap(); - assert_eq!(rx.peek(), Some(&Packet::Pong)); - assert_eq!(rx.recv().await, Some(Packet::Pong)); - assert!(rx.peek().is_none()); - - tx.send(Packet::Close).await.unwrap(); - assert_eq!(rx.peek(), Some(&Packet::Close)); - assert_eq!(rx.recv().await, Some(Packet::Close)); - assert!(rx.peek().is_none()); - - tx.send(Packet::Close).await.unwrap(); - assert_eq!(rx.peek(), Some(&Packet::Close)); - assert_eq!(rx.recv().await, Some(Packet::Close)); - assert!(rx.peek().is_none()); - } -} diff --git a/crates/engineioxide/src/service/mod.rs b/crates/engineioxide/src/service/mod.rs index 4bce8fe6..6b9d1c13 100644 --- a/crates/engineioxide/src/service/mod.rs +++ b/crates/engineioxide/src/service/mod.rs @@ -33,6 +33,8 @@ use std::{ }; use bytes::Bytes; +#[cfg(feature = "__test_harness")] +use engineioxide_core::ProtocolVersion; use futures_util::future::{self, Ready}; use http::{Request, Response}; use http_body::Body; @@ -47,7 +49,6 @@ use crate::{ mod futures; mod parser; -pub use self::parser::{ProtocolVersion, TransportType}; use self::{futures::ResponseFuture, parser::dispatch_req}; /// A `Service` that handles engine.io requests as a middleware. diff --git a/crates/engineioxide/src/service/parser.rs b/crates/engineioxide/src/service/parser.rs index f1235063..2cb5623a 100644 --- a/crates/engineioxide/src/service/parser.rs +++ b/crates/engineioxide/src/service/parser.rs @@ -1,11 +1,10 @@ //! A Parser module to parse any `EngineIo` query -use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use std::{future::Future, str::FromStr, sync::Arc}; +use std::{future::Future, sync::Arc}; use http::{Method, Request, Response}; -use engineioxide_core::Sid; +use engineioxide_core::{ProtocolVersion, Sid, TransportType}; use crate::{ body::ResponseBody, @@ -35,15 +34,8 @@ where sid: None, transport: TransportType::Polling, method: Method::GET, - #[cfg(feature = "v3")] b64, - }) => ResponseFuture::ready(polling::open_req( - engine, - protocol, - req, - #[cfg(feature = "v3")] - !b64, - )), + }) => ResponseFuture::ready(polling::open_req(engine, protocol, req, !b64)), Ok(RequestInfo { protocol, sid: Some(sid), @@ -118,102 +110,6 @@ impl From for Response> { } } -/// The engine.io protocol version -#[derive(Debug, Copy, Clone, PartialEq)] -pub enum ProtocolVersion { - /// The protocol version 3 - V3 = 3, - /// The protocol version 4 - V4 = 4, -} - -impl FromStr for ProtocolVersion { - type Err = ParseError; - - #[cfg(feature = "v3")] - fn from_str(s: &str) -> Result { - match s { - "3" => Ok(ProtocolVersion::V3), - "4" => Ok(ProtocolVersion::V4), - _ => Err(ParseError::UnsupportedProtocolVersion), - } - } - - #[cfg(not(feature = "v3"))] - fn from_str(s: &str) -> Result { - match s { - "4" => Ok(ProtocolVersion::V4), - _ => Err(ParseError::UnsupportedProtocolVersion), - } - } -} - -/// The type of `transport` used by the client. -#[derive(Debug, Copy, Clone, PartialEq, PartialOrd)] -pub enum TransportType { - /// Polling transport - Polling = 0x01, - /// Websocket transport - Websocket = 0x02, -} - -impl Serialize for TransportType { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_str((*self).into()) - } -} - -impl<'de> Deserialize<'de> for TransportType { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let s = String::deserialize(deserializer)?; - Self::from_str(&s).map_err(serde::de::Error::custom) - } -} - -impl From for TransportType { - fn from(t: u8) -> Self { - match t { - 0x01 => TransportType::Polling, - 0x02 => TransportType::Websocket, - _ => panic!("unknown transport type"), - } - } -} - -impl FromStr for TransportType { - type Err = ParseError; - - fn from_str(s: &str) -> Result { - match s { - "websocket" => Ok(TransportType::Websocket), - "polling" => Ok(TransportType::Polling), - _ => Err(ParseError::UnknownTransport), - } - } -} -impl From for &'static str { - fn from(t: TransportType) -> Self { - match t { - TransportType::Polling => "polling", - TransportType::Websocket => "websocket", - } - } -} -impl From for String { - fn from(t: TransportType) -> Self { - match t { - TransportType::Polling => "polling".into(), - TransportType::Websocket => "websocket".into(), - } - } -} - /// The request information extracted from the request URI. #[derive(Debug)] pub struct RequestInfo { @@ -226,7 +122,6 @@ pub struct RequestInfo { /// The request method. pub method: Method, /// If the client asked for base64 encoding only. - #[cfg(feature = "v3")] pub b64: bool, } @@ -240,8 +135,8 @@ impl RequestInfo { .split('&') .find(|s| s.starts_with("EIO=")) .and_then(|s| s.split('=').nth(1)) - .ok_or(UnsupportedProtocolVersion) - .and_then(|t| t.parse())?; + .and_then(|t| t.parse().ok()) + .ok_or(UnsupportedProtocolVersion)?; let sid = query .split('&') @@ -253,8 +148,8 @@ impl RequestInfo { .split('&') .find(|s| s.starts_with("transport=")) .and_then(|s| s.split('=').nth(1)) - .ok_or(UnknownTransport) - .and_then(|t| t.parse())?; + .and_then(|t| t.parse().ok()) + .ok_or(UnknownTransport)?; if !config.allowed_transport(transport) { return Err(TransportMismatch); @@ -267,6 +162,9 @@ impl RequestInfo { .map(|v| v == "1" || v == "true") .unwrap_or_default(); + #[cfg(not(feature = "v3"))] + let b64: bool = false; + let method = req.method().clone(); if !matches!(method, Method::GET) && sid.is_none() { Err(BadHandshakeMethod) @@ -276,7 +174,6 @@ impl RequestInfo { sid, transport, method, - #[cfg(feature = "v3")] b64, }) } diff --git a/crates/engineioxide/src/socket.rs b/crates/engineioxide/src/socket.rs index 222e3d8f..99c42edb 100644 --- a/crates/engineioxide/src/socket.rs +++ b/crates/engineioxide/src/socket.rs @@ -56,6 +56,8 @@ //! ``` use std::{ collections::VecDeque, + fmt::Debug, + ops::{Deref, DerefMut}, sync::{ Arc, atomic::{AtomicBool, AtomicU8, Ordering}, @@ -63,15 +65,9 @@ use std::{ time::Duration, }; -use crate::{ - config::EngineIoConfig, - errors::Error, - packet::Packet, - peekable::PeekableReceiver, - service::{ProtocolVersion, TransportType}, -}; +use crate::{config::EngineIoConfig, errors::Error}; use bytes::Bytes; -use engineioxide_core::Str; +use engineioxide_core::{Packet, PacketBuf, ProtocolVersion, Str, TransportType}; use futures_util::FutureExt; use http::request::Parts; use smallvec::{SmallVec, smallvec}; @@ -108,9 +104,8 @@ impl From<&Error> for Option { fn from(err: &Error) -> Self { use Error::*; match err { - WsTransport(_) | Io(_) => Some(DisconnectReason::TransportError), - BadPacket(_) | Base64(_) | StrUtf8(_) | PayloadTooLarge | InvalidPacketLength - | InvalidPacketType(_) => Some(DisconnectReason::PacketParsingError), + WsTransport(_) => Some(DisconnectReason::TransportError), + BadPacket(_) | PacketParse(_) => Some(DisconnectReason::PacketParsingError), HeartbeatTimeout => Some(DisconnectReason::HeartbeatTimeout), _ => None, } @@ -159,12 +154,28 @@ impl Permit<'_> { } } -/// Buffered packets to send to the client. -/// It is used to ensure atomicity when sending multiple packets to the client. -/// -/// The [`PacketBuf`] stack size will impact the dynamically allocated buffer -/// of the internal mpsc channel. -pub(crate) type PacketBuf = SmallVec<[Packet; 2]>; +#[derive(Debug)] +pub(crate) struct InternalRx { + pub rx: Receiver, + pub peeked: Option, +} +impl InternalRx { + fn new(rx: Receiver) -> Self { + Self { rx, peeked: None } + } +} +impl Deref for InternalRx { + type Target = Receiver; + + fn deref(&self) -> &Self::Target { + &self.rx + } +} +impl DerefMut for InternalRx { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.rx + } +} /// A [`Socket`] represents a client connection to the server. /// It is agnostic to the [`TransportType`]. @@ -200,17 +211,13 @@ where /// * In case of polling transport it will be locked and released for each request /// * In case of websocket transport it will always be locked until the connection is closed /// - /// It will be closed when a [`Close`](Packet::Close) packet is received: - /// * From the [encoder](crate::service::encoder) if the transport is polling - /// * From the fn [`on_ws_req_init`](crate::engine::EngineIo) if the transport is websocket - /// * Automatically via the [`close_session fn`](crate::engine::EngineIo::close_session) as a fallback. - /// Because with polling transport, if the client is not currently polling then the encoder will never be able to close the channel + /// Closing behavior is driven by the [`Socket::cancellation_token`]. /// /// The channel is made of a [`SmallVec`] of [`Packet`]s so that adjacent packets can be sent atomically. - pub(crate) internal_rx: Mutex>, + pub(crate) internal_rx: Mutex, /// Channel to send [PacketBuf] to the internal connection - internal_tx: mpsc::Sender, + pub(crate) internal_tx: mpsc::Sender, /// Internal channel to receive Pong [`Packets`](Packet) (v4 protocol) or Ping (v3 protocol) in the heartbeat job /// which is running in a separate task @@ -231,7 +238,6 @@ where pub req_parts: Parts, /// If the client supports binary packets (via polling XHR2) - #[cfg(feature = "v3")] pub(crate) supports_binary: bool, } @@ -245,7 +251,7 @@ where config: &EngineIoConfig, req_parts: Parts, close_fn: Box, - #[cfg(feature = "v3")] supports_binary: bool, + supports_binary: bool, ) -> Self { let (internal_tx, internal_rx) = mpsc::channel(config.max_buffer_size); let (heartbeat_tx, heartbeat_rx) = mpsc::channel(1); @@ -256,7 +262,7 @@ where transport: AtomicU8::new(transport as u8), upgrading: AtomicBool::new(false), - internal_rx: Mutex::new(PeekableReceiver::new(internal_rx)), + internal_rx: Mutex::new(InternalRx::new(internal_rx)), internal_tx, heartbeat_rx: Mutex::new(heartbeat_rx), @@ -268,7 +274,6 @@ where data: D::default(), req_parts, - #[cfg(feature = "v3")] supports_binary, } } @@ -555,7 +560,7 @@ where transport: AtomicU8::new(TransportType::Websocket as u8), upgrading: AtomicBool::new(false), - internal_rx: Mutex::new(PeekableReceiver::new(internal_rx)), + internal_rx: Mutex::new(InternalRx::new(internal_rx)), internal_tx, heartbeat_rx: Mutex::new(heartbeat_rx), @@ -566,7 +571,6 @@ where data: D::default(), req_parts: http::Request::<()>::default().into_parts().0, - #[cfg(feature = "v3")] supports_binary: true, }; let sock = Arc::new(sock); diff --git a/crates/engineioxide/src/transport/mod.rs b/crates/engineioxide/src/transport/mod.rs index 94d843a5..d0ad9a45 100644 --- a/crates/engineioxide/src/transport/mod.rs +++ b/crates/engineioxide/src/transport/mod.rs @@ -1,4 +1,23 @@ //! All transports modules available in engineioxide +use engineioxide_core::{OpenPacket, Sid, TransportType}; + +use crate::config::EngineIoConfig; + pub mod polling; pub mod ws; + +fn make_open_packet(transport: TransportType, id: Sid, config: &EngineIoConfig) -> OpenPacket { + let upgrades = if transport == TransportType::Polling { + smallvec::smallvec![TransportType::Websocket] + } else { + smallvec::smallvec![] + }; + OpenPacket { + sid: id, + upgrades, + ping_timeout: config.ping_timeout, + ping_interval: config.ping_interval, + max_payload: config.max_payload, + } +} diff --git a/crates/engineioxide/src/transport/polling/mod.rs b/crates/engineioxide/src/transport/polling.rs similarity index 57% rename from crates/engineioxide/src/transport/polling/mod.rs rename to crates/engineioxide/src/transport/polling.rs index e5207d35..1d618e79 100644 --- a/crates/engineioxide/src/transport/polling/mod.rs +++ b/crates/engineioxide/src/transport/polling.rs @@ -2,32 +2,22 @@ use std::sync::Arc; use bytes::Bytes; -use futures_util::StreamExt; -use http::{Request, Response, StatusCode}; +use engineioxide_core::payload::{self, Payload}; +use futures_util::{StreamExt, stream}; +use http::{Request, Response, StatusCode, header::CONTENT_TYPE}; use http_body::Body; use http_body_util::Full; +use tokio_util::future::FutureExt as _; -use engineioxide_core::Sid; +use engineioxide_core::{Packet, ProtocolVersion, Sid, TransportType}; use crate::{ - DisconnectReason, - body::ResponseBody, - engine::EngineIo, - errors::Error, - handler::EngineIoHandler, - packet::{OpenPacket, Packet}, - service::{ProtocolVersion, TransportType}, - transport::polling::payload::Payload, + DisconnectReason, body::ResponseBody, engine::EngineIo, errors::Error, + handler::EngineIoHandler, transport::make_open_packet, }; -mod payload; - /// Create a response for http request -fn http_response( - code: StatusCode, - data: D, - is_binary: bool, -) -> Result>, http::Error> +fn http_response(code: StatusCode, data: D, is_binary: bool) -> Response> where D: Into, { @@ -42,13 +32,14 @@ where res.header(CONTENT_TYPE, "text/plain; charset=UTF-8") } .body(ResponseBody::custom_response(Full::new(body))) + .unwrap() } pub fn open_req( engine: Arc>, protocol: ProtocolVersion, req: Request, - #[cfg(feature = "v3")] supports_binary: bool, + supports_binary: bool, ) -> Result>, Error> where H: EngineIoHandler, @@ -58,11 +49,10 @@ where protocol, TransportType::Polling, req.into_parts().0, - #[cfg(feature = "v3")] supports_binary, ); - let packet = OpenPacket::new(TransportType::Polling, socket.id, &engine.config); + let packet = make_open_packet(TransportType::Polling, socket.id, &engine.config); socket.spawn_heartbeat(engine.config.ping_interval, engine.config.ping_timeout); @@ -81,7 +71,7 @@ where #[cfg(not(feature = "v3"))] packet }; - http_response(StatusCode::OK, packet, false).map_err(Error::Http) + Ok(http_response(StatusCode::OK, packet, false)) } /// Handle http polling request @@ -106,39 +96,56 @@ where #[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)?); + 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() { + let mut rx = match socket.internal_rx.try_lock() { Ok(s) => s, Err(_) => { socket.close(DisconnectReason::MultipleHttpPollingError); - return Err(Error::HttpErrorResponse(StatusCode::BAD_REQUEST)); + return Err(Error::MultipleHttpPolling); } }; #[cfg(feature = "tracing")] - tracing::debug!("[sid={sid}] polling request"); + tracing::debug!(%sid, %protocol, supports_binary = socket.supports_binary, "polling request"); let max_payload = engine.config.max_payload; - #[cfg(feature = "v3")] - let Payload { data, has_binary } = - payload::encoder(rx, protocol, socket.supports_binary, max_payload).await?; - #[cfg(not(feature = "v3"))] - let Payload { data, has_binary } = payload::encoder(rx, protocol, max_payload).await?; + // Prepend the packet peeked during the previous + // polling request so it is encoded first, ahead of the newly received packets. + let rx_stream = stream::iter(rx.peeked.take()).chain(rx_stream::ReceiverStream::new(&mut rx)); + + // Stop waiting for the next packet as soon as the session is closed. The combinator is biased + // towards the encoder completion, so any buffered close packet is still flushed to the client. + let payload = payload::encoder(rx_stream, protocol, socket.supports_binary, max_payload) + .with_cancellation_token(&socket.cancellation_token) + .await; + + let Some(Payload { + data, + has_binary, + peeked, + }) = payload + else { + #[cfg(feature = "tracing")] + tracing::debug!(%sid, "session closed while polling, returning empty payload"); + return Ok(http_response(StatusCode::OK, "", false)); + }; + + // set back the peeked packet so it can be read again + // on the next polling request + rx.peeked = peeked; #[cfg(feature = "tracing")] - tracing::debug!("[sid={sid}] sending data: {:?}", data); - Ok(http_response(StatusCode::OK, data, has_binary)?) + tracing::trace!(%sid, %protocol, supports_binary = socket.supports_binary, "sending data: {:?}", data); + + Ok(http_response(StatusCode::OK, data, has_binary)) } /// Handle http polling post request @@ -148,7 +155,7 @@ pub async fn post_req( engine: Arc>, protocol: ProtocolVersion, sid: Sid, - body: Request, + req: Request, ) -> Result>, Error> where H: EngineIoHandler, @@ -162,15 +169,17 @@ where return Err(Error::TransportMismatch); } - let packets = payload::decoder(body, protocol, engine.config.max_payload); + let (parts, body) = req.into_parts(); + let content_type = parts.headers.get(CONTENT_TYPE); + let packets = payload::decoder(body, content_type, protocol, engine.config.max_payload); futures_util::pin_mut!(packets); while let Some(packet) = packets.next().await { match packet { Ok(Packet::Close) => { #[cfg(feature = "tracing")] - tracing::debug!("[sid={sid}] closing session"); - socket.send(Packet::Noop)?; + tracing::debug!(%sid, "received close packet, closing session"); + socket.send(Packet::Noop).ok(); // if the send fails, let's forcefully close the socket engine.close_session(sid, DisconnectReason::TransportClose); break; } @@ -188,16 +197,46 @@ where } Ok(p) => { #[cfg(feature = "tracing")] - tracing::debug!("[sid={sid}] bad packet received: {:?}", &p); + tracing::debug!(%sid, "invalid packet received: {:?}", &p); Err(Error::BadPacket(p)) } Err(e) => { #[cfg(feature = "tracing")] - tracing::debug!("[sid={sid}] error parsing packet: {:?}", e); + tracing::debug!(%sid, "could not parse packet: {e}"); engine.close_session(sid, DisconnectReason::PacketParsingError); - return Err(e); + return Err(e.into()); } }?; } - Ok(http_response(StatusCode::OK, "ok", false)?) + Ok(http_response(StatusCode::OK, "ok", false)) +} + +mod rx_stream { + use std::{ + pin::Pin, + task::{Context, Poll}, + }; + + use futures_core::Stream; + use tokio::sync::mpsc::Receiver; + + /// [`ReceiverStream`] is a stream that wraps a tokio::sync::mpsc::Receiver by reference. + /// Allowing to use it as a stream even if it is behind a mutex. + pub struct ReceiverStream<'a, T> { + inner: &'a mut Receiver, + } + + impl<'a, T> ReceiverStream<'a, T> { + pub fn new(inner: &'a mut Receiver) -> Self { + Self { inner } + } + } + + impl<'a, T> Stream for ReceiverStream<'a, T> { + type Item = T; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_recv(cx) + } + } } diff --git a/crates/engineioxide/src/transport/polling/payload/mod.rs b/crates/engineioxide/src/transport/polling/payload/mod.rs deleted file mode 100644 index db62cc96..00000000 --- a/crates/engineioxide/src/transport/polling/payload/mod.rs +++ /dev/null @@ -1,122 +0,0 @@ -//! Payload encoder and decoder for polling transport. - -use crate::{ - errors::Error, packet::Packet, peekable::PeekableReceiver, service::ProtocolVersion, - socket::PacketBuf, -}; -use bytes::Bytes; -use futures_core::Stream; -use http::Request; -use tokio::sync::MutexGuard; - -mod buf; -mod decoder; -mod encoder; - -const PACKET_SEPARATOR_V4: u8 = b'\x1e'; -#[cfg(feature = "v3")] -const STRING_PACKET_SEPARATOR_V3: u8 = b':'; -#[cfg(feature = "v3")] -const BINARY_PACKET_SEPARATOR_V3: u8 = 0xff; -#[cfg(feature = "v3")] -const STRING_PACKET_IDENTIFIER_V3: u8 = 0x00; -#[cfg(feature = "v3")] -const BINARY_PACKET_IDENTIFIER_V3: u8 = 0x01; - -pub fn decoder( - body: Request + Unpin>, - #[allow(unused_variables)] protocol: ProtocolVersion, - max_payload: u64, -) -> impl Stream> { - #[cfg(feature = "v3")] - { - use futures_util::future::Either; - use http::header::CONTENT_TYPE; - #[cfg(feature = "tracing")] - tracing::debug!("decoding payload {:?}", body.headers().get(CONTENT_TYPE)); - let is_binary = - body.headers().get(CONTENT_TYPE) == Some(&"application/octet-stream".parse().unwrap()); - match protocol { - ProtocolVersion::V4 => Either::Left(decoder::v4_decoder(body, max_payload)), - ProtocolVersion::V3 if is_binary => { - Either::Right(Either::Left(decoder::v3_binary_decoder(body, max_payload))) - } - ProtocolVersion::V3 => { - Either::Right(Either::Right(decoder::v3_string_decoder(body, max_payload))) - } - } - } - - #[cfg(not(feature = "v3"))] - { - decoder::v4_decoder(body, max_payload) - } -} - -/// A payload to transmit to the client through http polling -#[non_exhaustive] -pub struct Payload { - pub data: Bytes, - pub has_binary: bool, -} -impl Payload { - pub fn new(data: Bytes, has_binary: bool) -> Self { - Self { data, has_binary } - } -} - -pub async fn encoder( - rx: MutexGuard<'_, PeekableReceiver>, - #[allow(unused_variables)] protocol: ProtocolVersion, - #[cfg(feature = "v3")] supports_binary: bool, - max_payload: u64, -) -> Result { - #[cfg(feature = "v3")] - { - match protocol { - ProtocolVersion::V4 => encoder::v4_encoder(rx, max_payload).await, - ProtocolVersion::V3 if supports_binary => { - encoder::v3_binary_encoder(rx, max_payload).await - } - ProtocolVersion::V3 => encoder::v3_string_encoder(rx, max_payload).await, - } - } - - #[cfg(not(feature = "v3"))] - { - 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 && packet.is_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() -} diff --git a/crates/engineioxide/src/transport/ws.rs b/crates/engineioxide/src/transport/ws.rs index 565a9940..f579fac3 100644 --- a/crates/engineioxide/src/transport/ws.rs +++ b/crates/engineioxide/src/transport/ws.rs @@ -11,32 +11,27 @@ use futures_util::{ stream::{SplitSink, SplitStream}, }; use http::{HeaderValue, Request, Response, StatusCode, request::Parts}; +use smallvec::smallvec; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::{ WebSocketStream, tungstenite::{ - Message, + self, Message, handshake::derive_accept_key, protocol::{Role, WebSocketConfig}, }, }; +use tokio_util::future::FutureExt as _; -use engineioxide_core::{Sid, Str}; +use engineioxide_core::{Packet, ProtocolVersion, Sid, Str, TransportType}; use crate::{ - DisconnectReason, Socket, - body::ResponseBody, - config::EngineIoConfig, - engine::EngineIo, - errors::Error, - handler::EngineIoHandler, - packet::{OpenPacket, Packet}, - service::ProtocolVersion, - service::TransportType, + DisconnectReason, Socket, body::ResponseBody, config::EngineIoConfig, engine::EngineIo, + errors::Error, handler::EngineIoHandler, transport::make_open_packet, }; /// Create a response for websocket upgrade -fn ws_response(ws_key: &HeaderValue) -> Result>, http::Error> { +fn ws_response(ws_key: &HeaderValue) -> Response> { let derived = derive_accept_key(ws_key.as_bytes()); let sec = derived.parse::().unwrap(); Response::builder() @@ -48,6 +43,7 @@ fn ws_response(ws_key: &HeaderValue) -> Result>, htt ) .header(http::header::SEC_WEBSOCKET_ACCEPT, sec) .body(ResponseBody::empty_response()) + .unwrap() // we are only using constants } /// Upgrade a websocket request to create a websocket connection. @@ -66,7 +62,7 @@ pub fn new_req( let ws_key = parts .headers .get("Sec-WebSocket-Key") - .ok_or(Error::HttpErrorResponse(StatusCode::BAD_REQUEST))? + .ok_or(Error::InvalidWebSocketKey)? .clone(); tokio::spawn(async move { @@ -99,7 +95,7 @@ pub fn new_req( } }); - Ok(ws_response(&ws_key)?) + Ok(ws_response(&ws_key)) } /// Handle a websocket connection upgrade @@ -129,31 +125,30 @@ where Some(socket) if socket.is_ws() => return Err(Error::MultipleWebsocketRequests), Some(socket) => { let mut ws = ws_init().await; - match tokio::time::timeout( - engine.config.upgrade_timeout, - upgrade_handshake::(&socket, &mut ws), - ) - .await + + let upgrade_fut = upgrade_handshake::(&socket, &mut ws); + match tokio::time::timeout(engine.config.upgrade_timeout, upgrade_fut) + .with_cancellation_token(&socket.cancellation_token) + .await { - Ok(res) => res?, - Err(_) => { + Some(Ok(res)) => res?, + Some(Err(_)) => { #[cfg(feature = "tracing")] tracing::debug!(?sid, "ws upgrade timed out, closing session"); engine.close_session(socket.id, DisconnectReason::TransportError); return Err(Error::Upgrade); } + None => { + #[cfg(feature = "tracing")] + tracing::debug!(?sid, "socket is being closed"); + return Err(Error::Upgrade); + } } (socket, ws) } } } else { - let socket = engine.create_session( - protocol, - TransportType::Websocket, - req_data, - #[cfg(feature = "v3")] - false, - ); + let socket = engine.create_session(protocol, TransportType::Websocket, req_data, false); #[cfg(feature = "tracing")] tracing::debug!("new websocket connection"); @@ -167,28 +162,32 @@ where }; let (tx, rx) = ws.split(); - // Pipe between websocket and internal socket channel - let cancel_token = socket.cancellation_token.clone(); - tokio::spawn( - cancel_token - .clone() - .run_until_cancelled_owned(forward_to_socket::(socket.clone(), tx)), - ); + // Pipe between websocket and internal socket channel. + tokio::spawn(forward_to_socket::(socket.clone(), tx)); // pipe between internal socket channel and websocket - cancel_token - .run_until_cancelled_owned(async move { - if let Err(ref e) = forward_to_handler(&engine, rx, &socket).await { - #[cfg(feature = "tracing")] - tracing::debug!("error when handling packet: {:?}", e); - if let Some(reason) = e.into() { - engine.close_session(socket.id, reason); - } - } else { - engine.close_session(socket.id, DisconnectReason::TransportClose); - } - }) - .await; + match forward_to_handler(&engine, rx, &socket) + .with_cancellation_token(&socket.cancellation_token) + .await + { + Some(Err(ref e)) => { + let reason = + Option::::from(e).unwrap_or(DisconnectReason::TransportError); + + #[cfg(feature = "tracing")] + tracing::debug!("error when handling packet: {:?}", e); + engine.close_session(socket.id, reason); + } + Some(Ok(())) => { + #[cfg(feature = "tracing")] + tracing::debug!(sid = %socket.id, "ws transport was closed"); + engine.close_session(socket.id, DisconnectReason::TransportClose); + } + None => { + #[cfg(feature = "tracing")] + tracing::debug!(sid = %socket.id, "socket is closing"); + } + } Ok(()) } @@ -203,8 +202,13 @@ where { while let Some(msg) = rx.try_next().await? { match msg { - Message::Text(msg) => match Packet::parse(socket.protocol, utf8_bytes_to_str(msg))? { - Packet::Close => break, + Message::Text(msg) => match Packet::parse(socket.protocol, ws_bytes_to_str(msg))? { + Packet::Close => { + #[cfg(feature = "tracing")] + tracing::debug!("[sid={}] closing session", socket.id); + engine.close_session(socket.id, DisconnectReason::TransportClose); + break; + } Packet::Pong | Packet::Ping => socket .heartbeat_tx .try_send(()) @@ -267,11 +271,6 @@ async fn forward_to_socket( tx.feed(Message::Binary(bin)).await } } - Packet::Close => { - tx.send(Message::Close(None)).await.ok(); - internal_rx.close(); - break; - }, // A Noop Packet maybe sent by the server to upgrade from a polling connection // In the case that the packet was not poll in time it will remain in the buffer and therefore // it should be discarded here @@ -288,7 +287,12 @@ async fn forward_to_socket( }; } - while let Some(items) = internal_rx.recv().await { + // Forward buffered packets until the channel is closed or the session is cancelled. + while let Some(Some(items)) = internal_rx + .recv() + .with_cancellation_token(&socket.cancellation_token) + .await + { for item in items { map_fn!(item); } @@ -301,6 +305,15 @@ async fn forward_to_socket( tx.flush().await.ok(); } + + // Flush any remaining buffered packets before closing the websocket. + while let Ok(items) = internal_rx.try_recv() { + for item in items { + map_fn!(item); + } + } + tx.flush().await.ok(); + tx.close().await.ok(); } /// Send a Engine.IO [`OpenPacket`] to initiate a websocket connection async fn init_handshake( @@ -311,7 +324,8 @@ async fn init_handshake( where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - let packet = Packet::Open(OpenPacket::new(TransportType::Websocket, sid, config)); + let packet = Packet::Open(make_open_packet(TransportType::Websocket, sid, config)); + let packet: String = packet.into(); ws.send(Message::Text(packet.into())).await?; Ok(()) } @@ -358,7 +372,11 @@ where socket.start_upgrade(); // We send a last Noop request to close a potential waiting polling request. - socket.send(Packet::Noop)?; + socket + .internal_tx + .send(smallvec![Packet::Noop]) + .await + .map_err(|_| Error::Upgrade)?; // 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. @@ -369,13 +387,13 @@ where Some(Ok(Message::Text(d))) => d, _ => Err(Error::Upgrade)?, }; - match Packet::parse(socket.protocol, utf8_bytes_to_str(msg))? { + match Packet::parse(socket.protocol, ws_bytes_to_str(msg))? { Packet::PingUpgrade => { #[cfg(feature = "tracing")] tracing::debug!("received first ping upgrade"); - // Respond with a PongUpgrade packet - ws.send(Message::Text(Packet::PongUpgrade.into())).await?; + ws.send(Message::Text(String::from(Packet::PongUpgrade).into())) + .await?; } p => Err(Error::BadPacket(p))?, }; @@ -394,7 +412,7 @@ where Err(Error::Upgrade)? } }; - match Packet::parse(socket.protocol, utf8_bytes_to_str(msg))? { + match Packet::parse(socket.protocol, ws_bytes_to_str(msg))? { Packet::Upgrade => { #[cfg(feature = "tracing")] tracing::debug!("ws upgraded successfully") @@ -406,7 +424,7 @@ where Ok(()) } -fn utf8_bytes_to_str(bytes: tokio_tungstenite::tungstenite::Utf8Bytes) -> Str { +fn ws_bytes_to_str(bytes: tungstenite::Utf8Bytes) -> Str { // SAFETY: the bytes are guaranteed to be valid UTF-8 by the tungstenite parser unsafe { Str::from_bytes_unchecked(bytes.into()) } } diff --git a/crates/engineioxide/tests/disconnect_reason.rs b/crates/engineioxide/tests/disconnect_reason.rs index 1a162155..fee26081 100644 --- a/crates/engineioxide/tests/disconnect_reason.rs +++ b/crates/engineioxide/tests/disconnect_reason.rs @@ -52,6 +52,81 @@ impl EngineIoHandler for MyHandler { } } +/// A handler that forwards every connected socket to the test so it can drive its lifecycle. +#[derive(Debug, Clone)] +struct CloseHandler { + socket_tx: mpsc::Sender>>, +} + +impl EngineIoHandler for CloseHandler { + type Data = (); + + fn on_connect(self: Arc, socket: Arc>) { + self.socket_tx.try_send(socket).unwrap(); + } + fn on_disconnect(&self, _socket: Arc>, _reason: DisconnectReason) {} + fn on_message(self: &Arc, _msg: Str, _socket: Arc>) {} + fn on_binary(self: &Arc, _data: Bytes, _socket: Arc>) {} +} + +/// Closing a websocket session should resolve [`Socket::closed`]. +#[tokio::test] +pub async fn ws_server_closing() { + let (socket_tx, mut socket_rx) = mpsc::channel(1); + let mut svc = create_server(CloseHandler { socket_tx }).await; + let _stream = create_ws_connection(&mut svc).await; + + let socket = socket_rx.recv().await.unwrap(); + socket.close(DisconnectReason::ClosingServer); + + tokio::time::timeout(Duration::from_millis(100), socket.closed()) + .await + .expect("timeout: ws session did not close while the forward task held the channel lock"); + assert!(socket.is_closed()); +} + +/// Same as [`ws_server_closing`] but for polling with a long-poll request in-flight. +#[tokio::test] +pub async fn polling_server_closing_pending_poll() { + let (socket_tx, mut socket_rx) = mpsc::channel(1); + let mut svc = create_server(CloseHandler { socket_tx }).await; + let sid = create_polling_connection(&mut svc).await; + let socket = socket_rx.recv().await.unwrap(); + + // Drain the first buffered ping so the next poll actually blocks. + send_req( + &mut svc, + format!("transport=polling&sid={sid}"), + http::Method::GET, + None, + ) + .await; + + // Start a long-poll: the buffer is empty so it blocks waiting for the next packet while holding + // the internal channel lock. + let poll = tokio::spawn(send_req( + &mut svc, + format!("transport=polling&sid={sid}"), + http::Method::GET, + None, + )); + // Give the poll time to acquire the internal channel lock before closing. + tokio::time::sleep(Duration::from_millis(10)).await; + + socket.close(DisconnectReason::ClosingServer); + + tokio::time::timeout(Duration::from_millis(100), socket.closed()) + .await + .expect("timeout: polling session did not close while a poll held the channel lock"); + assert!(socket.is_closed()); + + // The pending poll should return rather than hang. + tokio::time::timeout(Duration::from_millis(100), poll) + .await + .expect("timeout: the pending poll never returned") + .unwrap(); +} + #[tokio::test] pub async fn polling_heartbeat_timeout() { let (disconnect_tx, mut rx) = mpsc::channel(10); diff --git a/crates/engineioxide/tests/fixture.rs b/crates/engineioxide/tests/fixture.rs index f1d5eef7..86158d2b 100644 --- a/crates/engineioxide/tests/fixture.rs +++ b/crates/engineioxide/tests/fixture.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + use std::{ collections::VecDeque, future::Future, @@ -74,6 +76,38 @@ pub fn send_req( } } +/// Same as [`send_req`] but returns the full response body without stripping the +/// leading packet-type byte. Useful to assert on raw polling payloads. +pub fn send_req_raw( + svc: &mut EngineIoService, + params: String, + method: http::Method, + body: Option, +) -> impl Future + 'static { + let body = match body { + Some(b) => Either::Left(Full::new(VecDeque::from(b.into_bytes()))), + None => Either::Right(Empty::>::new()), + }; + + let req = Request::builder() + .method(method) + .uri(format!("http://127.0.0.1/engine.io/?EIO=4&{params}")) + .body(body) + .unwrap(); + let res = svc.call(req); + async move { + let body = res + .await + .unwrap() + .body_mut() + .collect() + .await + .unwrap() + .to_bytes(); + String::from_utf8(body.to_vec()).unwrap() + } +} + pub async fn create_polling_connection(svc: &mut EngineIoService) -> String { let body = send_req( svc, diff --git a/crates/engineioxide/tests/peek_packet.rs b/crates/engineioxide/tests/peek_packet.rs new file mode 100644 index 00000000..08327b26 --- /dev/null +++ b/crates/engineioxide/tests/peek_packet.rs @@ -0,0 +1,115 @@ +//! Tests ensuring that a packet peeked but rejected because it would exceed the +//! configured `max_payload` is not lost: it must be sent on the next polling +//! request, in order, rather than dropped together with the encoder's peek +//! buffer. + +use std::{sync::Arc, time::Duration}; + +use bytes::Bytes; +use engineioxide::{ + Str, + config::EngineIoConfig, + handler::EngineIoHandler, + service::EngineIoService, + socket::{DisconnectReason, Socket}, +}; +use tokio::sync::mpsc; + +mod fixture; + +use fixture::{create_polling_connection, send_req_raw}; + +/// A handler that forwards every connected socket to the test so it can emit on it directly. +#[derive(Debug, Clone)] +struct SocketHandler { + socket_tx: mpsc::Sender>>, +} + +impl EngineIoHandler for SocketHandler { + type Data = (); + + fn on_connect(self: Arc, socket: Arc>) { + self.socket_tx.try_send(socket).unwrap(); + } + fn on_disconnect(&self, _socket: Arc>, _reason: DisconnectReason) {} + fn on_message(self: &Arc, _msg: Str, _socket: Arc>) {} + fn on_binary(self: &Arc, _data: Bytes, _socket: Arc>) {} +} + +/// Build a service with a small `max_payload` (so payloads are truncated after a +/// single packet) and heartbeat timings large enough not to interfere. +async fn create_small_payload_server( + max_payload: u64, +) -> ( + EngineIoService, + mpsc::Receiver>>, +) { + let (socket_tx, socket_rx) = mpsc::channel(1); + let config = EngineIoConfig::builder() + .ping_interval(Duration::from_secs(60)) + .ping_timeout(Duration::from_secs(60)) + .max_payload(max_payload) + .build(); + let svc = EngineIoService::with_config(Arc::new(SocketHandler { socket_tx }), config); + (svc, socket_rx) +} + +/// Drain the ping packet that the heartbeat job buffers right after connecting, +/// so subsequent polls only observe the packets emitted by the test. +async fn drain_initial_ping(svc: &mut EngineIoService, sid: &str) { + let ping = poll(svc, sid).await; + assert_eq!(ping, "2", "expected the initial heartbeat ping"); +} + +/// Perform a polling request and return the raw payload, failing fast instead of +/// hanging if the payload is unexpectedly empty (e.g. a peeked packet was lost). +async fn poll(svc: &mut EngineIoService, sid: &str) -> String { + let fut = send_req_raw( + svc, + format!("transport=polling&sid={sid}"), + http::Method::GET, + None, + ); + tokio::time::timeout(Duration::from_secs(1), fut) + .await + .expect("polling request timed out (a peeked packet was likely lost)") +} + +/// A packet rejected for exceeding `max_payload` must be delivered by the next poll. +#[tokio::test] +async fn polling_peeked_packet_sent_on_next_poll() { + // "4hello" is 6 bytes; two of them (+ separator) exceed the limit, so the + // first poll emits one packet and peeks-then-rejects the second. + let (mut svc, mut socket_rx) = create_small_payload_server(10).await; + let sid = create_polling_connection(&mut svc).await; + let socket = socket_rx.recv().await.unwrap(); + + drain_initial_ping(&mut svc, &sid).await; + + socket.emit("hello").unwrap(); + socket.emit("world").unwrap(); + + // First poll: only the first packet fits, the second is peeked and rejected. + assert_eq!(poll(&mut svc, &sid).await, "4hello"); + // Second poll: the rejected packet must be delivered here, not lost. + assert_eq!(poll(&mut svc, &sid).await, "4world"); +} + +/// Several packets rejected across consecutive polls must keep their order. +#[tokio::test] +async fn polling_peeked_packets_preserve_order() { + let (mut svc, mut socket_rx) = create_small_payload_server(10).await; + let sid = create_polling_connection(&mut svc).await; + let socket = socket_rx.recv().await.unwrap(); + + drain_initial_ping(&mut svc, &sid).await; + + socket.emit("hello").unwrap(); + socket.emit("world").unwrap(); + socket.emit("there").unwrap(); + + // Each poll drains exactly one packet, peeking and rejecting the next one. + assert_eq!(poll(&mut svc, &sid).await, "4hello"); + assert_eq!(poll(&mut svc, &sid).await, "4world"); + assert_eq!(poll(&mut svc, &sid).await, "4there"); +} diff --git a/crates/socketioxide/tests/disconnect_reason.rs b/crates/socketioxide/tests/disconnect_reason.rs index f9482b84..f5adf507 100644 --- a/crates/socketioxide/tests/disconnect_reason.rs +++ b/crates/socketioxide/tests/disconnect_reason.rs @@ -269,12 +269,66 @@ pub async fn server_ws_closing() { .await .expect("timeout waiting for server closing"); let packets = futures_util::future::join_all(streams.iter_mut().map(async |s| { - (s.next().await, s.next().await) // Closing packet / None + (s.next().await, s.next().await, s.next().await) // Closing packet / Closing WS msg / None })) .await; for packet in packets { - assert!(matches!(packet.0.unwrap().unwrap(), Message::Close(_))); - assert!(packet.1.is_none()); + assert!(matches!(packet.0.unwrap().unwrap(), Message::Text(_))); + assert!(matches!(packet.1.unwrap().unwrap(), Message::Close(_))); + assert!(packet.2.is_none()); + } +} + +/// Same as [`server_http_closing`] but with a long-polling request in-flight while the server closes +#[tokio::test] +pub async fn server_http_closing_pending_poll() { + let (svc, io) = create_server().await; + let _rx = attach_handler(&io, 100); + let sids = + futures_util::future::join_all((0..100).map(|_| create_polling_connection(&svc))).await; + + // Flush the buffered socket.io open packet so the next poll will actually block. + futures_util::future::join_all(sids.iter().map(|s| { + send_req( + &svc, + format!("transport=polling&sid={s}"), + http::Method::GET, + None, + ) + })) + .await; + + // Start a long-polling request for each session. As the buffer is now empty, each request blocks + // waiting for the next packet while holding the internal channel lock. + let pending_polls = sids + .iter() + .map(|s| { + let svc = svc.clone(); + let sid = s.clone(); + tokio::spawn(async move { + send_req( + &svc, + format!("transport=polling&sid={sid}"), + http::Method::GET, + None, + ) + .await + }) + }) + .collect::>(); + + // Give the pending polls time to acquire the internal channel lock before closing. + tokio::time::sleep(Duration::from_millis(10)).await; + + tokio::time::timeout(Duration::from_millis(100), io.close()) + .await + .expect("timeout waiting for server closing"); + + // The pending polls should return (with the close packet) rather than hang. + // The engine.io close packet is `1`; its packet-type char is stripped by `send_req`. + let packets = futures_util::future::join_all(pending_polls).await; + for packet in packets { + assert_eq!(packet.unwrap(), ""); } } diff --git a/examples/Cargo.lock b/examples/Cargo.lock index 24d5140e..57831639 100644 --- a/examples/Cargo.lock +++ b/examples/Cargo.lock @@ -1081,9 +1081,8 @@ dependencies = [ [[package]] name = "engineioxide" -version = "0.17.3" +version = "0.17.5" dependencies = [ - "base64", "bytes", "engineioxide-core", "futures-core", @@ -1112,8 +1111,15 @@ version = "0.2.1" dependencies = [ "base64", "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "pin-project-lite", "rand 0.10.1", "serde", + "serde_json", + "smallvec", ] [[package]] @@ -3789,7 +3795,7 @@ dependencies = [ [[package]] name = "socketioxide" -version = "0.18.3" +version = "0.18.4" dependencies = [ "bytes", "engineioxide", @@ -3814,7 +3820,7 @@ dependencies = [ [[package]] name = "socketioxide-core" -version = "0.18.0" +version = "0.18.1" dependencies = [ "arbitrary", "bytes",