open-standards/joule-code/joule-code-rs/crates/joule-moq-relay/src/wire.rs

//! Reading MOQT from the transport's streams: framed control messages on control and
//! request streams, and the items of data streams as their bytes arrive.

use crate::transport::{BoxRecv, BoxSend, TransportError};
use joule_moq::draft22::data::{Limits, UniStreamDecoder};
use joule_moq::draft22::{ControlMessage, GroupOrder, RequestKind, RequestStream, Side, StreamHeader, UniEvent, WireError};
use std::sync::{Arc, Mutex};

/// Why reading or writing a stream stopped.
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum StreamError {
    #[error("transport: {0}")]
    Transport(#[from] TransportError),
    #[error("moqt: {0}")]
    Wire(#[from] WireError),
}

/// Framed control messages (§9: Type, a 16-bit Length, the body) read from a stream.
pub struct MessageReader {
    recv: BoxRecv,
    buf: Vec<u8>,
    fin: bool,
}

impl MessageReader {
    pub fn new(recv: BoxRecv) -> Self {
        MessageReader { recv, buf: Vec::new(), fin: false }
    }

    /// The length of the whole message at the front of `b`, once its header is there.
    fn framed_len(b: &[u8]) -> Option<usize> {
        let first = *b.first()?;
        let tlen = joule_moq::draft22::varint::len_from_first(first);
        if b.len() < tlen + 2 {
            return None;
        }
        let len = u16::from_be_bytes([b[tlen], b[tlen + 1]]) as usize;
        Some(tlen + 2 + len)
    }

    /// The next message; `None` at a FIN on a message boundary.
    pub async fn next(&mut self) -> Result<Option<ControlMessage>, StreamError> {
        loop {
            if let Some(n) = Self::framed_len(&self.buf).filter(|n| *n <= self.buf.len()) {
                let m = ControlMessage::decode_exact(&self.buf[..n])?;
                self.buf.drain(..n);
                return Ok(Some(m));
            }
            if self.fin {
                if self.buf.is_empty() {
                    return Ok(None);
                }
                return Err(WireError::Truncated { what: "a stream ends inside a control message" }.into());
            }
            match self.recv.read_chunk().await? {
                Some(b) => self.buf.extend_from_slice(&b),
                None => self.fin = true,
            }
        }
    }

    pub fn stop(&mut self, code: u64) {
        self.recv.stop(code);
    }
}

/// The items of a unidirectional stream, each yielded once all of it has arrived.
pub struct UniReader {
    recv: BoxRecv,
    dec: UniStreamDecoder,
    fin: bool,
    /// What the stream began with, once read.
    pub header: Option<StreamHeader>,
}

impl UniReader {
    pub fn new(recv: BoxRecv, max_payload: u64) -> Self {
        let limits = Limits { max_payload, ..Limits::default() };
        UniReader { recv, dec: UniStreamDecoder::new(GroupOrder::Ascending, limits), fin: false, header: None }
    }

    /// Read up to the stream's header, and return it.
    pub async fn start(&mut self) -> Result<StreamHeader, StreamError> {
        if let Some(h) = &self.header {
            return Ok(h.clone());
        }
        match self.next().await? {
            Some(UniEvent::Header(h)) => Ok(h),
            _ => Err(WireError::Truncated { what: "a unidirectional stream ends before its type" }.into()),
        }
    }

    /// The Group Order of the request a fetch stream answers.
    pub fn set_group_order(&mut self, order: GroupOrder) {
        self.dec.set_group_order(order);
    }

    /// The next item; `None` at the FIN.
    pub async fn next(&mut self) -> Result<Option<UniEvent>, StreamError> {
        loop {
            if let Some(e) = self.dec.next()? {
                if let UniEvent::Header(h) = &e {
                    self.header = Some(h.clone());
                }
                return Ok(Some(e));
            }
            if self.fin {
                self.dec.finish()?;
                return Ok(None);
            }
            match self.recv.read_chunk().await? {
                Some(b) => self.dec.push(&b),
                None => self.fin = true,
            }
        }
    }

    pub fn stop(&mut self, code: u64) {
        self.recv.stop(code);
    }
}

/// The sending direction of a request stream, shareable between tasks: every message is
/// checked against the stream's rules before it goes ([`RequestStream::on_send`]).
#[derive(Clone)]
pub struct ReqTx {
    send: Arc<tokio::sync::Mutex<BoxSend>>,
    pub state: Arc<Mutex<RequestStream>>,
}

impl ReqTx {
    pub async fn send(&self, m: &ControlMessage) -> Result<(), StreamError> {
        self.state.lock().unwrap().on_send(m)?;
        let b = m.encode()?;
        self.send.lock().await.write_all(&b).await?;
        Ok(())
    }

    /// A FIN on this direction (§6.4.2.2).
    pub async fn fin(&self) {
        let _ = self.send.lock().await.finish();
    }

    /// Cancel this direction (§6.4.2.3).
    pub async fn reset(&self, code: u64) {
        self.send.lock().await.reset(code);
    }
}

/// The receiving direction of a request stream: each message admitted by the stream's
/// rules ([`RequestStream::on_receive`]).
pub struct ReqRx {
    pub reader: MessageReader,
    pub state: Arc<Mutex<RequestStream>>,
}

impl ReqRx {
    pub async fn next(&mut self) -> Result<Option<ControlMessage>, StreamError> {
        let m = self.reader.next().await?;
        if let Some(m) = &m {
            self.state.lock().unwrap().on_receive(m)?;
        }
        Ok(m)
    }
}

impl std::fmt::Debug for ReqTx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ReqTx").field("kind", &self.state.lock().unwrap().kind()).finish_non_exhaustive()
    }
}

impl std::fmt::Debug for ReqRx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ReqRx").field("kind", &self.state.lock().unwrap().kind()).finish_non_exhaustive()
    }
}

/// A request stream this endpoint opened, its request already sent.
pub fn requester(send: BoxSend, recv: BoxRecv, kind: RequestKind, my_limit: u64, peer_limit: u64) -> (ReqTx, ReqRx) {
    let state = Arc::new(Mutex::new(RequestStream::new(kind, Side::Requester, my_limit, peer_limit)));
    (ReqTx { send: Arc::new(tokio::sync::Mutex::new(send)), state: state.clone() }, ReqRx { reader: MessageReader::new(recv), state })
}

/// A request stream the peer opened with `first`.
pub fn responder(send: BoxSend, reader: MessageReader, first: &ControlMessage, my_limit: u64, peer_limit: u64) -> Result<(ReqTx, ReqRx), WireError> {
    let state = Arc::new(Mutex::new(RequestStream::opened_by_peer(first, my_limit, peer_limit)?));
    Ok((ReqTx { send: Arc::new(tokio::sync::Mutex::new(send)), state: state.clone() }, ReqRx { reader, state }))
}

/// Open a request stream on `conn` with `first` (§6.3).
pub async fn open_request(
    conn: &dyn crate::transport::Connection,
    first: &ControlMessage,
    my_limit: u64,
    peer_limit: u64,
) -> Result<(ReqTx, ReqRx), StreamError> {
    let kind = first.opens_request().ok_or(WireError::UnexpectedMessage("not a request"))?;
    let (mut send, recv) = conn.open_bi().await?;
    send.write_all(&first.encode()?).await?;
    Ok(requester(send, recv, kind, my_limit, peer_limit))
}