132 lines
3.2 KiB
Rust
132 lines
3.2 KiB
Rust
use webtransport_generic::{RecvStream, SendStream};
|
|
use crate::{Decode, DecodeError, Encode, Message};
|
|
use crate::network::stream::{recv, send};
|
|
|
|
use bytes::{Buf, BytesMut};
|
|
|
|
use std::io::Cursor;
|
|
use std::sync::Arc;
|
|
use tokio::sync::Mutex;
|
|
|
|
use anyhow::Context;
|
|
|
|
|
|
pub struct Control<S: SendStream, R: RecvStream> {
|
|
sender: ControlSend<S>,
|
|
recver: ControlRecv<R>,
|
|
}
|
|
|
|
impl<S: SendStream, R: RecvStream> Control<S, R>{
|
|
pub(crate) fn new(sender: Box<S>, recver: Box<R>) -> Self {
|
|
let sender = ControlSend::new(sender);
|
|
let recver = ControlRecv::new(recver);
|
|
|
|
Self { sender, recver }
|
|
}
|
|
|
|
pub fn split(self) -> (ControlSend<S>, ControlRecv<R>) {
|
|
(self.sender, self.recver)
|
|
}
|
|
|
|
pub async fn send<T: Into<Message>>(&mut self, msg: T) -> anyhow::Result<()> {
|
|
self.sender.send(msg)
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("{:?}", e))
|
|
.context("error sending control message")
|
|
}
|
|
|
|
pub async fn recv(&mut self) -> anyhow::Result<Message> {
|
|
self.recver.recv().await
|
|
}
|
|
}
|
|
|
|
pub struct ControlSend<S: SendStream> {
|
|
stream: Box<S>,
|
|
buf: BytesMut, // reuse a buffer to encode messages.
|
|
}
|
|
|
|
impl<S: SendStream> ControlSend<S> {
|
|
pub fn new(inner: Box<S>) -> Self {
|
|
Self {
|
|
buf: BytesMut::new(),
|
|
stream: inner,
|
|
}
|
|
}
|
|
|
|
pub async fn send<T: Into<Message>>(&mut self, msg: T) -> anyhow::Result<()> {
|
|
let msg = msg.into();
|
|
log::info!("sending message: {:?}", msg);
|
|
|
|
self.buf.clear();
|
|
msg.encode(&mut self.buf)?;
|
|
|
|
// TODO make this work with select!
|
|
send(self.stream.as_mut(), &mut self.buf)
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("{:?}", e.into()))
|
|
.context("error sending control message")?;
|
|
Ok(())
|
|
}
|
|
|
|
// Helper that lets multiple threads send control messages.
|
|
pub fn share(self) -> ControlShared<S> {
|
|
ControlShared {
|
|
stream: Arc::new(Mutex::new(self)),
|
|
}
|
|
}
|
|
}
|
|
|
|
// Helper that allows multiple threads to send control messages.
|
|
// There's no equivalent for receiving since only one thread should be receiving at a time.
|
|
#[derive(Clone)]
|
|
pub struct ControlShared<S: SendStream> {
|
|
stream: Arc<Mutex<ControlSend<S>>>,
|
|
}
|
|
|
|
impl<S: SendStream> ControlShared<S> {
|
|
pub async fn send<T: Into<Message>>(&mut self, msg: T) -> anyhow::Result<()> {
|
|
let mut stream = self.stream.lock().await;
|
|
stream.send(msg).await
|
|
}
|
|
}
|
|
|
|
pub struct ControlRecv<R: RecvStream> {
|
|
stream: Box<R>,
|
|
buf: BytesMut, // data we've read but haven't fully decoded yet
|
|
}
|
|
|
|
impl<R: RecvStream> ControlRecv<R> {
|
|
pub fn new(inner: Box<R>) -> Self {
|
|
Self {
|
|
buf: BytesMut::new(),
|
|
stream: inner,
|
|
}
|
|
}
|
|
|
|
// 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);
|
|
|
|
log::info!("received message: {:?}", msg);
|
|
return Ok(msg);
|
|
}
|
|
Err(DecodeError::UnexpectedEnd) => {
|
|
// The decode failed, so we need to append more data.
|
|
recv(self.stream.as_mut(), &mut self.buf)
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("{:?}", e.into()))
|
|
.context("error receiving control message")?;
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|
|
}
|