leyoda/rqbit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

read_buf: use unsafe but remove callback

Browse source
This commit is contained in:
Igor Katson 2024-05-03 12:38:24 +01:00
parent 01cbeabd51
commit e658b960e3
2 changed files with 18 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

28
crates/librqbit/src/peer_connection.rs
View file

@ -347,23 +347,21 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
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.

11
crates/librqbit/src/read_buf.rs
View file

@ -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<MessageBorrowed<'_>> {
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()),
};