Working WebSocket

examples/client.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use session_rs::ws::WebSocket;
+use session_rs::{SessionFrame, ws::WebSocket};
 
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> session_rs::Result<()> {
@@ -11,13 +11,10 @@ async fn main() -> session_rs::Result<()> {
     tokio::spawn(async move {
         loop {
             match read_session.read().await {
-                Ok(Some((opcode, payload))) => {
+                Ok(SessionFrame::Text(text)) => {
-                    if opcode == 0x1 {
+                    println!("Server says: {}", text);
-                        let text = String::from_utf8(payload).unwrap_or_default();
-                        println!("Server says: {}", text);
-                    }
                 }
-                Ok(None) => {}
+                Ok(_) => {}
                 Err(_) => break,
             }
         }

examples/server.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 use tokio::net::TcpListener;
 
-use session_rs::session::WebSocket;
+use session_rs::{SessionFrame, ws::WebSocket};
 
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> session_rs::Result<()> {
@@ -26,21 +26,17 @@ async fn main() -> session_rs::Result<()> {
 
             // Read loop
             loop {
-                match session.read_frame().await {
+                match session.read().await {
-                    Ok(Some((opcode, payload))) => {
+                    Ok(SessionFrame::Text(text)) => {
-                        if opcode == 0x1 {
+                        println!("Received text: {}", text);
-                            // Text frame → parse JSON if possible
-                            let text = String::from_utf8(payload).unwrap_or_default();
-                            println!("Received text: {}", text);
 
-                            // Echo back
+                        // Echo back
-                            if let Err(e) = session.send(&serde_json::json!({"echo": text})).await {
+                        if let Err(e) = session.send(&serde_json::json!({"echo": text})).await {
-                                eprintln!("Send error: {:?}", e);
+                            eprintln!("Send error: {:?}", e);
-                                break;
+                            break;
-                            }
                         }
                     }
-                    Ok(None) => {}
+                    Ok(_) => {}
                     Err(e) => {
                         eprintln!("{e:?}");
                         break;

src/lib.rs

@@ -1,13 +1,15 @@
+use std::string::FromUtf8Error;
+
 pub mod server;
 pub mod session;
 pub mod ws;
 
-pub enum SessionFrame<T> {
+pub enum SessionFrame {
-    Typed(T),
+    Text(String),
     Binary(Vec<u8>),
     Ping,
     Pong,
-    Close
+    Close,
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -19,6 +21,7 @@ pub enum Error {
     InvalidFrame(String),
     HandshakeFailed(String),
     ConnectionClosed,
+    Utf8(FromUtf8Error),
 }
 
 impl From<std::io::Error> for Error {
@@ -32,3 +35,9 @@ impl From<serde_json::Error> for Error {
         Self::Json(value)
     }
 }
+
+impl From<FromUtf8Error> for Error {
+    fn from(value: FromUtf8Error) -> Self {
+        Self::Utf8(value)
+    }
+}

src/ws/mod.rs

@@ -171,8 +171,36 @@ impl WebSocket {
         Ok((fin, opcode, payload))
     }
 
-    pub async fn read<T>(&self) -> crate::Result<SessionFrame<T>> {
+    pub async fn read(&self) -> crate::Result<SessionFrame> {
-        let (fin, opcode, payload) = self.read_frame().await?;
+        let (fin, opcode, mut payload) = self.read_frame().await?;
+
+        if !fin {
+            // Continuation loop
+            while let (fin, o, mut p) = self.read_frame().await?
+                && !fin
+            {
+                match o {
+                    // Continuation
+                    0x0 => payload.append(&mut p),
+                    // Close
+                    0x8 => {
+                        self.close().await.ok();
+                    }
+                    // Ping
+                    0x9 => {
+                        self.send_pong().await.ok();
+                    }
+                    // Pong
+                    0xA => {}
+                    _ => {
+                        self.close().await.ok();
+                        return Err(crate::Error::InvalidFrame(format!(
+                            "Unknown opcode: {opcode}"
+                        )));
+                    }
+                }
+            }
+        }
 
         match opcode {
             // Close
@@ -190,16 +218,11 @@ impl WebSocket {
             // Pong
             0xA => Ok(SessionFrame::Pong),
 
-            // Continuation
-            0x0 => Ok(SessionFrame::Pong),
-
             // Text
-            // 0x1 => {
+            0x1 => Ok(SessionFrame::Text(String::from_utf8(payload)?)),
-
-            // },
 
             // Binary
-            0x2 => Ok(SessionFrame::Pong),
+            0x2 => Ok(SessionFrame::Binary(payload)),
 
             _ => {
                 self.close().await.ok();