async Read frame

src/session.rs

@@ -1,10 +1,12 @@
-use crate::SessionMessage;
-
 use std::{
     hash::{Hash, Hasher},
     sync::Arc,
 };
-use tokio::{io::AsyncWriteExt, net::TcpStream, sync::Mutex};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpStream,
+    sync::Mutex,
+};
 
 pub struct Session {
     reader: Arc<Mutex<tokio::net::tcp::OwnedReadHalf>>,
@@ -50,6 +52,30 @@ impl Hash for Session {
     }
 }
 
+impl Session {
+    async fn send_frame(&self, opcode: u8, payload: &[u8]) -> crate::Result<()> {
+        let mut writer = self.writer.lock().await;
+
+        let mut header = Vec::with_capacity(10);
+        header.push(0x80 | opcode);
+
+        let len = payload.len();
+        if len < 126 {
+            header.push(len as u8);
+        } else if len <= 0xFFFF {
+            header.push(126);
+            header.extend_from_slice(&(len as u16).to_be_bytes());
+        } else {
+            header.push(127);
+            header.extend_from_slice(&(len as u64).to_be_bytes());
+        }
+
+        writer.write_all(&header).await?;
+        writer.write_all(payload).await?;
+        Ok(())
+    }
+}
+
 impl Session {
     pub async fn send<T: serde::Serialize>(&self, msg: &T) -> crate::Result<()> {
         let payload = serde_json::to_vec(msg)?;
@@ -88,25 +114,94 @@ impl Session {
         });
     }
 
-    async fn send_frame(&self, opcode: u8, payload: &[u8]) -> crate::Result<()> {
+    /// Send a close frame and flush
+    pub async fn close(&self) -> crate::Result<()> {
         let mut writer = self.writer.lock().await;
 
-        let mut header = Vec::with_capacity(10);
-        header.push(0x80 | opcode);
-
-        let len = payload.len();
-        if len < 126 {
-            header.push(len as u8);
-        } else if len <= 0xFFFF {
-            header.push(126);
-            header.extend_from_slice(&(len as u16).to_be_bytes());
-        } else {
-            header.push(127);
-            header.extend_from_slice(&(len as u64).to_be_bytes());
-        }
-
-        writer.write_all(&header).await?;
-        writer.write_all(payload).await?;
+        // FIN=1, opcode=0x8 (close), payload length=0
+        let frame = [0x88, 0x00];
+        writer.write_all(&frame).await?;
+        writer.flush().await?;
         Ok(())
     }
 }
+
+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>)>> {
+        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 opcode = header[0] & 0x0F;
+        let masked = header[1] & 0x80 != 0;
+        let mut payload_len = (header[1] & 0x7F) as u64;
+
+        // --- 2. Read extended payload length if necessary ---
+        if payload_len == 126 {
+            let mut buf = [0u8; 2];
+            reader.read_exact(&mut buf).await?;
+            payload_len = u16::from_be_bytes(buf) as u64;
+        } else if payload_len == 127 {
+            let mut buf = [0u8; 8];
+            reader.read_exact(&mut buf).await?;
+            payload_len = u64::from_be_bytes(buf);
+        }
+
+        // --- 3. Read mask key ---
+        if !masked {
+            // Per spec, client-to-server frames MUST be masked
+            self.close().await.ok();
+            return Err(crate::Error::InvalidFrame(
+                "Received unmasked frame from client".into(),
+            ));
+        }
+
+        let mut mask = [0u8; 4];
+        reader.read_exact(&mut mask).await?;
+
+        // --- 4. Read payload ---
+        let mut payload = vec![0u8; payload_len as usize];
+        if payload_len > 0 {
+            reader.read_exact(&mut payload).await?;
+            for i in 0..payload.len() {
+                payload[i] ^= mask[i % 4];
+            }
+        }
+
+        // --- 5. Handle control frames immediately ---
+        match opcode {
+            0x8 => {
+                // Close
+                self.close().await.ok();
+                return Err(crate::Error::ConnectionClosed);
+            }
+            0x9 => {
+                // Ping
+                self.send_pong().await.ok();
+                return Ok(None);
+            }
+            0xA => {
+                // Pong, ignore
+                return Ok(None);
+            }
+            0x0 | 0x1 | 0x2 => {
+                // Continuation / Text / Binary → valid payload
+            }
+            _ => {
+                self.close().await.ok();
+                return Err(crate::Error::InvalidFrame(format!(
+                    "Unknown opcode: {}",
+                    opcode
+                )));
+            }
+        }
+
+        // --- 6. Return opcode + payload ---
+        Ok(Some((opcode, payload)))
+    }
+}