feat(close): send after close returns a ConnectionClosed error

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

Author: JadKHaddad

src/error.rs

@@ -85,6 +85,9 @@ pub enum ReadError<I> {
 
 #[derive(Debug, thiserror::Error)]
 pub enum WriteError<I> {
+    /// Websocket connection is closed.
+    ///
+    /// To close the TCP connection, you should drop the [`WebSocket`](crate::WebSocket) instance.
     #[error("Connection closed")]
     ConnectionClosed,
     #[error("Write frame error: {0}")]

src/functions.rs

@@ -13,6 +13,7 @@ use crate::{
 pub struct ReadAutoCaller;
 
 impl ReadAutoCaller {
+    #[allow(clippy::too_many_arguments)]
     pub async fn call<'this, F, RW, Rng>(
         &self,
         auto: F,
@@ -21,6 +22,7 @@ impl ReadAutoCaller {
         read_state: &'this mut ReadState<'_>,
         write_state: &mut WriteState<'_>,
         fragments_state: &'this mut FragmentsState<'_>,
+        state: &mut ConnectionState,
     ) -> Option<Result<Option<Message<'this>>, Error<RW::Error>>>
     where
         RW: Read + Write,
@@ -37,10 +39,10 @@ impl ReadAutoCaller {
         let frame = match auto(frame) {
             Ok(on_frame) => match on_frame {
                 OnFrame::Send(message) => {
-                    let is_close = message.is_close();
+                    state.closed = message.is_close();
 
                     match framez::functions::send(write_state, codec, inner, message).await {
-                        Ok(_) => match is_close {
+                        Ok(_) => match state.closed {
                             false => return Some(Ok(None)),
                             true => return None,
                         },
@@ -61,6 +63,7 @@ impl ReadAutoCaller {
 pub struct ReadCaller;
 
 impl ReadCaller {
+    #[allow(clippy::too_many_arguments)]
     pub async fn call<'this, RW, Rng>(
         &self,
         _auto: (),
@@ -69,6 +72,7 @@ impl ReadCaller {
         read_state: &'this mut ReadState<'_>,
         _write_state: &mut WriteState<'_>,
         fragments_state: &'this mut FragmentsState<'_>,
+        _state: &mut ConnectionState,
     ) -> Option<Result<Option<Message<'this>>, Error<RW::Error>>>
     where
         RW: Read,
@@ -96,6 +100,10 @@ where
     RW: Write,
     Rng: RngCore,
 {
+    if state.closed {
+        return Err(Error::Write(WriteError::ConnectionClosed));
+    }
+
     state.closed = message.is_close();
 
     framez::functions::send(write_state, codec, inner, message)
@@ -109,13 +117,18 @@ pub async fn send_fragmented<RW, Rng>(
     codec: &mut FramesCodec<Rng>,
     inner: &mut RW,
     write_state: &mut WriteState<'_>,
+    state: &mut ConnectionState,
     message: Message<'_>,
     fragment_size: usize,
 ) -> Result<(), Error<RW::Error>>
 where
     RW: Write,
     Rng: RngCore,
 {
+    if state.closed {
+        return Err(Error::Write(WriteError::ConnectionClosed));
+    }
+
     for frame in message
         .fragments(fragment_size)
         .map_err(Error::Fragmentation)?

src/macros.rs

@@ -11,6 +11,7 @@ macro_rules! next {
                     &mut $websocketz.core.framed.core.state.read,
                     &mut $websocketz.core.framed.core.state.write,
                     &mut $websocketz.core.fragments_state,
+                    &mut $websocketz.core.state,
                 )
                 .await
             {
@@ -44,6 +45,7 @@ macro_rules! send_fragmented {
             &mut $websocketz.core.framed.core.codec,
             &mut $websocketz.core.framed.core.inner,
             &mut $websocketz.core.framed.core.state.write,
+            &mut $websocketz.core.state,
             $message,
             $fragment_size,
         )

src/tests.rs

@@ -956,7 +956,11 @@ mod fragmentation {
 }
 
 mod auto {
-    use crate::CloseFrame;
+
+    use crate::{
+        CloseFrame,
+        error::{Error, WriteError},
+    };
 
     use super::*;
 
@@ -1048,6 +1052,17 @@ mod auto {
 
             // Ensure the connection is closed
             assert!(next!(websocketz).is_none());
+
+            // Attempt to send another message after close should fail
+            match websocketz.send(Message::Text("test")).await {
+                Ok(_) => panic!("Expected error after close, but got Ok"),
+                Err(error) => {
+                    assert!(matches!(error, Error::Write(WriteError::ConnectionClosed)));
+                }
+            }
+
+            // Reading after eof should always return None
+            assert!(next!(websocketz).is_none());
         };
 
         let server = async move {
@@ -1064,6 +1079,14 @@ mod auto {
             );
 
             while next!(websocketz).is_some() {}
+
+            // Attempt to send another message after close should fail
+            match websocketz.send(Message::Text("test")).await {
+                Ok(_) => panic!("Expected error after close, but got Ok"),
+                Err(error) => {
+                    assert!(matches!(error, Error::Write(WriteError::ConnectionClosed)));
+                }
+            }
         };
 
         tokio::join!(server, client);

src/websocket_core.rs

@@ -66,20 +66,14 @@ 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 or the other side sends a close frame.
     ///
-    /// If the connection is closed, every read will return `None` and every write will return a [`WriteError::ConnectionClosed`].
+    /// If the connection is closed, every write will return a [`WriteError::ConnectionClosed`].
     pub closed: bool,
     /// Auto handling of ping/pong and close frames.
     auto: Auto,
 }
 
-// TODO: Set ConnectionState.closed to true if the user sends a close frame or the other side sends a close frame.
-// TODO: If ConnectionState.closed: Every read will then return (None, means connection closed) and every write will return a write error with ConnectionClosed.
-// TODO: 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)]
@@ -632,6 +626,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,
             fragment_size,
         )