7c3eae0a7a
The API is now synchronous and any quinn stuff has been moved to another package. The quinn stuff will be slowly moved into moq-transport with generic traits.
137 lines
3.1 KiB
Rust
137 lines
3.1 KiB
Rust
use anyhow::Context;
|
|
use bytes::{Buf, Bytes, BytesMut};
|
|
use moq_transport::{Decode, DecodeError, Encode, Object};
|
|
use std::{io::Cursor, sync::Arc};
|
|
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
|
// TODO support clients
|
|
type WebTransportSession = h3_webtransport::server::WebTransportSession<h3_quinn::Connection, Bytes>;
|
|
|
|
// Reduce some typing
|
|
pub type SendStream = h3_webtransport::stream::SendStream<h3_quinn::SendStream<Bytes>, Bytes>;
|
|
pub type RecvStream = h3_webtransport::stream::RecvStream<h3_quinn::RecvStream, Bytes>;
|
|
|
|
pub struct Objects {
|
|
send: SendObjects,
|
|
recv: RecvObjects,
|
|
}
|
|
|
|
impl Objects {
|
|
pub fn new(session: Arc<WebTransportSession>) -> Self {
|
|
let send = SendObjects::new(session.clone());
|
|
let recv = RecvObjects::new(session);
|
|
Self { send, recv }
|
|
}
|
|
|
|
pub fn split(self) -> (SendObjects, RecvObjects) {
|
|
(self.send, self.recv)
|
|
}
|
|
|
|
pub async fn recv(&mut self) -> anyhow::Result<(Object, RecvStream)> {
|
|
self.recv.recv().await
|
|
}
|
|
|
|
pub async fn send(&mut self, header: Object) -> anyhow::Result<SendStream> {
|
|
self.send.send(header).await
|
|
}
|
|
}
|
|
|
|
pub struct SendObjects {
|
|
session: Arc<WebTransportSession>,
|
|
|
|
// A reusable buffer for encoding messages.
|
|
buf: BytesMut,
|
|
}
|
|
|
|
impl SendObjects {
|
|
pub fn new(session: Arc<WebTransportSession>) -> Self {
|
|
Self {
|
|
session,
|
|
buf: BytesMut::new(),
|
|
}
|
|
}
|
|
|
|
pub async fn send(&mut self, header: Object) -> anyhow::Result<SendStream> {
|
|
self.buf.clear();
|
|
header.encode(&mut self.buf).unwrap();
|
|
|
|
let mut stream = self
|
|
.session
|
|
.open_uni(self.session.session_id())
|
|
.await
|
|
.context("failed to open uni stream")?;
|
|
|
|
// TODO support select! without making a new stream.
|
|
stream.write_all(&self.buf).await?;
|
|
|
|
Ok(stream)
|
|
}
|
|
}
|
|
|
|
impl Clone for SendObjects {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
session: self.session.clone(),
|
|
buf: BytesMut::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
// Not clone, so we don't accidentally have two listners.
|
|
pub struct RecvObjects {
|
|
session: Arc<WebTransportSession>,
|
|
|
|
// A uni stream that's been accepted but not fully read from yet.
|
|
stream: Option<RecvStream>,
|
|
|
|
// Data that we've read but haven't formed a full message yet.
|
|
buf: BytesMut,
|
|
}
|
|
|
|
impl RecvObjects {
|
|
pub fn new(session: Arc<WebTransportSession>) -> Self {
|
|
Self {
|
|
session,
|
|
stream: None,
|
|
buf: BytesMut::new(),
|
|
}
|
|
}
|
|
|
|
pub async fn recv(&mut self) -> anyhow::Result<(Object, RecvStream)> {
|
|
// Make sure any state is saved across await boundaries so this works with select!
|
|
|
|
let stream = match self.stream.as_mut() {
|
|
Some(stream) => stream,
|
|
None => {
|
|
let (_session_id, stream) = self
|
|
.session
|
|
.accept_uni()
|
|
.await
|
|
.context("failed to accept uni stream")?
|
|
.context("no uni stream")?;
|
|
|
|
self.stream.insert(stream)
|
|
}
|
|
};
|
|
|
|
loop {
|
|
// Read the contents of the buffer
|
|
let mut peek = Cursor::new(&self.buf);
|
|
|
|
match Object::decode(&mut peek) {
|
|
Ok(header) => {
|
|
let stream = self.stream.take().unwrap();
|
|
self.buf.advance(peek.position() as usize);
|
|
return Ok((header, stream));
|
|
}
|
|
Err(DecodeError::UnexpectedEnd) => {
|
|
// The decode failed, so we need to append more data.
|
|
stream.read_buf(&mut self.buf).await?;
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
}
|
|
}
|