test(auto): added tests for auto close

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

Author: JadKHaddad

src/functions.rs

@@ -3,7 +3,7 @@ use framez::state::{ReadState, WriteState};
 use rand::RngCore;
 
 use crate::{
-    Frame, Message, OnFrame, WebSocketCore,
+    ConnectionState, Frame, Message, OnFrame, WebSocketCore,
     codec::FramesCodec,
     error::{Error, ProtocolError, ReadError, WriteError},
     websocket_core::FragmentsState,
@@ -89,12 +89,15 @@ pub async fn send<RW, Rng>(
     codec: &mut FramesCodec<Rng>,
     inner: &mut RW,
     write_state: &mut WriteState<'_>,
+    state: &mut ConnectionState,
     message: Message<'_>,
 ) -> Result<(), Error<RW::Error>>
 where
     RW: Write,
     Rng: RngCore,
 {
+    state.closed = message.is_close();
+
     framez::functions::send(write_state, codec, inner, message)
         .await
         .map_err(|err| Error::Write(WriteError::WriteFrame(err)))?;

src/lib.rs

@@ -40,7 +40,7 @@ use opcode::OpCode;
 pub mod options;
 
 mod websocket_core;
-use websocket_core::{FragmentsState, OnFrame, WebSocketCore};
+use websocket_core::{ConnectionState, FragmentsState, OnFrame, WebSocketCore};
 
 mod websocket;
 pub use websocket::{WebSocket, WebSocketRead, WebSocketWrite};

src/macros.rs

@@ -30,6 +30,7 @@ macro_rules! send {
             &mut $websocketz.core.framed.core.codec,
             &mut $websocketz.core.framed.core.inner,
             &mut $websocketz.core.framed.core.state.write,
+            &mut $websocketz.core.state,
             $message,
         )
         .await

src/tests.rs

@@ -955,5 +955,119 @@ mod fragmentation {
     }
 }
 
+mod auto {
+    use crate::CloseFrame;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn pong() {
+        let (client, server) = tokio::io::duplex(16);
+
+        let client = async move {
+            let read_buf = &mut [0u8; SIZE];
+            let write_buf = &mut [0u8; SIZE];
+            let fragments_buf = &mut [0u8; SIZE];
+
+            let mut websocketz = WebSocket::client(
+                FromTokio::new(client),
+                StdRng::from_os_rng(),
+                read_buf,
+                write_buf,
+                fragments_buf,
+            );
+
+            // Send a ping frame
+            websocketz
+                .send(Message::Ping(b"ping"))
+                .await
+                .expect("Failed to send ping message");
+
+            // Expect a pong frame in response
+            match next!(websocketz) {
+                Some(Ok(Message::Pong(payload))) => {
+                    assert_eq!(payload, b"ping");
+                }
+                message => panic!("Unexpected message: {message:?}"),
+            }
+        };
+
+        let server = async move {
+            let read_buf = &mut [0u8; SIZE];
+            let write_buf = &mut [0u8; SIZE];
+            let fragments_buf = &mut [0u8; SIZE];
+
+            let mut websocketz = WebSocket::server(
+                FromTokio::new(server),
+                StdRng::from_os_rng(),
+                read_buf,
+                write_buf,
+                fragments_buf,
+            );
+
+            while next!(websocketz).is_some() {}
+        };
+
+        tokio::join!(server, client);
+    }
+
+    #[tokio::test]
+    async fn close() {
+        let (client, server) = tokio::io::duplex(16);
+
+        let client = async move {
+            let read_buf = &mut [0u8; SIZE];
+            let write_buf = &mut [0u8; SIZE];
+            let fragments_buf = &mut [0u8; SIZE];
+
+            let mut websocketz = WebSocket::client(
+                FromTokio::new(client),
+                StdRng::from_os_rng(),
+                read_buf,
+                write_buf,
+                fragments_buf,
+            );
+
+            // Send a close frame
+            websocketz
+                .send(Message::Close(Some(CloseFrame::new(
+                    CloseCode::Normal,
+                    "close",
+                ))))
+                .await
+                .expect("Failed to send close message");
+
+            // Expect a close frame in response
+            match next!(websocketz) {
+                Some(Ok(Message::Close(Some(frame)))) => {
+                    assert_eq!(frame.code(), CloseCode::Normal);
+                    assert_eq!(frame.reason(), "close");
+                }
+                message => panic!("Unexpected message: {message:?}"),
+            }
+
+            // Ensure the connection is closed
+            assert!(next!(websocketz).is_none());
+        };
+
+        let server = async move {
+            let read_buf = &mut [0u8; SIZE];
+            let write_buf = &mut [0u8; SIZE];
+            let fragments_buf = &mut [0u8; SIZE];
+
+            let mut websocketz = WebSocket::server(
+                FromTokio::new(server),
+                StdRng::from_os_rng(),
+                read_buf,
+                write_buf,
+                fragments_buf,
+            );
+
+            while next!(websocketz).is_some() {}
+        };
+
+        tokio::join!(server, client);
+    }
+}
+
 // TODO: test for read errors
-// TODO: test auto ping and auto close

src/websocket_core.rs

@@ -47,7 +47,9 @@ struct Fragmented {
 
 #[derive(Debug, Clone, Copy)]
 struct Auto {
+    /// Auto pong frame handling.
     pong: bool,
+    /// Auto close frame handling.
     close: bool,
 }
 
@@ -61,12 +63,38 @@ impl Auto {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+#[doc(hidden)]
+pub struct ConnectionState {
+    /// If the user sends a close frame, we should not send a close frame back.
+    ///
+    /// Must be set to `true` if the user sends a close frame.
+    pub closed: bool,
+    /// Auto handling of ping/pong and close frames.
+    auto: Auto,
+}
+
+// TODO: add a new write error: ConnectionClosed. closed must be set to true if the user sends a close frame, if the reader read a close frame.
+// Then the closed field must be set to true. Every read will then return (None, means connection closed) and every write will return a write error with ConnectionClosed.
+// And then add the tests for that. If the user closes the connection or the server closed the connection, and then the user tries to read or write a frame
+
+impl ConnectionState {
+    #[inline]
+    #[allow(clippy::new_without_default)]
+    pub const fn new() -> Self {
+        Self {
+            closed: false,
+            auto: Auto::positive(),
+        }
+    }
+}
+
 #[derive(Debug)]
 #[doc(hidden)]
 pub struct WebSocketCore<'buf, RW, Rng> {
     pub framed: Framed<'buf, FramesCodec<Rng>, RW>,
     pub fragments_state: FragmentsState<'buf>,
-    auto: Auto,
+    pub state: ConnectionState,
 }
 
 impl<'buf, RW, Rng> WebSocketCore<'buf, RW, Rng> {
@@ -78,7 +106,7 @@ impl<'buf, RW, Rng> WebSocketCore<'buf, RW, Rng> {
         Self {
             framed,
             fragments_state,
-            auto: Auto::positive(),
+            state: ConnectionState::new(),
         }
     }
 
@@ -142,12 +170,12 @@ impl<'buf, RW, Rng> WebSocketCore<'buf, RW, Rng> {
 
     #[inline]
     pub(crate) const fn set_auto_pong(&mut self, auto_pong: bool) {
-        self.auto.pong = auto_pong;
+        self.state.auto.pong = auto_pong;
     }
 
     #[inline]
     pub(crate) const fn set_auto_close(&mut self, auto_close: bool) {
-        self.auto.close = auto_close;
+        self.state.auto.close = auto_close;
     }
 
     /// Returns reference to the reader/writer.
@@ -401,14 +429,14 @@ impl<'buf, RW, Rng> WebSocketCore<'buf, RW, Rng> {
     pub const fn auto(
         &self,
     ) -> impl FnOnce(Frame<'_>) -> Result<OnFrame<'_>, ProtocolError> + 'static {
-        let auto = self.auto;
+        let state = self.state;
 
         move |frame| {
-            if auto.pong && frame.opcode() == OpCode::Ping {
+            if state.auto.pong && frame.opcode() == OpCode::Ping {
                 return Ok(OnFrame::Send(Message::Pong(frame.payload())));
             }
 
-            if auto.close && frame.opcode() == OpCode::Close {
+            if state.auto.close && frame.opcode() == OpCode::Close && !state.closed {
                 let close_frame = match Self::extract_close_frame(&frame) {
                     Ok(close_frame) => close_frame,
                     Err(err) => return Err(err),
@@ -583,6 +611,7 @@ impl<'buf, RW, Rng> WebSocketCore<'buf, RW, Rng> {
             &mut self.framed.core.codec,
             &mut self.framed.core.inner,
             &mut self.framed.core.state.write,
+            &mut self.state,
             message,
         )
         .await