diff --git a/src/utils/client.rs b/src/utils/client.rs index 5bfc79f..907738e 100644 --- a/src/utils/client.rs +++ b/src/utils/client.rs @@ -2,7 +2,6 @@ use std::{ hash::{Hash, Hasher}, io::{self, Read, Write}, net::{SocketAddr, TcpStream}, - time::Duration, }; use anyhow::anyhow; @@ -125,25 +124,11 @@ pub mod handshake { pub struct Client(TcpStream, Option, u64); impl Client { - /// Create a client with no timeouts + /// Create a client pub fn new(mut stream: TcpStream) -> crate::Result { handshake::handle_websocket_handshake(&mut stream)?; - Ok(Client(stream, None, rand::random())) - } - - /// 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)?; + 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())) } @@ -169,18 +154,20 @@ impl Client { Ok(()) } - /// Send a pong with given payload (control frames must be <=125) - fn send_pong(&self, payload: &[u8]) -> crate::Result<()> { + /// 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(()) + } - 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)?; + /// 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(()) } @@ -229,9 +216,14 @@ impl Client { // 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); // treat EOF as closed + return Ok(None); } return Err(e.into()); } @@ -287,9 +279,9 @@ impl Client { // Continuation / Text / Binary message_payload.extend(payload); if fin { - break; // got full message + break; } else { - continue; // wait for more fragments + continue; } } 0x8 => { @@ -309,12 +301,10 @@ impl Client { return Ok(None); } 0x9 => { - // Ping → respond with Pong - let _ = self.send_pong(&payload); + self.send_pong()?; continue; } 0xA => { - // Pong → ignore continue; } _ => {