50 lines
1.0 KiB
Rust
50 lines
1.0 KiB
Rust
use crate::coding::{Decode, DecodeError};
|
|
use crate::message::Message;
|
|
|
|
use bytes::{Buf, BytesMut};
|
|
|
|
use std::io::Cursor;
|
|
|
|
use webtransport_generic::AsyncRecvStream;
|
|
|
|
pub struct Receiver<R>
|
|
where
|
|
R: AsyncRecvStream, // TODO take RecvStream instead
|
|
{
|
|
stream: R,
|
|
buf: BytesMut, // data we've read but haven't fully decoded yet
|
|
}
|
|
|
|
impl<R> Receiver<R>
|
|
where
|
|
R: AsyncRecvStream,
|
|
{
|
|
pub fn new(stream: R) -> Self {
|
|
Self {
|
|
buf: BytesMut::new(),
|
|
stream,
|
|
}
|
|
}
|
|
|
|
// Read the next full message from the stream.
|
|
pub async fn recv(&mut self) -> anyhow::Result<Message> {
|
|
loop {
|
|
// Read the contents of the buffer
|
|
let mut peek = Cursor::new(&self.buf);
|
|
|
|
match Message::decode(&mut peek) {
|
|
Ok(msg) => {
|
|
// We've successfully decoded a message, so we can advance the buffer.
|
|
self.buf.advance(peek.position() as usize);
|
|
return Ok(msg);
|
|
}
|
|
Err(DecodeError::UnexpectedEnd) => {
|
|
// The decode failed, so we need to append more data.
|
|
self.stream.recv(&mut self.buf).await?;
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|
|
}
|