From 95670962406270ff49b99aae3c47d27840f0f743 Mon Sep 17 00:00:00 2001
From: ivmarkov <ivan.markov@gmail.com>
Date: Sat, 20 Jan 2024 14:31:08 +0000
Subject: [PATCH] Small extensions to the WS protocol

---
 edge-http/src/io/client.rs |  8 +++++++
 edge-ws/src/lib.rs         | 47 ++++++++++++++++++++++++++++++--------
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/edge-http/src/io/client.rs b/edge-http/src/io/client.rs
index 054baee..95c4706 100644
--- a/edge-http/src/io/client.rs
+++ b/edge-http/src/io/client.rs
@@ -131,6 +131,14 @@ where
         Ok(self.io_mut())
     }
 
+    pub fn release(mut self) -> (T::Connection<'b>, &'b mut [u8]) {
+        let mut state = self.unbind();
+
+        let io = state.io.take().unwrap();
+
+        (io, state.buf)
+    }
+
     async fn start_request(
         &mut self,
         http11: bool,
diff --git a/edge-ws/src/lib.rs b/edge-ws/src/lib.rs
index 6189a73..f7b769d 100644
--- a/edge-ws/src/lib.rs
+++ b/edge-ws/src/lib.rs
@@ -41,6 +41,19 @@ impl FrameType {
     }
 }
 
+impl fmt::Display for FrameType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Text(fragmented) => write!(f, "Text{}", if *fragmented { " (fragmented)" } else { "" }),
+            Self::Binary(fragmented) => write!(f, "Binary{}", if *fragmented { " (fragmented)" } else { "" }),
+            Self::Ping => write!(f, "Ping"),
+            Self::Pong => write!(f, "Pong"),
+            Self::Close => write!(f, "Close"),
+            Self::Continue(ffinal) => write!(f, "Continue{}", if *ffinal { " (final)" } else { "" }),
+        }
+    }
+}
+
 #[derive(Clone, Eq, PartialEq, Debug)]
 pub enum Error<E> {
     Incomplete(usize),
@@ -159,17 +172,21 @@ impl FrameHeader {
                 None
             };
 
+            let frame_type = match opcode {
+                0 => FrameType::Continue(final_frame),
+                1 => FrameType::Text(!final_frame),
+                2 => FrameType::Binary(!final_frame),
+                8 => FrameType::Close,
+                9 => FrameType::Ping,
+                10 => FrameType::Pong,
+                _ => unreachable!(),
+            };
+
             let frame_header = FrameHeader {
-                frame_type: match opcode {
-                    0 => FrameType::Continue(final_frame),
-                    1 => FrameType::Text(!final_frame),
-                    2 => FrameType::Binary(!final_frame),
-                    8 => FrameType::Close,
-                    9 => FrameType::Ping,
-                    10 => FrameType::Pong,
-                    _ => unreachable!(),
-                },
-                payload_len,
+                frame_type,
+                payload_len: matches!(frame_type, FrameType::Binary(_) | FrameType::Text(_) | FrameType::Continue(_))
+                    .then(|| payload_len)
+                    .unwrap_or(0),
                 mask_key,
             };
 
@@ -266,6 +283,16 @@ impl FrameHeader {
     }
 }
 
+impl fmt::Display for FrameHeader {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Frame {{ {}, payload len {}, mask {:?} }}",
+            self.frame_type, self.payload_len, self.mask_key
+        )
+    }
+}
+
 #[cfg(feature = "embedded-svc")]
 mod embedded_svc_compat {
     use core::convert::TryFrom;