dev(owned): trying to have a stream and a sink at the same time

Signed-off-by: Jad K. Haddad <jadkhaddad@gmail.com>

Author: JadKHaddad

Cargo.lock

@@ -19,6 +19,7 @@ thiserror = { version = "2.0.12", default-features = false }
 httparse = { version = "1.10.1", default-features = false }
 base64 = { version = "0.22.1", default-features = false }
 sha1 = { version = "0.10.6", default-features = false }
+heapless = { version = "0.9.1", default-features = false }
 
 [dev-dependencies]
 rand = { version = "0.9.1", features = ["std_rng"] }

Cargo.toml

@@ -16,6 +16,7 @@
         "fuzzingclient",
         "fuzzingserver",
         "Haddad",
+        "heapless",
         "httparse",
         "jadkhaddad",
         "linkall",
@@ -48,4 +49,4 @@
         "target",
         "autobahn/reports"
     ]
-}
+}

cspell.json

@@ -30,3 +30,45 @@ impl<'a> CloseFrame<'a> {
         self.reason
     }
 }
+
+/// An owned WebSocket Close frame.
+#[derive(Debug)]
+pub struct OwnedCloseFrame<const N: usize> {
+    /// The reason as a code.
+    code: CloseCode,
+    /// The reason as text string.
+    reason: heapless::String<N>,
+}
+
+impl<const N: usize> OwnedCloseFrame<N> {
+    /// Creates a new [`OwnedCloseFrame`].
+    pub const fn new(code: CloseCode, reason: heapless::String<N>) -> Self {
+        Self { code, reason }
+    }
+
+    /// Creates a new [`OwnedCloseFrame`] with no reason.
+    pub const fn no_reason(code: CloseCode) -> Self {
+        Self::new(code, heapless::String::new())
+    }
+
+    /// Returns the close code.
+    pub const fn code(&self) -> CloseCode {
+        self.code
+    }
+
+    /// Returns the reason as a string slice.
+    pub fn reason(&self) -> &str {
+        &self.reason
+    }
+}
+
+impl<const N: usize> TryFrom<CloseFrame<'_>> for OwnedCloseFrame<N> {
+    type Error = heapless::CapacityError;
+
+    fn try_from(value: CloseFrame<'_>) -> Result<Self, Self::Error> {
+        Ok(Self {
+            code: value.code(),
+            reason: heapless::String::try_from(value.reason())?,
+        })
+    }
+}

src/close_frame.rs

@@ -3,7 +3,7 @@ use framez::state::{ReadState, WriteState};
 use rand::RngCore;
 
 use crate::{
-    ConnectionState, Frame, Message, OnFrame, WebSocketCore,
+    ConnectionState, Frame, Message, OnFrame, OwnedMessage, WebSocketCore,
     codec::FramesCodec,
     error::{Error, ProtocolError, ReadError, WriteError},
     websocket_core::FragmentsState,
@@ -57,6 +57,64 @@ impl ReadAutoCaller {
         WebSocketCore::<RW, Rng>::on_frame(fragments_state, frame)
             .map(|result| result.map_err(Error::from))
     }
+
+    // TODO: see comments in on WebSocket::next()
+    // TODO: delete
+    #[allow(clippy::too_many_arguments)]
+    pub async fn call_next<const N: usize, F, RW, Rng>(
+        &self,
+        auto: F,
+        websocket: &mut WebSocketCore<'_, RW, Rng>,
+    ) -> Option<Result<Option<Result<OwnedMessage<N>, heapless::CapacityError>>, Error<RW::Error>>>
+    where
+        RW: Read + Write,
+        Rng: RngCore,
+        F: FnOnce(Frame<'_>) -> Result<OnFrame<'_>, ProtocolError> + 'static,
+    {
+        let frame = match framez::functions::maybe_next(
+            &mut websocket.framed.core.state.read,
+            &mut websocket.framed.core.codec,
+            &mut websocket.framed.core.inner,
+        )
+        .await
+        {
+            Some(Ok(Some(frame))) => frame,
+            Some(Ok(None)) => return Some(Ok(None)),
+            Some(Err(err)) => return Some(Err(Error::Read(ReadError::ReadFrame(err)))),
+            None => return None,
+        };
+
+        let frame = match auto(frame) {
+            Ok(on_frame) => match on_frame {
+                OnFrame::Send(message) => {
+                    websocket.state.closed = message.is_close();
+
+                    match framez::functions::send(
+                        &mut websocket.framed.core.state.write,
+                        &mut websocket.framed.core.codec,
+                        &mut websocket.framed.core.inner,
+                        message,
+                    )
+                    .await
+                    {
+                        Ok(_) => match websocket.state.closed {
+                            false => return Some(Ok(None)),
+                            true => return None,
+                        },
+                        Err(err) => return Some(Err(Error::Write(WriteError::WriteFrame(err)))),
+                    }
+                }
+                OnFrame::Noop(frame) => frame,
+            },
+            Err(err) => return Some(Err(Error::Read(ReadError::Protocol(err)))),
+        };
+
+        WebSocketCore::<RW, Rng>::on_frame(&mut websocket.fragments_state, frame).map(|result| {
+            result
+                .map(|opt| opt.map(OwnedMessage::<N>::try_from))
+                .map_err(Error::from)
+        })
+    }
 }
 
 #[derive(Debug)]

src/functions.rs

@@ -348,7 +348,7 @@ mod close_code;
 pub use close_code::CloseCode;
 
 mod close_frame;
-pub use close_frame::CloseFrame;
+pub use close_frame::{CloseFrame, OwnedCloseFrame};
 
 mod codec;
 use codec::FramesCodec;
@@ -368,7 +368,7 @@ pub mod http;
 mod mask;
 
 mod message;
-pub use message::Message;
+pub use message::{Message, OwnedMessage};
 
 #[doc(hidden)]
 pub mod mock;

src/lib.rs

@@ -1,4 +1,7 @@
-use crate::{CloseFrame, Frame, OpCode, error::FragmentationError, fragments::FragmentsIterator};
+use crate::{
+    CloseFrame, Frame, OpCode, OwnedCloseFrame, error::FragmentationError,
+    fragments::FragmentsIterator,
+};
 
 /// A WebSocket message.
 #[derive(Debug)]
@@ -122,3 +125,40 @@ impl<'a> Message<'a> {
         }
     }
 }
+
+/// An owned WebSocket message.
+#[derive(Debug)]
+pub enum OwnedMessage<const N: usize> {
+    /// A text WebSocket message
+    Text(heapless::String<N>),
+    /// A binary WebSocket message
+    Binary(heapless::Vec<u8, N>),
+    /// A ping message with the specified payload
+    ///
+    /// The payload here must have a length less than 125 bytes
+    Ping(heapless::Vec<u8, N>),
+    /// A pong message with the specified payload
+    ///
+    /// The payload here must have a length less than 125 bytes
+    Pong(heapless::Vec<u8, N>),
+    /// A close message with the optional close frame.
+    Close(Option<OwnedCloseFrame<N>>),
+}
+
+impl<const N: usize> TryFrom<Message<'_>> for OwnedMessage<N> {
+    type Error = heapless::CapacityError;
+
+    fn try_from(value: Message<'_>) -> Result<Self, Self::Error> {
+        match value {
+            Message::Text(payload) => Ok(Self::Text(heapless::String::try_from(payload)?)),
+            Message::Binary(payload) => Ok(Self::Binary(heapless::Vec::try_from(payload)?)),
+            Message::Ping(payload) => Ok(Self::Ping(heapless::Vec::try_from(payload)?)),
+            Message::Pong(payload) => Ok(Self::Pong(heapless::Vec::try_from(payload)?)),
+            Message::Close(Some(frame)) => Ok(Self::Close(Some(OwnedCloseFrame::new(
+                frame.code(),
+                heapless::String::try_from(frame.reason())?,
+            )))),
+            Message::Close(None) => Ok(Self::Close(None)),
+        }
+    }
+}

src/message.rs

@@ -1,3 +1,5 @@
+use core::cell::RefCell;
+
 use embedded_io_async::{Read, Write};
 use framez::{
     Framed,
@@ -6,9 +8,10 @@ use framez::{
 use rand::RngCore;
 
 use crate::{
-    FragmentsState, Frame, FramesCodec, Message, OnFrame, WebSocketCore,
+    FragmentsState, Frame, FramesCodec, Message, OnFrame, OwnedMessage, WebSocketCore,
     error::{Error, ProtocolError},
     http::{Request, Response},
+    next,
     options::{AcceptOptions, ConnectOptions},
 };
 
@@ -460,3 +463,41 @@ impl<'buf, RW, Rng> WebSocketWrite<'buf, RW, Rng> {
         self.core.send_fragmented(message, fragment_size).await
     }
 }
+
+#[derive(Debug)]
+pub struct OwnedWebSocket<'buf, const N: usize, RW, Rng> {
+    // XXX: the idea of the refcell is to have a stream and a sink at the same time using the same core
+    // but this would panic when using select on the stream and then using the sink to send a message, because the stream would have a mutable reference to the core and then the sink would try to borrow it again.
+    // See clippy warning
+    // I think we can only have a stream or a sink at a time, mutably borrowing the core
+    core: RefCell<WebSocket<'buf, RW, Rng>>,
+}
+
+impl<'buf, const N: usize, RW, Rng> OwnedWebSocket<'buf, N, RW, Rng>
+where
+    RW: Read + Write,
+    Rng: RngCore,
+{
+    pub async fn next(
+        &self,
+    ) -> Option<Result<Result<OwnedMessage<N>, heapless::CapacityError>, Error<RW::Error>>> {
+        let mut websocket = self.core.borrow_mut();
+
+        match 'next: loop {
+            match websocket
+                .caller()
+                .call_next::<N, _, _, _>(websocket.auto(), &mut websocket.core)
+                .await
+            {
+                Some(Ok(None)) => continue 'next,
+                Some(Ok(Some(item))) => break 'next Some(Ok(item)),
+                Some(Err(err)) => break 'next Some(Err(err)),
+                None => break 'next None,
+            }
+        } {
+            None => None,
+            Some(Ok(msg)) => Some(Ok(msg)),
+            Some(Err(err)) => Some(Err(err)),
+        }
+    }
+}