120 lines
3.1 KiB
Rust
120 lines
3.1 KiB
Rust
use webtransport_generic::{SendStream, RecvStream};
|
|
use tokio::sync::mpsc;
|
|
|
|
use moq_transport::{Announce, AnnounceError, AnnounceOk, Message, Subscribe, SubscribeError, SubscribeOk, Control};
|
|
|
|
pub struct Main<S: SendStream, R: RecvStream> {
|
|
control: Control<S, R>,
|
|
outgoing: mpsc::Receiver<Message>,
|
|
|
|
contribute: mpsc::Sender<Contribute>,
|
|
distribute: mpsc::Sender<Distribute>,
|
|
}
|
|
|
|
impl<S: SendStream, R: RecvStream> Main <S, R> {
|
|
pub async fn run(mut self) -> anyhow::Result<()> {
|
|
loop {
|
|
tokio::select! {
|
|
Some(msg) = self.outgoing.recv() => self.control.send(msg).await?,
|
|
Ok(msg) = self.control.recv() => self.handle(msg).await?,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn handle(&mut self, msg: Message) -> anyhow::Result<()> {
|
|
match msg.try_into() {
|
|
Ok(msg) => self.contribute.send(msg).await?,
|
|
Err(msg) => match msg.try_into() {
|
|
Ok(msg) => self.distribute.send(msg).await?,
|
|
Err(msg) => anyhow::bail!("unsupported control message: {:?}", msg),
|
|
},
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
pub struct Component<T> {
|
|
incoming: mpsc::Receiver<T>,
|
|
outgoing: mpsc::Sender<Message>,
|
|
}
|
|
|
|
impl<T> Component<T> {
|
|
pub async fn send<M: Into<Message>>(&mut self, msg: M) -> anyhow::Result<()> {
|
|
self.outgoing.send(msg.into()).await?;
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn recv(&mut self) -> Option<T> {
|
|
self.incoming.recv().await
|
|
}
|
|
}
|
|
|
|
// Splits a control stream into two components, based on if it's a message for contribution or distribution.
|
|
pub fn split<S: SendStream, R: RecvStream>(control: Control<S, R>) -> (Main<S, R>, Component<Contribute>, Component<Distribute>) {
|
|
let (outgoing_tx, outgoing_rx) = mpsc::channel(1);
|
|
let (contribute_tx, contribute_rx) = mpsc::channel(1);
|
|
let (distribute_tx, distribute_rx) = mpsc::channel(1);
|
|
|
|
let control = Main {
|
|
control,
|
|
outgoing: outgoing_rx,
|
|
contribute: contribute_tx,
|
|
distribute: distribute_tx,
|
|
};
|
|
|
|
let contribute = Component {
|
|
incoming: contribute_rx,
|
|
outgoing: outgoing_tx.clone(),
|
|
};
|
|
|
|
let distribute = Component {
|
|
incoming: distribute_rx,
|
|
outgoing: outgoing_tx,
|
|
};
|
|
|
|
(control, contribute, distribute)
|
|
}
|
|
|
|
// Messages we expect to receive from the client for contribution.
|
|
#[derive(Debug)]
|
|
pub enum Contribute {
|
|
Announce(Announce),
|
|
SubscribeOk(SubscribeOk),
|
|
SubscribeError(SubscribeError),
|
|
}
|
|
|
|
impl TryFrom<Message> for Contribute {
|
|
type Error = Message;
|
|
|
|
fn try_from(msg: Message) -> Result<Self, Self::Error> {
|
|
match msg {
|
|
Message::Announce(msg) => Ok(Self::Announce(msg)),
|
|
Message::SubscribeOk(msg) => Ok(Self::SubscribeOk(msg)),
|
|
Message::SubscribeError(msg) => Ok(Self::SubscribeError(msg)),
|
|
_ => Err(msg),
|
|
}
|
|
}
|
|
}
|
|
|
|
// Messages we expect to receive from the client for distribution.
|
|
#[derive(Debug)]
|
|
pub enum Distribute {
|
|
AnnounceOk(AnnounceOk),
|
|
AnnounceError(AnnounceError),
|
|
Subscribe(Subscribe),
|
|
}
|
|
|
|
impl TryFrom<Message> for Distribute {
|
|
type Error = Message;
|
|
|
|
fn try_from(value: Message) -> Result<Self, Self::Error> {
|
|
match value {
|
|
Message::AnnounceOk(msg) => Ok(Self::AnnounceOk(msg)),
|
|
Message::AnnounceError(msg) => Ok(Self::AnnounceError(msg)),
|
|
Message::Subscribe(msg) => Ok(Self::Subscribe(msg)),
|
|
_ => Err(value),
|
|
}
|
|
}
|
|
}
|