diff --git a/src/quad_socket/client/tcp.rs b/src/quad_socket/client/tcp.rs index 7e6c8de..9ff3e7a 100644 --- a/src/quad_socket/client/tcp.rs +++ b/src/quad_socket/client/tcp.rs @@ -14,8 +14,8 @@ impl TcpSocket { pub fn send(&mut self, data: &[u8]) { use std::io::Write; - self.stream.write(&[data.len() as u8]).unwrap(); - self.stream.write(data).unwrap(); + self.stream.write_all(&(data.len() as u32).to_be_bytes()).unwrap(); + self.stream.write_all(data).unwrap(); } pub fn try_recv(&mut self) -> Option> { diff --git a/src/quad_socket/protocol.rs b/src/quad_socket/protocol.rs index cf59531..41a8959 100644 --- a/src/quad_socket/protocol.rs +++ b/src/quad_socket/protocol.rs @@ -12,26 +12,29 @@ impl MessageReader { } pub fn next(&mut self, mut stream: impl std::io::Read) -> Result>, ()> { - let mut bytes = [0 as u8; 255]; - match self { - MessageReader::Empty => match stream.read_exact(&mut bytes[0..1]) { - Ok(_) => { - *self = MessageReader::Amount(bytes[0] as usize); - Ok(None) + MessageReader::Empty => { + let mut size = [0u8; 4]; + match stream.read_exact(&mut size) { + Ok(_) => { + *self = MessageReader::Amount(u32::from_be_bytes(size) as usize); + Ok(None) + } + Err(err) if err.kind() == ErrorKind::WouldBlock => Ok(None), + Err(_err) => Err(()), } - Err(err) if err.kind() == ErrorKind::WouldBlock => Ok(None), - Err(_err) => Err(()), - }, - MessageReader::Amount(len) => match stream.read_exact(&mut bytes[0..*len]) { - Ok(_) => { - let msg = bytes[0..*len].to_vec(); - *self = MessageReader::Empty; - Ok(Some(msg)) + } + MessageReader::Amount(len) => { + let mut buf = vec![0u8; *len]; + match stream.read_exact(&mut buf) { + Ok(_) => { + *self = MessageReader::Empty; + Ok(Some(buf)) + } + Err(err) if err.kind() == ErrorKind::WouldBlock => Ok(None), + Err(_) => Err(()), } - Err(err) if err.kind() == ErrorKind::WouldBlock => Ok(None), - Err(_) => Err(()), - }, + } } } } diff --git a/src/quad_socket/server.rs b/src/quad_socket/server.rs index b276847..be28619 100644 --- a/src/quad_socket/server.rs +++ b/src/quad_socket/server.rs @@ -40,8 +40,8 @@ impl<'a> Sender<'a> { out.send(data).ok()?; } Sender::Tcp(stream) => { - stream.write(&[data.len() as u8]).ok()?; - stream.write(data).ok()?; + stream.write_all(&(data.len() as u32).to_be_bytes()).ok()?; + stream.write_all(data).ok()?; } }