use std::{ hash::{Hash, Hasher}, io::{self, Read, Write}, net::{SocketAddr, TcpStream}, }; use anyhow::anyhow; use serde::{Deserialize, Serialize}; use crate::types::message::{ClientMessage, WsMessage}; 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 Client(TcpStream, Option, u64); impl Client { /// 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(Client(stream, None, rand::random())) } /// 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 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(()) } /// 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(); 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::WouldBlock || e.kind() == io::ErrorKind::TimedOut { self.send_ping()?; continue; } if e.kind() == io::ErrorKind::UnexpectedEof || e.kind() == 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 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; } else { continue; } } 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 => { self.send_pong()?; continue; } 0xA => { 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)) } /// 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>> { self.read_t() } pub fn get_uuid(&self) -> crate::Result { match &self.1 { Some(v) => Ok(v.clone()), None => Err(anyhow!("Client ({}) UUID not set", self.2).into()), } } pub fn set_uuid(&mut self, uuid: &str) { self.1 = Some(uuid.to_string()) } #[deprecated] pub fn addr(&self) -> crate::Result { Ok(self.0.peer_addr().unwrap_or(self.0.local_addr()?)) } } impl Clone for Client { fn clone(&self) -> Self { Client( self.0.try_clone().expect("failed to clone TcpStream"), self.1.clone(), self.2.clone(), ) } } impl PartialEq for Client { fn eq(&self, other: &Self) -> bool { self.2 == other.2 } } impl Eq for Client {} impl Hash for Client { fn hash(&self, state: &mut H) { self.2.hash(state); } }