diff --git a/src/crypto.rs b/src/crypto.rs index 3754ea4..133dc78 100644 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -47,9 +47,9 @@ impl Crypto { if go_first { handler.send_raw(&msg).await?; - buf = handler.recv_raw().await?; + buf = handler.recv_raw(DH_PBK_SIZE).await?; } else { - buf = handler.recv_raw().await?; + buf = handler.recv_raw(DH_PBK_SIZE).await?; handler.send_raw(&msg).await?; } diff --git a/src/sockets.rs b/src/sockets.rs index e92649c..ad01e3e 100644 --- a/src/sockets.rs +++ b/src/sockets.rs @@ -65,7 +65,7 @@ impl<'a> SocketHandler<'a> { } pub async fn recv(&mut self) -> Result, Box> { - let mut buf = self.recv_raw().await?; + let mut buf = self.recv_raw(1).await?; buf.pop(); buf = general_purpose::STANDARD_NO_PAD.decode(&buf)?.to_vec(); @@ -77,14 +77,25 @@ impl<'a> SocketHandler<'a> { Ok(data) } - pub async fn recv_raw(&mut self) -> Result, Box> { + pub async fn recv_raw( + &mut self, + min_limit: usize, + ) -> Result, Box> { let mut buf = Vec::new(); - let n = self.reader.read_until(b':', &mut buf).await?; - if n == 0 { - return Err("Received 0 bytes from the socket".into()); + while buf.len() <= min_limit { + let n = self.reader.read_until(b':', &mut buf).await?; + + if n == 0 { + return Err("Received 0 bytes from the socket".into()); + } } + /* + TODO: use min_limit to check whether read_until has reached EOF before reading all the necessary bytes + (e.g. regarding ecdh key exchange) --> loop and read until buf.len() == min_limit + */ + debug!("Received {} bytes from the socket", buf.len()); Ok(buf)