From 7940867c6a836f6580f2874eebb5d839f01d0b04 Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Wed, 18 Feb 2026 23:43:37 +0100
Subject: [PATCH] Refactoring

---
 src/lib.rs     |  3 +++
 src/session.rs | 51 +++++++++++++++++++++++++++++---------------------
 2 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0af8d9a..96c8850 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,6 +5,9 @@ pub mod session;
 pub enum SessionFrame<T> {
     Typed(T),
     Binary(Vec<u8>),
+    Ping,
+    Pong,
+    Close
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
diff --git a/src/session.rs b/src/session.rs
index 969ad35..6d96b2e 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -7,6 +7,8 @@ use tokio::{
     sync::Mutex,
 };
 
+use crate::SessionFrame;
+
 pub struct Session {
     pub(crate) reader: Arc<Mutex<tokio::net::tcp::OwnedReadHalf>>,
     pub(crate) writer: Arc<Mutex<tokio::net::tcp::OwnedWriteHalf>>,
@@ -119,14 +121,14 @@ impl Session {
 impl Session {
     /// Read a full WebSocket frame (handling masking and control frames)
     /// Returns (opcode, payload)
-    pub async fn read_frame(&self) -> crate::Result<Option<(u8, Vec<u8>)>> {
+    pub async fn read_frame(&self) -> crate::Result<(bool, u8, Vec<u8>)> {
         let mut reader = self.reader.lock().await;
 
         // --- 1. Read first 2-byte header ---
         let mut header = [0u8; 2];
         reader.read_exact(&mut header).await?;
 
-        // let fin = header[0] & 0x80 != 0;
+        let fin = header[0] & 0x80 != 0;
         let opcode = header[0] & 0x0F;
         let masked = header[1] & 0x80 != 0;
         let mut payload_len = (header[1] & 0x7F) as u64;
@@ -163,35 +165,42 @@ impl Session {
             }
         }
 
-        // --- 5. Handle control frames immediately ---
+        // --- 6. Return opcode + payload ---
+        Ok((fin, opcode, payload))
+    }
+
+    pub async fn read<T>(&self) -> crate::Result<SessionFrame<T>> {
+        let (fin, opcode, payload) = self.read_frame().await?;
+
         match opcode {
+            // Close
             0x8 => {
-                // Close
                 self.close().await.ok();
-                return Err(crate::Error::ConnectionClosed);
+                Ok(SessionFrame::Close)
             }
+
+            // Ping
             0x9 => {
-                // Ping
                 self.send_pong().await.ok();
-                return Ok(None);
-            }
-            0xA => {
-                // Pong, ignore
-                return Ok(None);
-            }
-            0x0 | 0x1 | 0x2 => {
-                // Continuation / Text / Binary → valid payload
+                Ok(SessionFrame::Ping)
             }
+
+            // Pong, ignore
+            0xA => Ok(SessionFrame::Pong),
+
+            // Continuation / Text / Binary → valid payload
+            0x0 => Ok(None),
+
+            0x1 => Ok(None),
+
+            0x2 => Ok(None),
+
             _ => {
                 self.close().await.ok();
-                return Err(crate::Error::InvalidFrame(format!(
-                    "Unknown opcode: {}",
-                    opcode
-                )));
+                Err(crate::Error::InvalidFrame(format!(
+                    "Unknown opcode: {opcode}"
+                )))
             }
         }
-
-        // --- 6. Return opcode + payload ---
-        Ok(Some((opcode, payload)))
     }
 }