43 lines
1008 B
Rust
43 lines
1008 B
Rust
|
use crate::object::Header;
|
||
|
|
|
||
|
|
use anyhow::Context;
|
||
|
|
|
||
|
|
use tokio::task::JoinSet;
|
||
|
|
|
||
|
|
use webtransport_generic::Session;
|
||
|
|
|
||
|
|
pub struct Receiver<S: Session> {
|
||
|
|
session: S,
|
||
|
|
|
||
|
|
// Streams that we've accepted but haven't read the header from yet.
|
||
|
|
streams: JoinSet<anyhow::Result<(Header, S::RecvStream)>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl<S: Session> Receiver<S> {
|
||
|
|
pub fn new(session: S) -> Self {
|
||
|
|
Self {
|
||
|
|
session,
|
||
|
|
streams: JoinSet::new(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn recv(&mut self) -> anyhow::Result<(Header, S::RecvStream)> {
|
||
|
|
loop {
|
||
|
|
tokio::select! {
|
||
|
|
res = self.session.accept_uni() => {
|
||
|
|
let stream = res.context("failed to accept stream")?;
|
||
|
|
self.streams.spawn(async move { Self::read(stream).await });
|
||
|
|
},
|
||
|
|
res = self.streams.join_next(), if !self.streams.is_empty() => {
|
||
|
|
return res.unwrap().context("failed to run join set")?;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn read(mut stream: S::RecvStream) -> anyhow::Result<(Header, S::RecvStream)> {
|
||
|
|
let header = Header::decode(&mut stream).await?;
|
||
|
|
Ok((header, stream))
|
||
|
|
}
|
||
|
|
}
|