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 between two MoQT peers. Both ends
//! traverse the same lifecycle: `Init → Handshaking → Established → Closing
//! → Closed`. The role (publisher / subscriber / pubsub / relay) is
//! negotiated during handshake.

use crate::error::{MoqError, MoqResult};
use crate::messages::SetupMessage;
use crate::MOQT_DRAFT_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 {
    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,
    pub state: SessionState,
    pub negotiated_version: Option<u32>,
    pub jwp_receipts_enabled: bool,
    /// Active subscriptions, keyed by subscribe_id.
    pub subscriptions: HashMap<u64, SubscriptionState>,
    /// Track alias → subscribe_id reverse index.
    pub alias_index: HashMap<u64, u64>,
}

#[derive(Clone, Debug)]
pub struct SubscriptionState {
    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,
}

impl MoqSession {
    pub fn new(role: SessionRole) -> Self {
        Self {
            role,
            state: SessionState::Init,
            negotiated_version: None,
            jwp_receipts_enabled: false,
            subscriptions: HashMap::new(),
            alias_index: HashMap::new(),
        }
    }

    /// Produce the SETUP message this endpoint will send first.
    pub fn build_setup(&self, request_jwp: bool, label: Option<String>) -> SetupMessage {
        SetupMessage {
            versions: vec![MOQT_DRAFT_VERSION],
            role: self.role.role_bits(),
            jwp_receipts: request_jwp,
            peer_label: label,
        }
    }

    /// Process a peer SETUP and transition to Handshaking → Established.
    pub fn accept_setup(&mut self, peer: &SetupMessage) -> MoqResult<SetupMessage> {
        if !peer.versions.contains(&MOQT_DRAFT_VERSION) {
            return Err(MoqError::UnsupportedVersion(
                peer.versions.first().copied().unwrap_or(0),
            ));
        }
        self.state = SessionState::Handshaking;
        self.negotiated_version = Some(MOQT_DRAFT_VERSION);
        self.jwp_receipts_enabled = peer.jwp_receipts;
        let response = self.build_setup(self.jwp_receipts_enabled, None);
        self.state = SessionState::Established;
        Ok(response)
    }

    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::*;

    #[test]
    fn handshake_round_trip() {
        let mut srv = MoqSession::new(SessionRole::PubSub);
        let client_setup = SetupMessage {
            versions: vec![MOQT_DRAFT_VERSION],
            role: 0x02,
            jwp_receipts: true,
            peer_label: Some("client".into()),
        };
        let resp = srv.accept_setup(&client_setup).unwrap();
        assert_eq!(srv.state, SessionState::Established);
        assert!(srv.jwp_receipts_enabled);
        assert!(resp.jwp_receipts);
    }

    #[test]
    fn rejects_unsupported_version() {
        let mut srv = MoqSession::new(SessionRole::Publisher);
        let bad_setup = SetupMessage {
            versions: vec![0x1234_5678],
            role: 0x02,
            jwp_receipts: false,
            peer_label: None,
        };
        let result = srv.accept_setup(&bad_setup);
        assert!(matches!(result, Err(MoqError::UnsupportedVersion(_))));
    }
}