Add serilazer to where it needs

shenjackyuanjie · shenjackyuanjie · commit 5e96948274fc · 2024-05-14T00:33:18.000+08:00
diff --git a/engineio/src/asynchronous/async_socket.rs b/engineio/src/asynchronous/async_socket.rs
@@ -120,7 +120,10 @@ impl Socket {
     }
 
     /// Helper method that parses bytes and returns an iterator over the elements.
-    fn parse_payload(bytes: Bytes, serializer: Arc<PacketSerializer>) -> impl Stream<Item = Result<Packet>> {
+    fn parse_payload(
+        bytes: Bytes,
+        serializer: Arc<PacketSerializer>,
+    ) -> impl Stream<Item = Result<Packet>> {
         try_stream! {
             // let payload = Payload::try_from(bytes);
             let payload = serializer.decode_payload(bytes);
diff --git a/engineio/src/asynchronous/async_transports/websocket.rs b/engineio/src/asynchronous/async_transports/websocket.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 
 use crate::asynchronous::transport::AsyncTransport;
 use crate::error::Result;
+use crate::PacketSerializer;
 use async_trait::async_trait;
 use bytes::Bytes;
 use futures_util::stream::StreamExt;
@@ -27,7 +28,11 @@ pub struct WebsocketTransport {
 
 impl WebsocketTransport {
     /// Creates a new instance over a request that might hold additional headers and an URL.
-    pub async fn new(base_url: Url, headers: Option<HeaderMap>) -> Result<Self> {
+    pub async fn new(
+        base_url: Url,
+        headers: Option<HeaderMap>,
+        serializer: Arc<PacketSerializer>,
+    ) -> Result<Self> {
         let mut url = base_url;
         url.query_pairs_mut().append_pair("transport", "websocket");
         url.set_scheme("ws").unwrap();
@@ -41,7 +46,7 @@ impl WebsocketTransport {
         let (ws_stream, _) = connect_async(req).await?;
         let (sen, rec) = ws_stream.split();
 
-        let inner = AsyncWebsocketGeneralTransport::new(sen, rec).await;
+        let inner = AsyncWebsocketGeneralTransport::new(sen, rec, serializer).await;
         Ok(WebsocketTransport {
             inner,
             base_url: Arc::new(RwLock::new(url)),
@@ -118,7 +123,12 @@ mod test {
         let url = crate::test::engine_io_server()?.to_string()
             + "engine.io/?EIO="
             + &ENGINE_IO_VERSION.to_string();
-        WebsocketTransport::new(Url::from_str(&url[..])?, None).await
+        WebsocketTransport::new(
+            Url::from_str(&url[..])?,
+            None,
+            PacketSerializer::default_arc(),
+        )
+        .await
     }
 
     #[tokio::test]
diff --git a/engineio/src/asynchronous/async_transports/websocket_general.rs b/engineio/src/asynchronous/async_transports/websocket_general.rs
@@ -1,6 +1,6 @@
 use std::{borrow::Cow, str::from_utf8, sync::Arc, task::Poll};
 
-use crate::{error::Result, Error, Packet, PacketId};
+use crate::{error::Result, Error, Packet, PacketId, PacketSerializer};
 use bytes::{BufMut, Bytes, BytesMut};
 use futures_util::{
     ready,
@@ -22,16 +22,19 @@ type AsyncWebsocketReceiver = SplitStream<WebSocketStream<MaybeTlsStream<TcpStre
 pub(crate) struct AsyncWebsocketGeneralTransport {
     sender: Arc<Mutex<AsyncWebsocketSender>>,
     receiver: Arc<Mutex<AsyncWebsocketReceiver>>,
+    serializer: Arc<PacketSerializer>,
 }
 
 impl AsyncWebsocketGeneralTransport {
     pub(crate) async fn new(
         sender: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
         receiver: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
+        serializer: Arc<PacketSerializer>,
     ) -> Self {
         AsyncWebsocketGeneralTransport {
             sender: Arc::new(Mutex::new(sender)),
             receiver: Arc::new(Mutex::new(receiver)),
+            serializer,
         }
     }
 
@@ -41,25 +44,30 @@ impl AsyncWebsocketGeneralTransport {
         let mut receiver = self.receiver.lock().await;
         let mut sender = self.sender.lock().await;
 
+        let ping_packet = Packet::new(PacketId::Ping, Bytes::from("probe"));
+        let ping_packet = self.serializer.encode(ping_packet);
+
         sender
-            .send(Message::text(Cow::Borrowed(from_utf8(&Bytes::from(
-                Packet::new(PacketId::Ping, Bytes::from("probe")),
-            ))?)))
+            .send(Message::text(Cow::Borrowed(from_utf8(&ping_packet)?)))
             .await?;
 
         let msg = receiver
             .next()
             .await
             .ok_or(Error::IllegalWebsocketUpgrade())??;
 
-        if msg.into_data() != Bytes::from(Packet::new(PacketId::Pong, Bytes::from("probe"))) {
+        let pong_packet = Packet::new(PacketId::Pong, Bytes::from("probe"));
+        let pong_packet = self.serializer.encode(pong_packet);
+
+        if msg.into_data() != pong_packet {
             return Err(Error::InvalidPacket());
         }
 
+        let upgrade_packet = Packet::new(PacketId::Upgrade, Bytes::from(""));
+        let upgrade_packet = self.serializer.encode(upgrade_packet);
+
         sender
-            .send(Message::text(Cow::Borrowed(from_utf8(&Bytes::from(
-                Packet::new(PacketId::Upgrade, Bytes::from("")),
-            ))?)))
+            .send(Message::text(Cow::Borrowed(from_utf8(&upgrade_packet)?)))
             .await?;
 
         Ok(())
diff --git a/engineio/src/asynchronous/async_transports/websocket_secure.rs b/engineio/src/asynchronous/async_transports/websocket_secure.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 
 use crate::asynchronous::transport::AsyncTransport;
 use crate::error::Result;
+use crate::PacketSerializer;
 use async_trait::async_trait;
 use bytes::Bytes;
 use futures_util::Stream;
@@ -34,6 +35,7 @@ impl WebsocketSecureTransport {
         base_url: Url,
         tls_config: Option<TlsConnector>,
         headers: Option<HeaderMap>,
+        serializer: Arc<PacketSerializer>,
     ) -> Result<Self> {
         let mut url = base_url;
         url.query_pairs_mut().append_pair("transport", "websocket");
@@ -61,7 +63,7 @@ impl WebsocketSecureTransport {
         .await?;
 
         let (sen, rec) = ws_stream.split();
-        let inner = AsyncWebsocketGeneralTransport::new(sen, rec).await;
+        let inner = AsyncWebsocketGeneralTransport::new(sen, rec, serializer).await;
 
         Ok(WebsocketSecureTransport {
             inner,
@@ -143,6 +145,7 @@ mod test {
             Url::from_str(&url[..])?,
             Some(crate::test::tls_connector()?),
             None,
+            PacketSerializer::default_arc(),
         )
         .await
     }
diff --git a/engineio/src/asynchronous/client/builder.rs b/engineio/src/asynchronous/client/builder.rs
@@ -227,7 +227,9 @@ impl ClientBuilder {
 
         match self.url.scheme() {
             "http" | "ws" => {
-                let mut transport = WebsocketTransport::new(self.url.clone(), headers).await?;
+                let mut transport =
+                    WebsocketTransport::new(self.url.clone(), headers, self.serializer.clone())
+                        .await?;
 
                 if self.handshake.is_some() {
                     transport.upgrade().await?;
@@ -252,6 +254,7 @@ impl ClientBuilder {
                     self.url.clone(),
                     self.tls_config.clone(),
                     headers,
+                    self.serializer.clone(),
                 )
                 .await?;
 
diff --git a/engineio/src/client/client.rs b/engineio/src/client/client.rs
@@ -65,8 +65,8 @@ impl ClientBuilder {
     }
 
     /// Specify Packet Serializer
-    pub fn packet_serializer(mut self, packet_serializer: PacketSerializer) -> Self {
-        self.serializer = Arc::new(packet_serializer);
+    pub fn packet_serializer(mut self, packet_serializer: Arc<PacketSerializer>) -> Self {
+        self.serializer = packet_serializer;
 
         self
     }
@@ -228,7 +228,7 @@ impl ClientBuilder {
 
         match url.scheme() {
             "http" | "ws" => {
-                let transport = WebsocketTransport::new(url, headers)?;
+                let transport = WebsocketTransport::new(url, headers, self.serializer.clone())?;
                 if self.handshake.is_some() {
                     transport.upgrade()?;
                 } else {
@@ -250,8 +250,12 @@ impl ClientBuilder {
                 })
             }
             "https" | "wss" => {
-                let transport =
-                    WebsocketSecureTransport::new(url, self.tls_config.clone(), headers)?;
+                let transport = WebsocketSecureTransport::new(
+                    url,
+                    self.tls_config.clone(),
+                    headers,
+                    self.serializer.clone(),
+                )?;
                 if self.handshake.is_some() {
                     transport.upgrade()?;
                 } else {
diff --git a/engineio/src/lib.rs b/engineio/src/lib.rs
@@ -97,7 +97,7 @@ pub mod error;
 
 pub use client::{Client, ClientBuilder};
 pub use error::Error;
-pub use packet::{Packet, PacketId};
+pub use packet::{Packet, PacketId, PacketSerializer};
 
 #[cfg(test)]
 pub(crate) mod test {
diff --git a/engineio/src/packet.rs b/engineio/src/packet.rs
@@ -91,6 +91,10 @@ impl PacketSerializer {
         let _ = buf.split_off(buf.len() - 1);
         buf.freeze()
     }
+
+    pub fn default_arc() -> std::sync::Arc<Self> {
+        std::sync::Arc::new(Self::default())
+    }
 }
 
 impl Default for PacketSerializer {
diff --git a/engineio/src/transports/websocket.rs b/engineio/src/transports/websocket.rs
@@ -4,7 +4,7 @@ use crate::{
     },
     error::Result,
     transport::Transport,
-    Error,
+    Error, PacketSerializer,
 };
 use bytes::Bytes;
 use http::HeaderMap;
@@ -20,12 +20,17 @@ pub struct WebsocketTransport {
 
 impl WebsocketTransport {
     /// Creates an instance of `WebsocketTransport`.
-    pub fn new(base_url: Url, headers: Option<HeaderMap>) -> Result<Self> {
+    pub fn new(
+        base_url: Url,
+        headers: Option<HeaderMap>,
+        serializer: Arc<PacketSerializer>,
+    ) -> Result<Self> {
         let runtime = tokio::runtime::Builder::new_current_thread()
             .enable_all()
             .build()?;
 
-        let inner = runtime.block_on(AsyncWebsocketTransport::new(base_url, headers))?;
+        let inner =
+            runtime.block_on(AsyncWebsocketTransport::new(base_url, headers, serializer))?;
 
         Ok(WebsocketTransport {
             runtime: Arc::new(runtime),
@@ -90,7 +95,11 @@ mod test {
         let url = crate::test::engine_io_server()?.to_string()
             + "engine.io/?EIO="
             + &ENGINE_IO_VERSION.to_string();
-        WebsocketTransport::new(Url::from_str(&url[..])?, None)
+        WebsocketTransport::new(
+            Url::from_str(&url[..])?,
+            None,
+            PacketSerializer::default_arc(),
+        )
     }
 
     #[test]
diff --git a/engineio/src/transports/websocket_secure.rs b/engineio/src/transports/websocket_secure.rs
@@ -5,7 +5,7 @@ use crate::{
     },
     error::Result,
     transport::Transport,
-    Error,
+    Error, PacketSerializer,
 };
 use bytes::Bytes;
 use http::HeaderMap;
@@ -26,13 +26,14 @@ impl WebsocketSecureTransport {
         base_url: Url,
         tls_config: Option<TlsConnector>,
         headers: Option<HeaderMap>,
+        serializer: Arc<PacketSerializer>,
     ) -> Result<Self> {
         let runtime = tokio::runtime::Builder::new_current_thread()
             .enable_all()
             .build()?;
 
         let inner = runtime.block_on(AsyncWebsocketSecureTransport::new(
-            base_url, tls_config, headers,
+            base_url, tls_config, headers, serializer,
         ))?;
 
         Ok(WebsocketSecureTransport {
@@ -99,6 +100,7 @@ mod test {
             Url::from_str(&url[..])?,
             Some(crate::test::tls_connector()?),
             None,
+            PacketSerializer::default_arc(),
         )
     }
 
diff --git a/socketio/src/client/builder.rs b/socketio/src/client/builder.rs
@@ -33,7 +33,7 @@ impl Default for TransportType {
     }
 }
 
-pub use rust_engineio::Packet::PacketSerializer;
+pub use rust_engineio::PacketSerializer;
 
 /// A builder class for a `socket.io` socket. This handles setting up the client and
 /// configuring the callback, the namespace and metadata of the socket. If no
@@ -48,7 +48,7 @@ pub struct ClientBuilder {
     tls_config: Option<TlsConnector>,
     opening_headers: Option<HeaderMap>,
     transport_type: TransportType,
-    packet_serializer: PacketSerializer,
+    packet_serializer: Arc<PacketSerializer>,
     auth: Option<serde_json::Value>,
     pub(crate) reconnect: bool,
     pub(crate) reconnect_on_disconnect: bool,
@@ -100,7 +100,7 @@ impl ClientBuilder {
             tls_config: None,
             opening_headers: None,
             transport_type: TransportType::default(),
-            packet_serializer: PacketSerializer::default(),
+            packet_serializer: PacketSerializer::default_arc(),
             auth: None,
             reconnect: true,
             reconnect_on_disconnect: false,
@@ -335,7 +335,7 @@ impl ClientBuilder {
     /// }
     /// ```
     pub fn packet_serializer(mut self, packet_serializer: PacketSerializer) -> Self {
-        self.packet_serializer = packet_serializer;
+        self.packet_serializer = Arc::new(packet_serializer);
 
         self
     }
@@ -376,7 +376,7 @@ impl ClientBuilder {
         }
 
         let mut builder =
-            EngineIoClientBuilder::new(url).packet_serializer(self.packet_serializer.into());
+            EngineIoClientBuilder::new(url).packet_serializer(self.packet_serializer.clone());
 
         if let Some(tls_config) = self.tls_config {
             builder = builder.tls_config(tls_config);

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,10 @@ impl PacketSerializer {`
`91`	`91`	`let _ = buf.split_off(buf.len() - 1);`
`92`	`92`	`buf.freeze()`
`93`	`93`	`}`
	`94`	`+`
	`95`	`+ pub fn default_arc() -> std::sync::Arc<Self> {`
	`96`	`+ std::sync::Arc::new(Self::default())`
	`97`	`+ }`
`94`	`98`	`}`
`95`	`99`
`96`	`100`	`impl Default for PacketSerializer {`