dspk_stream/protocol/
codec.rs

1//! Length-prefixed messages on QUIC streams: a [`LENGTH_PREFIX_BYTES`]-byte big-endian length,
2//! then the protobuf encoding.
3
4use prost::Message;
5
6use super::LENGTH_PREFIX_BYTES;
7use crate::error::{StreamError, StreamResult};
8
9/// Writes `message` with its length prefix.
10pub async fn write_message(stream: &mut quinn::SendStream, message: &impl Message) -> StreamResult<()> {
11    let len = message.encoded_len();
12    let prefix = u32::try_from(len).map_err(|_| StreamError::TooLarge { len, max: u32::MAX as usize })?;
13    let mut buf = Vec::with_capacity(LENGTH_PREFIX_BYTES + len);
14    buf.extend_from_slice(&prefix.to_be_bytes());
15    message.encode(&mut buf).expect("the buffer holds the encoded length");
16    stream.write_all(&buf).await?;
17    Ok(())
18}
19
20/// Reads the next message. `Ok(None)` when the stream ended cleanly before a message; an error
21/// when it ended inside one or the length exceeds `max_bytes` (checked before allocating).
22pub async fn read_message<M: Message + Default>(stream: &mut quinn::RecvStream, max_bytes: usize) -> StreamResult<Option<M>> {
23    let mut prefix = [0u8; LENGTH_PREFIX_BYTES];
24    match stream.read_exact(&mut prefix).await {
25        Ok(()) => {}
26        Err(quinn::ReadExactError::FinishedEarly(0)) => return Ok(None),
27        Err(quinn::ReadExactError::FinishedEarly(_)) => return Err(StreamError::Truncated),
28        Err(e) => return Err(e.into()),
29    }
30    let len = u32::from_be_bytes(prefix) as usize;
31    if len > max_bytes {
32        return Err(StreamError::TooLarge { len, max: max_bytes });
33    }
34    let mut buf = vec![0u8; len];
35    match stream.read_exact(&mut buf).await {
36        Ok(()) => Ok(Some(M::decode(buf.as_slice())?)),
37        Err(quinn::ReadExactError::FinishedEarly(_)) => Err(StreamError::Truncated),
38        Err(e) => Err(e.into()),
39    }
40}
dspk_stream/protocol/codec.rs

dspk_stream/protocol/
codec.rs
