diff --git a/crates/librqbit/src/peer_connection.rs b/crates/librqbit/src/peer_connection.rs index 127966d..c588135 100644 --- a/crates/librqbit/src/peer_connection.rs +++ b/crates/librqbit/src/peer_connection.rs @@ -347,23 +347,21 @@ impl PeerConnection { let reader = async move { loop { - read_buf - .read_message(&mut read_half, rwtimeout, |message| { - trace!("received: {:?}", &message); - - if let Message::Extended(ExtendedMessage::Handshake(h)) = &message { - *extended_handshake_ref.write() = Some(h.clone_to_owned()); - self.handler.on_extended_handshake(h)?; - trace!("remembered extended handshake for future serializing"); - } else { - self.handler - .on_received_message(message) - .context("error in handler.on_received_message()")?; - } - Ok(()) - }) + let message = read_buf + .read_message(&mut read_half, rwtimeout) .await .context("error reading message")?; + trace!("received: {:?}", &message); + + if let Message::Extended(ExtendedMessage::Handshake(h)) = &message { + *extended_handshake_ref.write() = Some(h.clone_to_owned()); + self.handler.on_extended_handshake(h)?; + trace!("remembered extended handshake for future serializing"); + } else { + self.handler + .on_received_message(message) + .context("error in handler.on_received_message()")?; + } } // For type inference. diff --git a/crates/librqbit/src/read_buf.rs b/crates/librqbit/src/read_buf.rs index 5867db7..71c6e66 100644 --- a/crates/librqbit/src/read_buf.rs +++ b/crates/librqbit/src/read_buf.rs @@ -68,18 +68,17 @@ impl ReadBuf { &mut self, mut conn: impl AsyncReadExt + Unpin, timeout: Duration, - on_message: impl for<'a> FnOnce(MessageBorrowed<'a>) -> anyhow::Result<()>, - ) -> anyhow::Result<()> { + ) -> anyhow::Result> { loop { let need_additional_bytes = match MessageBorrowed::deserialize(&self.buf[self.processed..self.filled]) { Err(MessageDeserializeError::NotEnoughData(d, _)) => d, Ok((msg, size)) => { self.processed += size; - // Rust's borrow checker can't do this early return. So we are using a callback instead. - // return Ok(msg); - on_message(msg)?; - return Ok(()); + + // Rust's borrow checker can't do this early return so resort to unsafe. + let msg = unsafe { std::mem::transmute(msg) }; + return Ok(msg); } Err(e) => return Err(e.into()), };