Skip to content

Commit

Permalink
fix bad parsing
Browse files Browse the repository at this point in the history
  • Loading branch information
menghaoyu2002 committed Jun 6, 2024
1 parent cf47e1c commit 402e5bc
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions src/client/message.rs
Original file line number Diff line number Diff line change
Expand Up @@ -202,11 +202,10 @@ pub async fn send_message(stream: &TcpStream, message: &Message) -> Result<(), S

pub async fn receive_message(stream: &TcpStream) -> Result<Message, ReceiveError> {
let mut len = [0u8; 4];

let mut bytes_read = 0;
while bytes_read < 4 {
stream.readable().await.unwrap();
match stream.try_read(&mut len[bytes_read..]) {
match stream.try_read(&mut len) {
Ok(0) => {
return Err(ReceiveError::ReceiveError(ReceiveMessageError {
error: "stream was closed".to_string(),
Expand Down Expand Up @@ -234,7 +233,7 @@ pub async fn receive_message(stream: &TcpStream) -> Result<Message, ReceiveError
});
}

let mut message = vec![0u8; len as usize];
let mut message = Vec::new();
let mut bytes_read = 0;
while bytes_read < len as usize {
let mut buffer = vec![0u8; len as usize];
Expand All @@ -247,7 +246,7 @@ pub async fn receive_message(stream: &TcpStream) -> Result<Message, ReceiveError
}
Ok(n) => {
bytes_read += n;
message.extend_from_slice(&buffer[..n - 1]);
message.extend_from_slice(&buffer[..n]);
}
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
yield_now().await;
Expand Down

0 comments on commit 402e5bc

Please sign in to comment.