Refactor WebSocket implementation to separate inner representation into InnerWebSocket

Signed-off-by: Yuki Kishimoto <yukikishimoto@protonmail.com>

Author: yukibtc

src/native/mod.rs

@@ -57,7 +57,7 @@ async fn connect_direct(url: &Url, timeout: Duration) -> Result<WebSocket, Error
     ))
     .await
     .map_err(|_| Error::Timeout)??;
-    Ok(WebSocket::Tokio(Box::new(stream)))
+    Ok(WebSocket::tokio(Box::new(stream)))
 }
 
 #[cfg(feature = "socks")]
@@ -81,7 +81,7 @@ async fn connect_proxy(
     ))
     .await
     .map_err(|_| Error::Timeout)??;
-    Ok(WebSocket::Tokio(Box::new(stream)))
+    Ok(WebSocket::tokio(Box::new(stream)))
 }
 
 #[cfg(feature = "tor")]
@@ -104,7 +104,7 @@ async fn connect_tor(
     ))
     .await
     .map_err(|_| Error::Timeout)??;
-    Ok(WebSocket::Tor(Box::new(stream)))
+    Ok(WebSocket::tor(Box::new(stream)))
 }
 
 #[inline]

src/socket.rs

@@ -1,7 +1,6 @@
 // Copyright (c) 2022-2024 Yuki Kishimoto
 // Distributed under the MIT software license
 
-use std::ops::DerefMut;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use std::time::Duration;
@@ -22,7 +21,7 @@ use crate::{ConnectionMode, Error, Message};
 #[cfg(not(target_arch = "wasm32"))]
 type WsStream<T> = WebSocketStream<MaybeTlsStream<T>>;
 
-pub enum WebSocket {
+enum InnerWebSocket {
     #[cfg(not(target_arch = "wasm32"))]
     Tokio(Box<WsStream<TcpStream>>),
     #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
@@ -31,7 +30,34 @@ pub enum WebSocket {
     Wasm(WsStream),
 }
 
+pub struct WebSocket {
+    inner: InnerWebSocket,
+}
+
 impl WebSocket {
+    #[inline]
+    fn new(inner: InnerWebSocket) -> Self {
+        Self { inner }
+    }
+
+    #[inline]
+    #[cfg(not(target_arch = "wasm32"))]
+    pub(crate) fn tokio(inner: Box<WsStream<TcpStream>>) -> Self {
+        Self::new(InnerWebSocket::Tokio(inner))
+    }
+
+    #[inline]
+    #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
+    pub(crate) fn tor(inner: Box<WsStream<DataStream>>) -> Self {
+        Self::new(InnerWebSocket::Tor(inner))
+    }
+
+    #[inline]
+    #[cfg(target_arch = "wasm32")]
+    pub(crate) fn wasm(inner: WsStream) -> Self {
+        Self::new(InnerWebSocket::Wasm(inner))
+    }
+
     pub async fn connect(
         url: &Url,
         _mode: &ConnectionMode,
@@ -51,50 +77,50 @@ impl Sink<Message> for WebSocket {
     type Error = Error;
 
     fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        match self.deref_mut() {
+        match &mut self.inner {
             #[cfg(not(target_arch = "wasm32"))]
-            Self::Tokio(s) => Pin::new(s.as_mut()).poll_ready(cx).map_err(Into::into),
+            InnerWebSocket::Tokio(s) => Pin::new(s.as_mut()).poll_ready(cx).map_err(Into::into),
             #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
-            Self::Tor(s) => Pin::new(s.as_mut()).poll_ready(cx).map_err(Into::into),
+            InnerWebSocket::Tor(s) => Pin::new(s.as_mut()).poll_ready(cx).map_err(Into::into),
             #[cfg(target_arch = "wasm32")]
-            Self::Wasm(s) => Pin::new(s).poll_ready(cx),
+            InnerWebSocket::Wasm(s) => Pin::new(s).poll_ready(cx),
         }
     }
 
     fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
-        match self.deref_mut() {
+        match &mut self.inner {
             #[cfg(not(target_arch = "wasm32"))]
-            Self::Tokio(s) => Pin::new(s.as_mut())
+            InnerWebSocket::Tokio(s) => Pin::new(s.as_mut())
                 .start_send(item.into())
                 .map_err(Into::into),
             #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
-            Self::Tor(s) => Pin::new(s.as_mut())
+            InnerWebSocket::Tor(s) => Pin::new(s.as_mut())
                 .start_send(item.into())
                 .map_err(Into::into),
             #[cfg(target_arch = "wasm32")]
-            Self::Wasm(s) => Pin::new(s).start_send(item),
+            InnerWebSocket::Wasm(s) => Pin::new(s).start_send(item),
         }
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        match self.deref_mut() {
+        match &mut self.inner {
             #[cfg(not(target_arch = "wasm32"))]
-            Self::Tokio(s) => Pin::new(s.as_mut()).poll_flush(cx).map_err(Into::into),
+            InnerWebSocket::Tokio(s) => Pin::new(s.as_mut()).poll_flush(cx).map_err(Into::into),
             #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
-            Self::Tor(s) => Pin::new(s.as_mut()).poll_flush(cx).map_err(Into::into),
+            InnerWebSocket::Tor(s) => Pin::new(s.as_mut()).poll_flush(cx).map_err(Into::into),
             #[cfg(target_arch = "wasm32")]
-            Self::Wasm(s) => Pin::new(s).poll_flush(cx),
+            InnerWebSocket::Wasm(s) => Pin::new(s).poll_flush(cx),
         }
     }
 
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        match self.deref_mut() {
+        match &mut self.inner {
             #[cfg(not(target_arch = "wasm32"))]
-            Self::Tokio(s) => Pin::new(s.as_mut()).poll_close(cx).map_err(Into::into),
+            InnerWebSocket::Tokio(s) => Pin::new(s.as_mut()).poll_close(cx).map_err(Into::into),
             #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
-            Self::Tor(s) => Pin::new(s.as_mut()).poll_close(cx).map_err(Into::into),
+            InnerWebSocket::Tor(s) => Pin::new(s.as_mut()).poll_close(cx).map_err(Into::into),
             #[cfg(target_arch = "wasm32")]
-            Self::Wasm(s) => Pin::new(s).poll_close(cx).map_err(Into::into),
+            InnerWebSocket::Wasm(s) => Pin::new(s).poll_close(cx).map_err(Into::into),
         }
     }
 }
@@ -103,30 +129,30 @@ impl Stream for WebSocket {
     type Item = Result<Message, Error>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        match self.deref_mut() {
+        match &mut self.inner {
             #[cfg(not(target_arch = "wasm32"))]
-            Self::Tokio(s) => Pin::new(s)
+            InnerWebSocket::Tokio(s) => Pin::new(s)
                 .poll_next(cx)
                 .map(|i| i.map(|res| res.map(Message::from_native)))
                 .map_err(Into::into),
             #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
-            Self::Tor(s) => Pin::new(s)
+            InnerWebSocket::Tor(s) => Pin::new(s)
                 .poll_next(cx)
                 .map(|i| i.map(|res| res.map(Message::from_native)))
                 .map_err(Into::into),
             #[cfg(target_arch = "wasm32")]
-            Self::Wasm(s) => Pin::new(s).poll_next(cx).map_err(Into::into),
+            InnerWebSocket::Wasm(s) => Pin::new(s).poll_next(cx).map_err(Into::into),
         }
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        match self {
+        match &self.inner {
             #[cfg(not(target_arch = "wasm32"))]
-            Self::Tokio(s) => s.size_hint(),
+            InnerWebSocket::Tokio(s) => s.size_hint(),
             #[cfg(all(feature = "tor", not(target_arch = "wasm32")))]
-            Self::Tor(s) => s.size_hint(),
+            InnerWebSocket::Tor(s) => s.size_hint(),
             #[cfg(target_arch = "wasm32")]
-            Self::Wasm(s) => s.size_hint(),
+            InnerWebSocket::Wasm(s) => s.size_hint(),
         }
     }
 }

src/wasm/mod.rs

@@ -31,7 +31,7 @@ pub async fn connect(url: &Url, timeout: Duration) -> Result<WebSocket, Error> {
     let (_ws, stream) = time::timeout(Some(timeout), WasmWebSocket::connect(url))
         .await
         .ok_or(Error::Timeout)??;
-    Ok(WebSocket::Wasm(stream))
+    Ok(WebSocket::wasm(stream))
 }
 
 /// Helper function to reduce code bloat