chore(socket): simplify handle_engineio_packet

1c3t3a · Aug 12, 2022 · f3dc134 · f3dc134
1 parent 31996cf
commit f3dc134
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 52 deletions.
diff --git a/engineio/src/asynchronous/async_transports/websocket.rs b/engineio/src/asynchronous/async_transports/websocket.rs
@@ -53,6 +53,10 @@ impl WebsocketTransport {
     pub(crate) async fn upgrade(&self) -> Result<()> {
         self.inner.upgrade().await
     }
+
+    pub(crate) async fn poll_next(&self) -> Result<Option<Bytes>> {
+        self.inner.poll_next().await
+    }
 }
 
 #[async_trait]

diff --git a/engineio/src/asynchronous/async_transports/websocket_general.rs b/engineio/src/asynchronous/async_transports/websocket_general.rs
@@ -11,14 +11,17 @@ use tokio::{net::TcpStream, sync::Mutex};
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 use tungstenite::Message;
 
+type AsyncWebsocketSender = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>;
+type AsyncWebsocketReceiver = SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>;
+
 /// A general purpose asynchronous websocket transport type. Holds
 /// the sender and receiver stream of a websocket connection
 /// and implements the common methods `update` and `emit`. This also
 /// implements `Stream`.
 #[derive(Clone)]
 pub(crate) struct AsyncWebsocketGeneralTransport {
-    sender: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>,
-    receiver: Arc<Mutex<SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>>>,
+    sender: Arc<Mutex<AsyncWebsocketSender>>,
+    receiver: Arc<Mutex<AsyncWebsocketReceiver>>,
 }
 
 impl AsyncWebsocketGeneralTransport {
@@ -75,6 +78,27 @@ impl AsyncWebsocketGeneralTransport {
 
         Ok(())
     }
+
+    pub(crate) async fn poll_next(&self) -> Result<Option<Bytes>> {
+        loop {
+            let mut receiver = self.receiver.lock().await;
+            let next = receiver.next().await;
+            match next {
+                Some(Ok(Message::Text(str))) => return Ok(Some(Bytes::from(str))),
+                Some(Ok(Message::Binary(data))) => {
+                    let mut msg = BytesMut::with_capacity(data.len() + 1);
+                    msg.put_u8(PacketId::Message as u8);
+                    msg.put(data.as_ref());
+
+                    return Ok(Some(msg.freeze()));
+                }
+                // ignore packets other than text and binary
+                Some(Ok(_)) => (),
+                Some(Err(err)) => return Err(err.into()),
+                None => return Ok(None),
+            }
+        }
+    }
 }
 
 impl Stream for AsyncWebsocketGeneralTransport {

diff --git a/engineio/src/asynchronous/async_transports/websocket_secure.rs b/engineio/src/asynchronous/async_transports/websocket_secure.rs
@@ -62,6 +62,10 @@ impl WebsocketSecureTransport {
     pub(crate) async fn upgrade(&self) -> Result<()> {
         self.inner.upgrade().await
     }
+
+    pub(crate) async fn poll_next(&self) -> Result<Option<Bytes>> {
+        self.inner.poll_next().await
+    }
 }
 
 impl Stream for WebsocketSecureTransport {

diff --git a/engineio/src/transports/websocket.rs b/engineio/src/transports/websocket.rs
@@ -7,16 +7,15 @@ use crate::{
     Error,
 };
 use bytes::Bytes;
-use futures_util::StreamExt;
 use http::HeaderMap;
 use std::sync::Arc;
-use tokio::{runtime::Runtime, sync::Mutex};
+use tokio::runtime::Runtime;
 use url::Url;
 
 #[derive(Clone)]
 pub struct WebsocketTransport {
     runtime: Arc<Runtime>,
-    inner: Arc<Mutex<AsyncWebsocketTransport>>,
+    inner: Arc<AsyncWebsocketTransport>,
 }
 
 impl WebsocketTransport {
@@ -30,47 +29,40 @@ impl WebsocketTransport {
 
         Ok(WebsocketTransport {
             runtime: Arc::new(runtime),
-            inner: Arc::new(Mutex::new(inner)),
+            inner: Arc::new(inner),
         })
     }
 
     /// Sends probe packet to ensure connection is valid, then sends upgrade
     /// request
     pub(crate) fn upgrade(&self) -> Result<()> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.upgrade().await
-        })
+        self.runtime.block_on(async { self.inner.upgrade().await })
     }
 }
 
 impl Transport for WebsocketTransport {
     fn emit(&self, data: Bytes, is_binary_att: bool) -> Result<()> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.emit(data, is_binary_att).await
-        })
+        self.runtime
+            .block_on(async { self.inner.emit(data, is_binary_att).await })
     }
 
     fn poll(&self) -> Result<Bytes> {
         self.runtime.block_on(async {
-            let mut lock = self.inner.lock().await;
-            lock.next().await.ok_or(Error::IncompletePacket())?
+            let r = self.inner.poll_next().await;
+            match r {
+                Ok(b) => b.ok_or(Error::IncompletePacket()),
+                Err(_) => Err(Error::IncompletePacket()),
+            }
         })
     }
 
     fn base_url(&self) -> Result<url::Url> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.base_url().await
-        })
+        self.runtime.block_on(async { self.inner.base_url().await })
     }
 
     fn set_base_url(&self, url: url::Url) -> Result<()> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.set_base_url(url).await
-        })
+        self.runtime
+            .block_on(async { self.inner.set_base_url(url).await })
     }
 }
 

diff --git a/engineio/src/transports/websocket_secure.rs b/engineio/src/transports/websocket_secure.rs
@@ -8,17 +8,16 @@ use crate::{
     Error,
 };
 use bytes::Bytes;
-use futures_util::StreamExt;
 use http::HeaderMap;
 use native_tls::TlsConnector;
 use std::sync::Arc;
-use tokio::{runtime::Runtime, sync::Mutex};
+use tokio::runtime::Runtime;
 use url::Url;
 
 #[derive(Clone)]
 pub struct WebsocketSecureTransport {
     runtime: Arc<Runtime>,
-    inner: Arc<Mutex<AsyncWebsocketSecureTransport>>,
+    inner: Arc<AsyncWebsocketSecureTransport>,
 }
 
 impl WebsocketSecureTransport {
@@ -38,47 +37,40 @@ impl WebsocketSecureTransport {
 
         Ok(WebsocketSecureTransport {
             runtime: Arc::new(runtime),
-            inner: Arc::new(Mutex::new(inner)),
+            inner: Arc::new(inner),
         })
     }
 
     /// Sends probe packet to ensure connection is valid, then sends upgrade
     /// request
     pub(crate) fn upgrade(&self) -> Result<()> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.upgrade().await
-        })
+        self.runtime.block_on(async { self.inner.upgrade().await })
     }
 }
 
 impl Transport for WebsocketSecureTransport {
     fn emit(&self, data: Bytes, is_binary_att: bool) -> Result<()> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.emit(data, is_binary_att).await
-        })
+        self.runtime
+            .block_on(async { self.inner.emit(data, is_binary_att).await })
     }
 
     fn poll(&self) -> Result<Bytes> {
         self.runtime.block_on(async {
-            let mut lock = self.inner.lock().await;
-            lock.next().await.ok_or(Error::IncompletePacket())?
+            let r = self.inner.poll_next().await;
+            match r {
+                Ok(b) => b.ok_or(Error::IncompletePacket()),
+                Err(_) => Err(Error::IncompletePacket()),
+            }
         })
     }
 
     fn base_url(&self) -> Result<url::Url> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.base_url().await
-        })
+        self.runtime.block_on(async { self.inner.base_url().await })
     }
 
     fn set_base_url(&self, url: url::Url) -> Result<()> {
-        self.runtime.block_on(async {
-            let lock = self.inner.lock().await;
-            lock.set_base_url(url).await
-        })
+        self.runtime
+            .block_on(async { self.inner.set_base_url(url).await })
     }
 }
 

diff --git a/socketio/src/socket.rs b/socketio/src/socket.rs
@@ -159,12 +159,8 @@ impl Socket {
 
     /// Handles new incoming engineio packets
     fn handle_engineio_packet(&self, packet: EnginePacket) -> Result<Packet> {
-        let socket_packet = Packet::try_from(&packet.data);
-        if let Err(err) = socket_packet {
-            return Err(err);
-        }
-        // SAFETY: checked above to see if it was Err
-        let mut socket_packet = socket_packet.unwrap();
+        let mut socket_packet = Packet::try_from(&packet.data)?;
+
         // Only handle attachments if there are any
         if socket_packet.attachment_count > 0 {
             let mut attachments_left = socket_packet.attachment_count;