moqtap_client/

connection.rs

use std::sync::Arc;

use bytes::{Buf, BytesMut};

use moqtap_codec::data_stream::{DatagramHeader, FetchHeader, ObjectHeader, SubgroupHeader};
use moqtap_codec::message::{CodecError, ControlMessage};
use moqtap_codec::types::*;
use moqtap_codec::varint::VarInt;
use moqtap_session::request_id::Role;

use crate::endpoint::{Endpoint, EndpointError};

/// MoQT ALPN identifier (draft-14).
pub const MOQT_ALPN: &[u8] = b"moq-00";

/// Errors from the connection layer.
#[derive(Debug, thiserror::Error)]
pub enum ConnectionError {
    #[error("endpoint error: {0}")]
    Endpoint(#[from] EndpointError),
    #[error("codec error: {0}")]
    Codec(#[from] CodecError),
    #[error("quinn connection error: {0}")]
    Connection(#[from] quinn::ConnectionError),
    #[error("quinn write error: {0}")]
    Write(#[from] quinn::WriteError),
    #[error("quinn read error: {0}")]
    Read(#[from] quinn::ReadExactError),
    #[error("quinn connect error: {0}")]
    Connect(#[from] quinn::ConnectError),
    #[error("closed stream: {0}")]
    ClosedStream(#[from] quinn::ClosedStream),
    #[error("send datagram error: {0}")]
    SendDatagram(#[from] quinn::SendDatagramError),
    #[error("varint error: {0}")]
    VarInt(#[from] moqtap_codec::varint::VarIntError),
    #[error("control stream not open")]
    NoControlStream,
    #[error("unexpected end of stream")]
    UnexpectedEnd,
    #[error("stream finished")]
    StreamFinished,
    #[error("invalid server address: {0}")]
    InvalidAddress(String),
    #[error("TLS config error: {0}")]
    TlsConfig(String),
}

/// Configuration for a MoQT client connection.
pub struct ClientConfig {
    /// MoQT versions the client supports.
    pub supported_versions: Vec<VarInt>,
    /// Whether to skip TLS certificate verification (for testing).
    pub skip_cert_verification: bool,
}

impl Default for ClientConfig {
    fn default() -> Self {
        Self {
            supported_versions: vec![VarInt::from_u64(0xff000000 + 14).unwrap()], // draft-14
            skip_cert_verification: false,
        }
    }
}

/// A framed writer for a QUIC send stream. Handles MoQT length-prefixed framing.
pub struct FramedSendStream {
    inner: quinn::SendStream,
}

impl FramedSendStream {
    pub fn new(inner: quinn::SendStream) -> Self {
        Self { inner }
    }

    /// Write a control message to the stream with type+length framing.
    pub async fn write_control(&mut self, msg: &ControlMessage) -> Result<(), ConnectionError> {
        let mut buf = Vec::new();
        msg.encode(&mut buf)?;
        self.inner.write_all(&buf).await?;
        Ok(())
    }

    /// Write a subgroup stream header.
    pub async fn write_subgroup_header(
        &mut self,
        header: &SubgroupHeader,
    ) -> Result<(), ConnectionError> {
        let mut buf = Vec::new();
        header.encode(&mut buf);
        self.inner.write_all(&buf).await?;
        Ok(())
    }

    /// Write a fetch response header.
    pub async fn write_fetch_header(
        &mut self,
        header: &FetchHeader,
    ) -> Result<(), ConnectionError> {
        let mut buf = Vec::new();
        header.encode(&mut buf);
        self.inner.write_all(&buf).await?;
        Ok(())
    }

    /// Write an object header followed by payload.
    pub async fn write_object(
        &mut self,
        header: &ObjectHeader,
        payload: &[u8],
    ) -> Result<(), ConnectionError> {
        let mut buf = Vec::new();
        header.encode(&mut buf);
        self.inner.write_all(&buf).await?;
        if !payload.is_empty() {
            self.inner.write_all(payload).await?;
        }
        Ok(())
    }

    /// Finish the stream (send FIN).
    pub async fn finish(&mut self) -> Result<(), ConnectionError> {
        self.inner.finish()?;
        Ok(())
    }
}

/// A framed reader for a QUIC recv stream. Handles MoQT varint-length decoding.
pub struct FramedRecvStream {
    inner: quinn::RecvStream,
    buf: BytesMut,
}

impl FramedRecvStream {
    pub fn new(inner: quinn::RecvStream) -> Self {
        Self { inner, buf: BytesMut::with_capacity(4096) }
    }

    /// Read more data from the stream into the internal buffer.
    async fn fill(&mut self) -> Result<bool, ConnectionError> {
        let mut tmp = [0u8; 4096];
        match self.inner.read(&mut tmp).await {
            Ok(Some(n)) => {
                self.buf.extend_from_slice(&tmp[..n]);
                Ok(true)
            }
            Ok(None) => Ok(false),
            Err(e) => Err(ConnectionError::Read(quinn::ReadExactError::ReadError(e))),
        }
    }

    /// Ensure at least `n` bytes are available in the buffer.
    async fn ensure(&mut self, n: usize) -> Result<(), ConnectionError> {
        while self.buf.len() < n {
            if !self.fill().await? {
                return Err(ConnectionError::UnexpectedEnd);
            }
        }
        Ok(())
    }

    /// Read a control message from the stream.
    pub async fn read_control(&mut self) -> Result<ControlMessage, ConnectionError> {
        // Read type ID varint
        self.ensure(1).await?;
        let type_len = varint_len(self.buf[0]);
        self.ensure(type_len).await?;

        let mut cursor = &self.buf[..type_len];
        let _type_id = VarInt::decode(&mut cursor)?;

        // Read payload length varint
        self.ensure(type_len + 1).await?;
        let payload_len_start = type_len;
        let payload_len_varint_len = varint_len(self.buf[payload_len_start]);
        self.ensure(type_len + payload_len_varint_len).await?;

        let mut cursor = &self.buf[payload_len_start..type_len + payload_len_varint_len];
        let payload_len = VarInt::decode(&mut cursor)?.into_inner() as usize;

        // Read full payload
        let total = type_len + payload_len_varint_len + payload_len;
        self.ensure(total).await?;

        // Now decode the whole message
        let mut frame = &self.buf[..total];
        let msg = ControlMessage::decode(&mut frame)?;
        self.buf.advance(total);
        Ok(msg)
    }

    /// Read a subgroup stream header.
    pub async fn read_subgroup_header(&mut self) -> Result<SubgroupHeader, ConnectionError> {
        // SubgroupHeader: track_alias(v) + group(v) + subgroup(v) + priority(1)
        // Need at least 4 bytes, but varints can be up to 8 each
        self.ensure(4).await?;
        loop {
            let mut cursor = &self.buf[..];
            match SubgroupHeader::decode(&mut cursor) {
                Ok(header) => {
                    let consumed = self.buf.len() - cursor.remaining();
                    self.buf.advance(consumed);
                    return Ok(header);
                }
                Err(_) => {
                    if !self.fill().await? {
                        return Err(ConnectionError::UnexpectedEnd);
                    }
                }
            }
        }
    }

    /// Read a fetch response header.
    pub async fn read_fetch_header(&mut self) -> Result<FetchHeader, ConnectionError> {
        self.ensure(4).await?;
        loop {
            let mut cursor = &self.buf[..];
            match FetchHeader::decode(&mut cursor) {
                Ok(header) => {
                    let consumed = self.buf.len() - cursor.remaining();
                    self.buf.advance(consumed);
                    return Ok(header);
                }
                Err(_) => {
                    if !self.fill().await? {
                        return Err(ConnectionError::UnexpectedEnd);
                    }
                }
            }
        }
    }

    /// Read an object header.
    pub async fn read_object_header(&mut self) -> Result<ObjectHeader, ConnectionError> {
        self.ensure(2).await?;
        loop {
            let mut cursor = &self.buf[..];
            match ObjectHeader::decode(&mut cursor) {
                Ok(header) => {
                    let consumed = self.buf.len() - cursor.remaining();
                    self.buf.advance(consumed);
                    return Ok(header);
                }
                Err(_) => {
                    if !self.fill().await? {
                        return Err(ConnectionError::UnexpectedEnd);
                    }
                }
            }
        }
    }

    /// Read exactly `n` bytes of object payload.
    pub async fn read_payload(&mut self, n: usize) -> Result<Vec<u8>, ConnectionError> {
        self.ensure(n).await?;
        let data = self.buf[..n].to_vec();
        self.buf.advance(n);
        Ok(data)
    }
}

/// A live MoQT connection over QUIC, combining the endpoint state machine
/// with actual network I/O.
pub struct Connection {
    quic: quinn::Connection,
    endpoint: Endpoint,
    control_send: Option<FramedSendStream>,
    control_recv: Option<FramedRecvStream>,
}

impl Connection {
    /// Connect to a MoQT server as a client.
    ///
    /// Opens a QUIC connection, performs the bidirectional control stream
    /// setup handshake (CLIENT_SETUP / SERVER_SETUP), and returns a
    /// ready-to-use connection.
    pub async fn connect(addr: &str, config: ClientConfig) -> Result<Self, ConnectionError> {
        let server_addr = addr.parse().map_err(|e: std::net::AddrParseError| {
            ConnectionError::InvalidAddress(e.to_string())
        })?;

        let mut endpoint = Endpoint::new(Role::Client);

        // Build TLS config
        let mut tls_config = if config.skip_cert_verification {
            rustls::ClientConfig::builder()
                .dangerous()
                .with_custom_certificate_verifier(Arc::new(SkipVerification))
                .with_no_client_auth()
        } else {
            let mut roots = rustls::RootCertStore::empty();
            roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
            rustls::ClientConfig::builder().with_root_certificates(roots).with_no_client_auth()
        };

        tls_config.alpn_protocols = vec![MOQT_ALPN.to_vec()];

        let quic_config: quinn::crypto::rustls::QuicClientConfig =
            tls_config.try_into().map_err(|e| ConnectionError::TlsConfig(format!("{e}")))?;
        let client_config = quinn::ClientConfig::new(Arc::new(quic_config));

        let mut quinn_endpoint = quinn::Endpoint::client("0.0.0.0:0".parse().unwrap())
            .map_err(|e| ConnectionError::InvalidAddress(e.to_string()))?;
        quinn_endpoint.set_default_client_config(client_config);

        let server_name = addr.split(':').next().unwrap_or("localhost").to_string();

        let quic = quinn_endpoint.connect(server_addr, &server_name)?.await?;

        // Open bidirectional control stream
        let (send, recv) = quic.open_bi().await?;
        let mut control_send = FramedSendStream::new(send);
        let mut control_recv = FramedRecvStream::new(recv);

        // Perform setup handshake
        endpoint.connect()?;
        let setup_msg = endpoint.send_client_setup(config.supported_versions)?;
        control_send.write_control(&setup_msg).await?;

        let server_setup = control_recv.read_control().await?;
        if let ControlMessage::ServerSetup(ref ss) = server_setup {
            endpoint.receive_server_setup(ss)?;
        } else {
            return Err(ConnectionError::Endpoint(EndpointError::NotActive));
        }

        Ok(Self {
            quic,
            endpoint,
            control_send: Some(control_send),
            control_recv: Some(control_recv),
        })
    }

    /// Accept a MoQT connection as a server (given an already-accepted QUIC connection).
    pub async fn accept(
        quic: quinn::Connection,
        selected_version: VarInt,
    ) -> Result<Self, ConnectionError> {
        let mut endpoint = Endpoint::new(Role::Server);
        endpoint.connect()?;

        // Accept the bidirectional control stream from the client
        let (send, recv) = quic.accept_bi().await?;
        let mut control_send = FramedSendStream::new(send);
        let mut control_recv = FramedRecvStream::new(recv);

        // Read CLIENT_SETUP
        let client_setup_msg = control_recv.read_control().await?;
        if let ControlMessage::ClientSetup(ref cs) = client_setup_msg {
            let server_setup = endpoint.receive_client_setup_and_respond(cs, selected_version)?;
            control_send.write_control(&server_setup).await?;
        } else {
            return Err(ConnectionError::Endpoint(EndpointError::NotActive));
        }

        Ok(Self {
            quic,
            endpoint,
            control_send: Some(control_send),
            control_recv: Some(control_recv),
        })
    }

    // ── Control message I/O ─────────────────────────────────

    /// Send a control message on the control stream.
    pub async fn send_control(&mut self, msg: &ControlMessage) -> Result<(), ConnectionError> {
        let send = self.control_send.as_mut().ok_or(ConnectionError::NoControlStream)?;
        send.write_control(msg).await
    }

    /// Read the next control message from the control stream.
    pub async fn recv_control(&mut self) -> Result<ControlMessage, ConnectionError> {
        let recv = self.control_recv.as_mut().ok_or(ConnectionError::NoControlStream)?;
        recv.read_control().await
    }

    /// Read and dispatch the next incoming control message through the endpoint
    /// state machine. Returns the decoded message for inspection.
    pub async fn recv_and_dispatch(&mut self) -> Result<ControlMessage, ConnectionError> {
        let msg = self.recv_control().await?;
        self.endpoint.receive_message(msg.clone())?;
        Ok(msg)
    }

    // ── Subscribe flow ──────────────────────────────────────

    /// Send a SUBSCRIBE and return the allocated request ID.
    pub async fn subscribe(
        &mut self,
        track_namespace: TrackNamespace,
        track_name: Vec<u8>,
        subscriber_priority: u8,
        group_order: GroupOrder,
        filter_type: FilterType,
    ) -> Result<VarInt, ConnectionError> {
        let (req_id, msg) = self.endpoint.subscribe(
            track_namespace,
            track_name,
            subscriber_priority,
            group_order,
            filter_type,
        )?;
        self.send_control(&msg).await?;
        Ok(req_id)
    }

    /// Send an UNSUBSCRIBE for the given request ID.
    pub async fn unsubscribe(&mut self, request_id: VarInt) -> Result<(), ConnectionError> {
        let msg = self.endpoint.unsubscribe(request_id)?;
        self.send_control(&msg).await
    }

    // ── Fetch flow ──────────────────────────────────────────

    /// Send a FETCH and return the allocated request ID.
    pub async fn fetch(
        &mut self,
        track_namespace: TrackNamespace,
        track_name: Vec<u8>,
        start_group: VarInt,
        start_object: VarInt,
    ) -> Result<VarInt, ConnectionError> {
        let (req_id, msg) =
            self.endpoint.fetch(track_namespace, track_name, start_group, start_object)?;
        self.send_control(&msg).await?;
        Ok(req_id)
    }

    /// Send a FETCH_CANCEL for the given request ID.
    pub async fn fetch_cancel(&mut self, request_id: VarInt) -> Result<(), ConnectionError> {
        let msg = self.endpoint.fetch_cancel(request_id)?;
        self.send_control(&msg).await
    }

    // ── Namespace flows ─────────────────────────────────────

    /// Send a SUBSCRIBE_NAMESPACE and return the request ID.
    pub async fn subscribe_namespace(
        &mut self,
        track_namespace: TrackNamespace,
    ) -> Result<VarInt, ConnectionError> {
        let (req_id, msg) = self.endpoint.subscribe_namespace(track_namespace)?;
        self.send_control(&msg).await?;
        Ok(req_id)
    }

    /// Send a PUBLISH_NAMESPACE and return the request ID.
    pub async fn publish_namespace(
        &mut self,
        track_namespace: TrackNamespace,
    ) -> Result<VarInt, ConnectionError> {
        let (req_id, msg) = self.endpoint.publish_namespace(track_namespace)?;
        self.send_control(&msg).await?;
        Ok(req_id)
    }

    // ── Data streams ────────────────────────────────────────

    /// Open a new unidirectional stream for sending subgroup data.
    pub async fn open_subgroup_stream(
        &self,
        header: &SubgroupHeader,
    ) -> Result<FramedSendStream, ConnectionError> {
        let send = self.quic.open_uni().await?;
        let mut framed = FramedSendStream::new(send);
        framed.write_subgroup_header(header).await?;
        Ok(framed)
    }

    /// Accept an incoming unidirectional data stream and read its subgroup header.
    pub async fn accept_subgroup_stream(
        &self,
    ) -> Result<(SubgroupHeader, FramedRecvStream), ConnectionError> {
        let recv = self.quic.accept_uni().await?;
        let mut framed = FramedRecvStream::new(recv);
        let header = framed.read_subgroup_header().await?;
        Ok((header, framed))
    }

    /// Send an object via datagram.
    pub fn send_datagram(
        &self,
        header: &DatagramHeader,
        payload: &[u8],
    ) -> Result<(), ConnectionError> {
        let mut buf = Vec::new();
        header.encode(&mut buf);
        buf.extend_from_slice(payload);
        self.quic.send_datagram(bytes::Bytes::from(buf))?;
        Ok(())
    }

    /// Receive a datagram and decode its header.
    pub async fn recv_datagram(&self) -> Result<(DatagramHeader, Vec<u8>), ConnectionError> {
        let data = self.quic.read_datagram().await?;
        let mut cursor = &data[..];
        let header = DatagramHeader::decode(&mut cursor)?;
        let payload = cursor.to_vec();
        Ok((header, payload))
    }

    // ── Accessors ───────────────────────────────────────────

    /// Access the underlying endpoint state machine.
    pub fn endpoint(&self) -> &Endpoint {
        &self.endpoint
    }

    /// Mutable access to the endpoint state machine.
    pub fn endpoint_mut(&mut self) -> &mut Endpoint {
        &mut self.endpoint
    }

    /// Get the negotiated MoQT version.
    pub fn negotiated_version(&self) -> Option<VarInt> {
        self.endpoint.negotiated_version()
    }

    /// Close the connection.
    pub fn close(&self, code: u32, reason: &[u8]) {
        self.quic.close(quinn::VarInt::from_u32(code), reason);
    }
}

/// Determine the encoded length of a varint from its first byte.
fn varint_len(first_byte: u8) -> usize {
    1 << (first_byte >> 6)
}

/// TLS certificate verifier that skips all verification (for testing only).
#[derive(Debug)]
struct SkipVerification;

impl rustls::client::danger::ServerCertVerifier for SkipVerification {
    fn verify_server_cert(
        &self,
        _end_entity: &rustls::pki_types::CertificateDer<'_>,
        _intermediates: &[rustls::pki_types::CertificateDer<'_>],
        _server_name: &rustls::pki_types::ServerName<'_>,
        _ocsp_response: &[u8],
        _now: rustls::pki_types::UnixTime,
    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::danger::ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::pki_types::CertificateDer<'_>,
        _dcs: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::pki_types::CertificateDer<'_>,
        _dcs: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        vec![
            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
            rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
            rustls::SignatureScheme::ED25519,
            rustls::SignatureScheme::RSA_PSS_SHA256,
            rustls::SignatureScheme::RSA_PSS_SHA384,
            rustls::SignatureScheme::RSA_PSS_SHA512,
        ]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn varint_len_single_byte() {
        // 0b00xxxxxx -> 1 byte
        assert_eq!(varint_len(0x00), 1);
        assert_eq!(varint_len(0x3F), 1);
    }

    #[test]
    fn varint_len_two_bytes() {
        // 0b01xxxxxx -> 2 bytes
        assert_eq!(varint_len(0x40), 2);
        assert_eq!(varint_len(0x7F), 2);
    }

    #[test]
    fn varint_len_four_bytes() {
        // 0b10xxxxxx -> 4 bytes
        assert_eq!(varint_len(0x80), 4);
        assert_eq!(varint_len(0xBF), 4);
    }

    #[test]
    fn varint_len_eight_bytes() {
        // 0b11xxxxxx -> 8 bytes
        assert_eq!(varint_len(0xC0), 8);
        assert_eq!(varint_len(0xFF), 8);
    }

    #[test]
    fn default_client_config() {
        let config = ClientConfig::default();
        assert_eq!(config.supported_versions.len(), 1);
        assert!(!config.skip_cert_verification);
    }

    #[test]
    fn moqt_alpn_value() {
        assert_eq!(MOQT_ALPN, b"moq-00");
    }
}