open-standards/joule-code/joule-code-rs/crates/joule-moq/src/session.rs

//! Session state machine for a MoqT endpoint.
//!
//! A session is a single QUIC connection (or WebTransport session) between two
//! MoqT peers. Both ends traverse the same lifecycle: `Init → Handshaking →
//! Established → Closing → Closed`.
//!
//! Under `draft-ietf-moq-transport-22` the draft is fixed before the first MOQT
//! byte, by ALPN or WebTransport protocol ([`MoqSession::select_alpn`]); each
//! endpoint then opens a control stream with a SETUP of its own, carrying Setup
//! Options and no version and no role. Extensions are declared there: this crate
//! declares the JWP receipt extension by
//! [`crate::carriage::SETUP_OPTION_JWP_RECEIPTS`]. The role
//! ([`SessionRole`]) is this endpoint's own business; -22 does not exchange it.

use crate::carriage::SETUP_OPTION_JWP_RECEIPTS;
use crate::draft::Draft;
use crate::draft22::kvp::{opt, KeyValuePair, SetupOptions};
use crate::draft22::{PeerRequestIds, Perspective, RequestIds, Setup, WireError};
use crate::error::{MoqError, MoqResult};
use crate::JWP_MOQ_EXTENSION_VERSION;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SessionRole {
    Publisher,
    Subscriber,
    PubSub,
    Relay,
}

impl SessionRole {
    /// The role as the bits draft-08's SETUP carried. -22 exchanges no role;
    /// kept for callers that log or store it.
    pub fn role_bits(&self) -> u8 {
        match self {
            SessionRole::Publisher => 0x01,
            SessionRole::Subscriber => 0x02,
            SessionRole::PubSub => 0x03,
            SessionRole::Relay => 0x03,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SessionState {
    Init,
    Handshaking,
    Established,
    Closing,
    Closed,
}

/// Cross-cutting state for a single MoqT session.
pub struct MoqSession {
    pub role: SessionRole,
    /// Whether this endpoint opened the transport session.
    pub perspective: Perspective,
    pub state: SessionState,
    /// The draft the ALPN selected.
    pub draft: Option<Draft>,
    /// That draft's 32-bit code, once the session is established.
    pub negotiated_version: Option<u32>,
    pub jwp_receipts_enabled: bool,
    /// Active subscriptions, keyed by the Request ID of their SUBSCRIBE or PUBLISH.
    pub subscriptions: HashMap<u64, SubscriptionState>,
    /// Track alias → Request ID reverse index.
    pub alias_index: HashMap<u64, u64>,
    /// This endpoint's own Request IDs.
    pub request_ids: RequestIds,
    /// The peer's Request IDs seen so far.
    pub peer_request_ids: PeerRequestIds,
    /// The peer's Setup Options, once its SETUP arrived.
    pub peer_options: Option<SetupOptions>,
    /// MAX_REQUEST_UPDATES this endpoint declares (0: no limit).
    pub max_request_updates: u64,
    /// MAX_FILTER_RANGES this endpoint declares (0: the peer sends no Range Filter).
    pub max_filter_ranges: u64,
}

#[derive(Clone, Debug)]
pub struct SubscriptionState {
    /// The Request ID of the SUBSCRIBE (or PUBLISH) that made it.
    pub subscribe_id: u64,
    pub track_alias: u64,
    pub require_receipts: bool,
    pub last_group: Option<u64>,
    pub last_object: Option<u64>,
    pub joules_consumed_micro: u64,
}

/// The MOQT_IMPLEMENTATION this crate sends (§9.1.5).
pub fn implementation_name() -> String {
    format!("joule-moq/{}", crate::VERSION)
}

impl MoqSession {
    /// A session as the server end of its transport session.
    pub fn new(role: SessionRole) -> Self {
        Self::with_perspective(role, Perspective::Server)
    }

    pub fn with_perspective(role: SessionRole, perspective: Perspective) -> Self {
        Self {
            role,
            perspective,
            state: SessionState::Init,
            draft: None,
            negotiated_version: None,
            jwp_receipts_enabled: false,
            subscriptions: HashMap::new(),
            alias_index: HashMap::new(),
            request_ids: RequestIds::new(perspective),
            peer_request_ids: PeerRequestIds::new(perspective.peer()),
            peer_options: None,
            max_request_updates: 0,
            max_filter_ranges: 0,
        }
    }

    /// Choose the draft from the ALPNs a client offered (a server), or check the
    /// one a server selected (a client): refused when this crate speaks none of
    /// them.
    pub fn select_alpn<'a>(&mut self, offered: impl IntoIterator<Item = &'a [u8]>) -> MoqResult<Draft> {
        let d = Draft::select(offered).ok_or(MoqError::UnsupportedVersion(0))?;
        self.draft = Some(d);
        Ok(d)
    }

    /// The SETUP this endpoint sends: its implementation, the limits it
    /// declares, and the JWP receipt extension when `request_jwp`.
    pub fn build_setup(&self, request_jwp: bool) -> Setup {
        let mut o = SetupOptions::new();
        o.push(KeyValuePair::bytes(opt::MOQT_IMPLEMENTATION, implementation_name().into_bytes()));
        if self.max_filter_ranges != 0 {
            o.push(KeyValuePair::varint(opt::MAX_FILTER_RANGES, self.max_filter_ranges));
        }
        if self.max_request_updates != 0 {
            o.push(KeyValuePair::varint(opt::MAX_REQUEST_UPDATES, self.max_request_updates));
        }
        if request_jwp {
            o.push(KeyValuePair::varint(SETUP_OPTION_JWP_RECEIPTS, JWP_MOQ_EXTENSION_VERSION as u64));
        }
        Setup { options: o }
    }

    /// Whether a SETUP declares the JWP receipt extension at this crate's version.
    pub fn declares_jwp(setup: &Setup) -> bool {
        setup.options.extension(SETUP_OPTION_JWP_RECEIPTS).and_then(KeyValuePair::as_varint) == Some(JWP_MOQ_EXTENSION_VERSION as u64)
    }

    /// Take the peer's SETUP: refused where its options break the draft for a
    /// peer at its end (§9.1.1, §9.1.2, §9.1.4). The JWP receipt extension is on
    /// when the peer declares it at this crate's version and `we_offer` it.
    pub fn on_peer_setup(&mut self, peer: &Setup, we_offer_jwp: bool) -> MoqResult<()> {
        peer.check_from(self.perspective.peer(), false)
            .map_err(|code| MoqError::Wire(WireError::Session { code, reason: "the peer's SETUP breaks a rule for its end" }))?;
        self.state = SessionState::Handshaking;
        let d = *self.draft.get_or_insert(Draft::D22);
        self.negotiated_version = Some(d.version_code());
        self.jwp_receipts_enabled = we_offer_jwp && Self::declares_jwp(peer);
        self.peer_options = Some(peer.options.clone());
        self.state = SessionState::Established;
        Ok(())
    }

    /// Take the peer's SETUP and produce this endpoint's, offering the JWP
    /// extension exactly when the peer did.
    pub fn accept_setup(&mut self, peer: &Setup) -> MoqResult<Setup> {
        let offer = Self::declares_jwp(peer);
        self.on_peer_setup(peer, offer)?;
        Ok(self.build_setup(self.jwp_receipts_enabled))
    }

    /// The next Request ID this endpoint gives a request.
    pub fn next_request_id(&mut self) -> MoqResult<u64> {
        self.request_ids.next().ok_or_else(|| MoqError::InvalidSessionState("Request IDs exhausted".into()))
    }

    /// Admit a Request ID the peer sent (INVALID_REQUEST_ID otherwise).
    pub fn accept_peer_request_id(&mut self, id: u64) -> MoqResult<()> {
        Ok(self.peer_request_ids.accept(id)?)
    }

    pub fn add_subscription(&mut self, subscribe_id: u64, track_alias: u64, require_receipts: bool) -> MoqResult<()> {
        if self.state != SessionState::Established {
            return Err(MoqError::InvalidSessionState(format!("subscribe requires Established, was {:?}", self.state)));
        }
        self.subscriptions.insert(
            subscribe_id,
            SubscriptionState { subscribe_id, track_alias, require_receipts, last_group: None, last_object: None, joules_consumed_micro: 0 },
        );
        self.alias_index.insert(track_alias, subscribe_id);
        Ok(())
    }

    pub fn remove_subscription(&mut self, subscribe_id: u64) {
        if let Some(s) = self.subscriptions.remove(&subscribe_id) {
            self.alias_index.remove(&s.track_alias);
        }
    }

    pub fn record_object(&mut self, track_alias: u64, group_id: u64, object_id: u64, joules_micro: u64) -> MoqResult<()> {
        let sub_id = *self.alias_index.get(&track_alias).ok_or(MoqError::SubscriptionNotFound(track_alias))?;
        let sub = self.subscriptions.get_mut(&sub_id).ok_or(MoqError::SubscriptionNotFound(track_alias))?;
        sub.last_group = Some(group_id);
        sub.last_object = Some(object_id);
        sub.joules_consumed_micro = sub.joules_consumed_micro.saturating_add(joules_micro);
        Ok(())
    }

    pub fn close(&mut self) {
        self.state = SessionState::Closing;
        self.subscriptions.clear();
        self.alias_index.clear();
        self.state = SessionState::Closed;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::draft22::codes::SessionError;
    use crate::draft22::Token;
    use crate::MOQT_DRAFT_VERSION;

    /// The handshake: the client's SETUP declares the receipt extension, the
    /// server's answers it, and both sessions are established on -22.
    #[test]
    fn handshake_round_trip() {
        let mut srv = MoqSession::new(SessionRole::PubSub);
        let client = MoqSession::with_perspective(SessionRole::Subscriber, Perspective::Client);
        let client_setup = client.build_setup(true);
        let resp = srv.accept_setup(&client_setup).unwrap();
        assert_eq!(srv.state, SessionState::Established);
        assert!(srv.jwp_receipts_enabled);
        assert!(MoqSession::declares_jwp(&resp));
        assert_eq!(srv.negotiated_version, Some(MOQT_DRAFT_VERSION));
        assert_eq!(resp.options.implementation(), Some(implementation_name().as_bytes()));
        // Through the codec and back.
        let wire = crate::draft22::ControlMessage::Setup(resp.clone()).encode().unwrap();
        assert_eq!(crate::draft22::ControlMessage::decode_exact(&wire).unwrap(), crate::draft22::ControlMessage::Setup(resp));
    }

    /// A version is refused at the ALPN, before any SETUP: -22 carries none in
    /// SETUP. And a SETUP that breaks a rule for its sender's end is refused.
    #[test]
    fn rejects_unsupported_version() {
        let mut srv = MoqSession::new(SessionRole::Publisher);
        let old: [&[u8]; 2] = [b"moq-00", b"moqt-18"];
        assert!(matches!(srv.select_alpn(old), Err(MoqError::UnsupportedVersion(_))));
        let ok: [&[u8]; 2] = [b"moqt-18", b"moqt-22"];
        assert_eq!(srv.select_alpn(ok).unwrap(), Draft::D22);

        // A client receiving PATH from the server.
        let mut client = MoqSession::with_perspective(SessionRole::Subscriber, Perspective::Client);
        let mut bad = srv.build_setup(false);
        bad.options.push(KeyValuePair::bytes(opt::PATH, b"/x".to_vec()));
        assert!(matches!(client.on_peer_setup(&bad, false), Err(MoqError::Wire(WireError::Session { code: SessionError::INVALID_PATH, .. }))));
        // A server receiving USE_ALIAS in a client's SETUP.
        let mut c = client.build_setup(false);
        c.options.push(KeyValuePair::bytes(opt::AUTHORIZATION_TOKEN, Token::UseAlias { alias: 1 }.to_bytes()));
        assert!(matches!(srv.on_peer_setup(&c, false), Err(MoqError::Wire(WireError::Session { code: SessionError::PROTOCOL_VIOLATION, .. }))));
    }

    #[test]
    fn request_ids_by_end() {
        let mut c = MoqSession::with_perspective(SessionRole::Subscriber, Perspective::Client);
        let mut s = MoqSession::new(SessionRole::Publisher);
        assert_eq!((c.next_request_id().unwrap(), c.next_request_id().unwrap()), (0, 2));
        s.accept_peer_request_id(0).unwrap();
        assert!(s.accept_peer_request_id(0).is_err());
        assert!(s.accept_peer_request_id(1).is_err());
        assert_eq!(s.next_request_id().unwrap(), 1);
    }
}