From 46b7a60dd430bb7e340fceb5715f78f3efe19b56 Mon Sep 17 00:00:00 2001 From: Leo dev Date: Sun, 14 Sep 2025 15:48:29 +0200 Subject: [PATCH] Updated client.rs with better prod --- src/utils/client.rs | 381 +++++++++++++++++++++++++++++++------------- 1 file changed, 266 insertions(+), 115 deletions(-) diff --git a/src/utils/client.rs b/src/utils/client.rs index d20988f..4f1331f 100644 --- a/src/utils/client.rs +++ b/src/utils/client.rs @@ -1,7 +1,8 @@ use std::{ hash::{Hash, Hasher}, - io::{Read, Write}, + io::{self, Read, Write}, net::TcpStream, + time::Duration, }; use anyhow::anyhow; @@ -13,54 +14,303 @@ pub mod handshake { use base64::Engine; use base64::engine::general_purpose::STANDARD as Base64; use sha1::{Digest, Sha1}; - use std::io::{Read, Write}; + 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 buffer = [0; 1024]; - let size = stream.read(&mut buffer)?; - let request = String::from_utf8_lossy(&buffer[..size]); + let mut reader = BufReader::new(stream.try_clone()?); + let mut request_line = String::new(); + reader.read_line(&mut request_line)?; - let key_line = request - .lines() - .find(|line| line.to_lowercase().starts_with("sec-websocket-key")) - .ok_or_else(|| { - std::io::Error::new(std::io::ErrorKind::InvalidData, "Missing Sec-WebSocket-Key") - })?; + if !request_line.starts_with("GET") { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Invalid HTTP method", + )); + } - let key = key_line.splitn(2, ':').nth(1).unwrap().trim(); + 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()); + } + } + let key = headers.get("sec-websocket-key").ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "Missing Sec-WebSocket-Key") + })?; + + if headers.get("upgrade").map(|v| v.to_lowercase()) != Some("websocket".to_string()) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Missing or invalid Upgrade header", + )); + } + + 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", + )); + } + + // Optional: validate Sec-WebSocket-Version == 13 (most common) + if let Some(ver) = headers.get("sec-websocket-version") { + if ver.trim() != "13" { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Unsupported Sec-WebSocket-Version", + )); + } + } + + // 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); + // Note: include Sec-WebSocket-Protocol handling if you support subprotocols 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", + Sec-WebSocket-Accept: {}\r\n\ + Sec-WebSocket-Version: 13\r\n\r\n", accept_key ); stream.write_all(response.as_bytes())?; stream.flush()?; - Ok(()) } } -pub struct Client(TcpStream); +pub struct Client(pub TcpStream); impl Client { + /// Create a client with no timeouts pub fn new(mut stream: TcpStream) -> crate::Result { handshake::handle_websocket_handshake(&mut stream)?; Ok(Client(stream)) } + + /// Create a client and set read/write timeouts (useful in prod) + pub fn with_timeouts( + mut stream: TcpStream, + read_timeout: Option, + write_timeout: Option, + ) -> crate::Result { + if let Some(t) = read_timeout { + stream.set_read_timeout(Some(t))?; + } + if let Some(t) = write_timeout { + stream.set_write_timeout(Some(t))?; + } + handshake::handle_websocket_handshake(&mut stream)?; + Ok(Client(stream)) + } + + /// Send a close frame and flush. `code` is a WebSocket close code (e.g., 1000 normal). + pub fn send_close(&self, code: u16, reason: &str) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + + // control frames must be <= 125 bytes + let mut payload = Vec::new(); + payload.extend_from_slice(&code.to_be_bytes()); + payload.extend_from_slice(reason.as_bytes()); + if payload.len() > 125 { + return Err(anyhow!("close reason too long").into()); + } + + let mut frame = Vec::with_capacity(2 + payload.len()); + frame.push(0x88); // FIN=1, opcode=0x8 (Close) + frame.push(payload.len() as u8); // server->client MUST NOT mask + frame.extend_from_slice(&payload); + + stream.write_all(&frame)?; + stream.flush()?; + Ok(()) + } + + /// Send a pong with given payload (control frames must be <=125) + fn send_pong(&self, payload: &[u8]) -> crate::Result<()> { + let mut stream = self.0.try_clone()?; + + if payload.len() > 125 { + return Err(anyhow!("pong payload too long").into()); + } + let mut frame = Vec::with_capacity(2 + payload.len()); + frame.push(0x8A); // FIN=1, opcode=0xA (Pong) + frame.push(payload.len() as u8); + frame.extend_from_slice(payload); + stream.write_all(&frame)?; + 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(()) + } + + /// 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(&self) -> crate::Result>> { + let mut stream = self.0.try_clone()?; + + let mut message_payload = Vec::new(); + + loop { + // read the 2-byte header + let mut header = [0u8; 2]; + if let Err(e) = stream.read_exact(&mut header) { + if e.kind() == io::ErrorKind::UnexpectedEof || e.kind() == io::ErrorKind::BrokenPipe + { + return Ok(None); // treat EOF as closed + } + 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 lengths + 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 (client→server MUST be masked) + let mut mask = [0u8; 4]; + if masked { + stream.read_exact(&mut mask)?; + } else { + let _ = self.send_close(1002, "Client frames must be masked"); + return Ok(None); + } + + // Control frame checks + if matches!(opcode, 0x8 | 0x9 | 0xA) { + if payload_len > 125 { + let _ = self.send_close(1002, "Control frame too large"); + return Ok(None); + } + if !fin { + let _ = self.send_close(1002, "Control frames must not be fragmented"); + return Ok(None); + } + } + + // Read payload + unmask + 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 | 0x1 | 0x2 => { + // Continuation / Text / Binary + message_payload.extend(payload); + if fin { + break; // got full message + } else { + continue; // wait for more fragments + } + } + 0x8 => { + // Close + let (code, reason) = if payload.len() >= 2 { + let code = u16::from_be_bytes([payload[0], payload[1]]); + let reason = if payload.len() > 2 { + String::from_utf8_lossy(&payload[2..]).into_owned() + } else { + String::new() + }; + (code, reason) + } else { + (1000, String::new()) + }; + let _ = self.send_close(code, &reason); + return Ok(None); + } + 0x9 => { + // Ping → respond with Pong + let _ = self.send_pong(&payload); + continue; + } + 0xA => { + // Pong → ignore + continue; + } + _ => { + let _ = self.send_close(1002, "Unsupported opcode"); + return Ok(None); + } + } + } + + // Try parsing JSON into ClientMessage + let message = match String::from_utf8(message_payload.clone()) { + Ok(text) => match serde_json::from_str(&text) { + Ok(msg) => WsMessage::Message(msg), + Err(_) => WsMessage::String(text), + }, + Err(_) => WsMessage::Binary(message_payload), + }; + + Ok(Some(message)) + } } impl Clone for Client { @@ -82,102 +332,3 @@ impl Hash for Client { self.0.peer_addr().unwrap().hash(state); } } - -impl Client { - /// Read a full WebSocket message, handling fragmentation (FIN) - pub fn read(&self) -> crate::Result>> { - let mut stream = &self.0; - let mut message_payload = Vec::new(); - let mut final_frame = false; - - while !final_frame { - let mut header = [0u8; 2]; - if stream.read_exact(&mut header).is_err() { - return Ok(None); // connection closed - } - - 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 lengths - 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 (client → server) - let mut mask = [0u8; 4]; - if masked { - stream.read_exact(&mut mask)?; - } - - // Read payload - let mut payload = vec![0u8; payload_len as usize]; - stream.read_exact(&mut payload)?; - - if masked { - for i in 0..payload.len() { - payload[i] ^= mask[i % 4]; - } - } - - match opcode { - 0x0 | 0x1 | 0x2 => { - // Continuation / Text / Binary - message_payload.extend(payload); - } - 0x8 => return Ok(None), // Close - 0x9 => continue, // Ping → ignore - 0xA => continue, // Pong → ignore - _ => return Err(anyhow!("Unsupported WebSocket opcode: {}", opcode).into()), - } - - final_frame = fin; - } - - // Try parsing JSON into ClientMessage - let message = match String::from_utf8(message_payload.clone()) { - Ok(text) => match serde_json::from_str(&text) { - Ok(msg) => WsMessage::Message(msg), - Err(_) => WsMessage::String(text), - }, - Err(_) => WsMessage::Binary(message_payload), - }; - - Ok(Some(message)) - } - - /// Send a JSON-serializable object as a WebSocket text frame - pub fn send(&self, m: T) -> crate::Result<()> { - let payload = serde_json::to_string(&m)?; - let payload_bytes = payload.as_bytes(); - - let mut stream = self.0.try_clone()?; - let mut header = Vec::new(); - header.push(0x81); // FIN=1, opcode=0x1 (text) - - let len = payload_bytes.len(); - 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(()) - } -}