diff --git a/Cargo.lock b/Cargo.lock index 45f11d5..a7b9a42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,491 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "anyhow" +version = "1.0.101" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "rand_core", + "wasip2", + "wasip3", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "indexmap" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +dependencies = [ + "equivalent", + "hashbrown 0.16.1", + "serde", + "serde_core", +] + +[[package]] +name = "itoa" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" + +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + +[[package]] +name = "libc" +version = "0.2.182" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" +dependencies = [ + "chacha20", + "getrandom", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" + +[[package]] +name = "semver" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + [[package]] name = "session-rs" version = "0.1.0" +dependencies = [ + "base64", + "rand", + "serde", + "serde_json", + "sha1", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest", +] + +[[package]] +name = "syn" +version = "2.0.116" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3df424c70518695237746f84cede799c9c58fcb37450d7b23716568cc8bc69cb" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.2+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/Cargo.toml b/Cargo.toml index 2f38e0f..fd21142 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,3 +4,8 @@ version = "0.1.0" edition = "2024" [dependencies] +base64 = "0.22.1" +rand = "0.10.0" +serde = "1.0.228" +serde_json = "1.0.149" +sha1 = "0.10.6" diff --git a/src/lib.rs b/src/lib.rs index 1543da2..a149d19 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,2 +1,26 @@ -pub mod session; pub mod server; +pub mod session; + +pub enum SessionMessage { + SessionMessage(T), + Binary(Vec), +} + +pub type Result = std::result::Result; + +pub enum Error { + Io(std::io::Error), + Json(serde_json::Error), +} + +impl From for Error { + fn from(value: std::io::Error) -> Self { + Self::Io(value) + } +} + +impl From for Error { + fn from(value: serde_json::Error) -> Self { + Self::Json(value) + } +} diff --git a/src/session.rs b/src/session.rs index 704d549..e9a9242 100644 --- a/src/session.rs +++ b/src/session.rs @@ -1 +1,398 @@ -pub struct Session(); +use std::{ + hash::{Hash, Hasher}, + io::{self, Read, Write}, + net::TcpStream, +}; + +use serde::{Deserialize, Serialize}; + +use crate::SessionMessage; + +pub mod handshake { + use base64::Engine; + use base64::engine::general_purpose::STANDARD as Base64; + use sha1::{Digest, Sha1}; + use std::collections::HashMap; + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpStream; + + const WS_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; + + pub fn handle_websocket_handshake(stream: &mut TcpStream) -> std::io::Result<()> { + let mut reader = BufReader::new(stream.try_clone()?); + let mut request_line = String::new(); + reader.read_line(&mut request_line)?; + + // Trim CRLF to make sure comparisons are clean + let request_line = request_line.trim_end(); + + // Allow HEAD (used by Render for health checks) + if request_line.starts_with("HEAD") { + let response = "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n"; + stream.write_all(response.as_bytes())?; + stream.flush()?; + return Ok(()); + } + + // Only proceed if it’s a GET + if !request_line.starts_with("GET") { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid HTTP method: {request_line}"), + )); + } + + // Read headers + let mut headers = HashMap::new(); + let mut line = String::new(); + loop { + line.clear(); + let bytes = reader.read_line(&mut line)?; + if bytes == 0 || line == "\r\n" { + break; + } + if let Some((k, v)) = line.split_once(':') { + headers.insert(k.trim().to_lowercase(), v.trim().to_string()); + } + } + + // Check if it's actually a WebSocket upgrade request + let is_websocket_upgrade = headers + .get("upgrade") + .map(|v| v.eq_ignore_ascii_case("websocket")) + .unwrap_or(false); + + if !is_websocket_upgrade { + // Not a WebSocket request — probably a normal HTTP GET (e.g. health check) + let response = + "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\n\r\nOK"; + stream.write_all(response.as_bytes())?; + stream.flush()?; + return Ok(()); + } + + // Validate "Connection: Upgrade" + if !headers + .get("connection") + .map(|v| v.to_lowercase().contains("upgrade")) + .unwrap_or(false) + { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Missing or invalid Connection header", + )); + } + + // Validate WebSocket key + let key = headers.get("sec-websocket-key").ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "Missing Sec-WebSocket-Key") + })?; + + // Validate version + if let Some(ver) = headers.get("sec-websocket-version") { + if ver.trim() != "13" { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Unsupported Sec-WebSocket-Version: {}", ver), + )); + } + } + + // Compute accept key + let mut hasher = Sha1::new(); + hasher.update(key.as_bytes()); + hasher.update(WS_GUID.as_bytes()); + let hash = hasher.finalize(); + let accept_key = Base64.encode(hash); + + // Send response + let response = format!( + "HTTP/1.1 101 Switching Protocols\r\n\ + Upgrade: websocket\r\n\ + Connection: Upgrade\r\n\ + Sec-WebSocket-Accept: {}\r\n\r\n", + accept_key + ); + + stream.write_all(response.as_bytes())?; + stream.flush()?; + Ok(()) + } +} + +pub struct Session(TcpStream, u64); + +impl Session { + /// Create a client + pub fn new(mut stream: TcpStream) -> crate::Result { + handshake::handle_websocket_handshake(&mut stream)?; + stream.set_read_timeout(Some(std::time::Duration::from_secs(10)))?; + stream.set_write_timeout(Some(std::time::Duration::from_secs(10)))?; + Ok(Session(stream, rand::random())) + } + + /// Send a close frame and flush. + pub fn send_close(&self) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + stream.write_all(&[0x88])?; + stream.flush()?; + Ok(()) + } + + /// Send a ping (no payload) + fn send_ping(&self) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + // FIN + opcode (ping = 0x89), payload length = 0x00 + stream.write_all(&[0x89, 0x00])?; + stream.flush()?; + Ok(()) + } + + /// Send a pong (no payload) + fn send_pong(&self) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + // FIN + opcode (pong = 0x8A), payload length = 0x00 + stream.write_all(&[0x8A, 0x00])?; + stream.flush()?; + Ok(()) + } + + /// Send a text/binary frame (server->client must NOT mask) + pub fn send(&self, m: T) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + + let payload = serde_json::to_string(&m)?; + let payload_bytes = payload.as_bytes(); + let len = payload_bytes.len(); + + let mut header = Vec::new(); + header.push(0x81); // FIN=1, opcode=0x1 (text) + + if len < 126 { + header.push(len as u8); + } else if len <= 65535 { + 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()); + } + + stream.write_all(&header)?; + stream.write_all(payload_bytes)?; + stream.flush()?; + Ok(()) + } + + /// Send a binary WebSocket frame (server -> client) + pub fn send_bin(&self, payload: &[u8]) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + + let mut header = Vec::with_capacity(10); + + // FIN=1, opcode=2 (binary) + header.push(0x82); + + let len = payload.len(); + + if len < 126 { + header.push(len as u8); // mask bit = 0 + } 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()); + } + + stream.write_all(&header)?; + stream.write_all(payload)?; + stream.flush()?; + + Ok(()) + } + + /// Read a full WebSocket message, handling fragmentation and control frames. + /// + /// Returns: + /// - Ok(Some(WsMessage)) on an application message (text/binary) + /// - Ok(None) if the connection should be closed (close received / read EOF) + /// - Err on protocol or IO errors. + pub fn read_t Deserialize<'de>>( + &self, + ) -> crate::Result>> { + let mut stream = self.0.try_clone()?; + + let mut message_payload = Vec::new(); + let mut expecting_continuation = false; + let mut message_type: Option = None; // 0x1 for text, 0x2 for binary + + loop { + // Read 2-byte header + let mut header = [0u8; 2]; + match stream.read_exact(&mut header) { + Ok(_) => {} + Err(e) => match e.kind() { + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut => { + self.send_ping()?; + continue; + } + io::ErrorKind::UnexpectedEof | io::ErrorKind::BrokenPipe => return Ok(None), + _ => return Err(e.into()), + }, + } + + 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; + + // Extended payload length + if payload_len == 126 { + let mut ext_len = [0u8; 2]; + stream.read_exact(&mut ext_len)?; + payload_len = u16::from_be_bytes(ext_len) as u64; + } else if payload_len == 127 { + let mut ext_len = [0u8; 8]; + stream.read_exact(&mut ext_len)?; + payload_len = u64::from_be_bytes(ext_len); + } + + // Mask key + let mut mask = [0u8; 4]; + if masked { + stream.read_exact(&mut mask)?; + } else { + let _ = self.send_close(); + return Ok(None); + } + + // Control frame checks + if matches!(opcode, 0x8 | 0x9 | 0xA) { + if payload_len > 125 { + let _ = self.send_close(); + return Ok(None); + } + if !fin { + let _ = self.send_close(); + return Ok(None); + } + } + + // Read payload + let mut payload = vec![0u8; payload_len as usize]; + if payload_len > 0 { + stream.read_exact(&mut payload)?; + for i in 0..payload.len() { + payload[i] ^= mask[i % 4]; + } + } + + match opcode { + 0x0 => { + // Continuation + if !expecting_continuation { + let _ = self.send_close(); + return Ok(None); + } + message_payload.extend(payload); + if fin { + break; + } + } + 0x1 => { + // Text + if expecting_continuation { + let _ = self.send_close(); + return Ok(None); + } + message_payload.extend(payload); + message_type = Some(0x1); + if fin { + break; + } else { + expecting_continuation = true; + } + } + 0x2 => { + // Binary + if expecting_continuation { + let _ = self.send_close(); + return Ok(None); + } + message_payload.extend(payload); + message_type = Some(0x2); + if fin { + break; + } else { + expecting_continuation = true; + } + } + 0x8 => { + // Close + let _ = self.send_close(); + return Ok(None); + } + 0x9 => { + // Ping + self.send_pong()?; + continue; + } + 0xA => { + // Pong + continue; + } + _ => { + let _ = self.send_close(); + return Ok(None); + } + } + } + + // Convert payload into proper message type + let message = match message_type { + Some(0x1) => { + // Text frame → try JSON, otherwise keep text + match String::from_utf8(message_payload.clone()) { + Ok(text) => match serde_json::from_str(&text) { + Ok(msg) => SessionMessage::SessionMessage(msg), + Err(e) => return Err(crate::Error::Json(e)), + }, + Err(_) => SessionMessage::Binary(message_payload), + } + } + Some(0x2) => SessionMessage::Binary(message_payload), + _ => return Ok(None), // Should not happen + }; + + Ok(Some(message)) + } + + pub fn close(&self) -> crate::Result<()> { + self.0.shutdown(std::net::Shutdown::Both)?; + Ok(()) + } +} + +impl Clone for Session { + fn clone(&self) -> Self { + Session( + self.0.try_clone().expect("failed to clone TcpStream"), + self.1.clone(), + ) + } +} + +impl PartialEq for Session { + fn eq(&self, other: &Self) -> bool { + self.1 == other.1 + } +} + +impl Eq for Session {} + +impl Hash for Session { + fn hash(&self, state: &mut H) { + self.1.hash(state); + } +}