umsh_node/

host.rs

use alloc::rc::Rc;
use alloc::vec::Vec;
use core::cell::RefCell;

use umsh_mac::{LocalIdentityId, SendOptions};

use crate::dispatch::EventDispatcher;
use crate::identity_responder::IdentityResponsePlan;
use crate::mac::MacBackend;
use crate::node::{LocalNode, LocalNodeState, NodeMembership, PfsLifecycle};
use crate::receive::ReceivedPacketRef;
use crate::{NodeIdentityPayload, OwnedMacCommand, mac_command};
use umsh_core::PayloadType;

/// Error returned when a [`Host`] cannot make progress.
///
/// Wraps the underlying MAC runtime failure.
#[derive(Debug)]
pub enum HostError<E> {
    Mac(E),
}

/// Multi-identity orchestration layer for the node API.
///
/// `Host` owns the shared MAC run loop and routes inbound events to the correct
/// [`LocalNode`], including ephemeral PFS identities that belong to a long-term node.
/// Applications typically:
///
/// 1. construct a `Host` from a shared [`MacHandle`](umsh_mac::MacHandle)
/// 2. register one or more [`LocalNode`]s with [`add_node`](Self::add_node)
/// 3. attach callbacks to nodes / peers / wrappers
/// 4. drive progress with [`run`](Self::run) or [`pump_once`](Self::pump_once)
///
/// `run()` is the preferred long-lived driver. `pump_once()` exists for callers that need to
/// multiplex UMSH progress with other async work using `select!`.
///
/// # Callback re-entrancy
///
/// Receive and control callbacks (`on_receive`, `on_text`, `on_ack_received`, …) are invoked
/// synchronously from inside [`pump_once`](Self::pump_once) while the coordinator borrow is
/// held. A callback therefore MUST NOT call back into the MAC — it cannot `await` a `send`,
/// and a blocking or re-entrant borrow of the coordinator would deadlock. Keep callbacks
/// short and side-effect-free with respect to the MAC: record what you need (e.g. push to a
/// queue or set a flag) and perform any follow-up sends after `pump_once` / `run` returns
/// control to your own task.
pub struct Host<M: MacBackend> {
    mac: M,
    dispatcher: Rc<RefCell<EventDispatcher>>,
    nodes: Vec<(LocalIdentityId, LocalNode<M>)>,
    pfs_control_options: SendOptions,
}

impl<M: MacBackend> Host<M> {
    /// Create a host around a shared MAC backend.
    pub fn new(mac: M) -> Self {
        Self {
            mac,
            dispatcher: Rc::new(RefCell::new(EventDispatcher::new())),
            nodes: Vec::new(),
            pfs_control_options: SendOptions::default()
                .with_ack_requested(true)
                .with_flood_hops(5),
        }
    }

    /// Return a clone of the underlying shared MAC backend.
    pub fn mac(&self) -> M {
        self.mac.clone()
    }

    /// Borrow the send options used for node-managed PFS control messages.
    pub fn pfs_control_options(&self) -> &SendOptions {
        &self.pfs_control_options
    }

    /// Replace the send options used for node-managed PFS control messages.
    pub fn set_pfs_control_options(&mut self, options: SendOptions) {
        self.pfs_control_options = options;
    }

    /// Create and register a [`LocalNode`] for an already-registered local identity.
    ///
    /// The returned handle is cheap to clone and becomes the application-facing entry point
    /// for sending traffic and attaching callbacks for that identity.
    pub fn add_node(&mut self, identity_id: LocalIdentityId) -> LocalNode<M> {
        let membership = Rc::new(RefCell::new(NodeMembership::new()));
        let state = Rc::new(RefCell::new(LocalNodeState::new()));
        let node = LocalNode::new(
            identity_id,
            self.mac.clone(),
            self.dispatcher.clone(),
            membership,
            state,
        );
        self.nodes.push((identity_id, node.clone()));
        node
    }

    /// Look up a previously added node by identity id.
    pub fn node(&self, identity_id: LocalIdentityId) -> Option<LocalNode<M>> {
        self.nodes
            .iter()
            .find(|(id, _)| *id == identity_id)
            .map(|(_, node)| node.clone())
    }

    fn route_node(&self, identity_id: LocalIdentityId) -> Option<LocalNode<M>> {
        if let Some(node) = self.node(identity_id) {
            return Some(node);
        }

        #[cfg(feature = "software-crypto")]
        {
            return self
                .nodes
                .iter()
                .find(|(_, node)| node.owns_ephemeral_identity(identity_id))
                .map(|(_, node)| node.clone());
        }

        #[cfg(not(feature = "software-crypto"))]
        {
            None
        }
    }

    /// Drive the shared MAC until one wake cycle completes.
    ///
    /// This is a single wake-driven step, not a non-blocking poll. It waits until the MAC has
    /// meaningful work to do (radio activity or a protocol deadline), dispatches any resulting
    /// callbacks, services PFS command handling, and then returns.
    ///
    /// Use this when you need to multiplex UMSH progress with other async sources using
    /// `select!`. If UMSH owns the task, prefer [`run`](Self::run).
    pub async fn pump_once(&mut self) -> Result<(), HostError<M::RunError>> {
        let now_ms = self.mac.now_ms().await;
        let pending_pfs = Rc::new(RefCell::new(Vec::<(
            LocalIdentityId,
            umsh_core::PublicKey,
            OwnedMacCommand,
        )>::new()));
        let pending_pfs_ref = pending_pfs.clone();
        let pending_identity = Rc::new(RefCell::new(
            Vec::<(LocalNode<M>, IdentityResponsePlan)>::new(),
        ));
        let pending_identity_ref = pending_identity.clone();
        let dispatcher = self.dispatcher.clone();
        let nodes = self.nodes.clone();
        self.mac
            .next_event(move |identity_id, event| {
                dispatcher
                    .borrow_mut()
                    .dispatch_ticket_state(identity_id, &event);
                let Some(node) = route_node(&nodes, identity_id) else {
                    return;
                };
                match event {
                    umsh_mac::MacEventRef::Received(packet) => {
                        let _ = node.dispatch_received_packet(&packet);
                        if packet.packet_type() == umsh_core::PacketType::Broadcast
                            && packet.payload().is_empty()
                        {
                            if let Some(from_hint) = packet.from_hint() {
                                node.dispatch_beacon(from_hint, packet.from_key());
                            }
                        } else if let Some(from) = packet.from_key() {
                            dispatch_payload_callbacks(
                                &node,
                                &packet,
                                from,
                                &pending_pfs_ref,
                                &pending_identity_ref,
                                now_ms,
                            );
                        }
                    }
                    umsh_mac::MacEventRef::AckReceived { peer, receipt } => {
                        node.dispatch_ack_received(
                            peer,
                            crate::SendToken::new(identity_id, receipt),
                        );
                    }
                    umsh_mac::MacEventRef::AckTimeout { peer, receipt } => {
                        node.dispatch_ack_timeout(
                            peer,
                            crate::SendToken::new(identity_id, receipt),
                        );
                    }
                    umsh_mac::MacEventRef::Transmitted { wire_bytes, .. } => {
                        node.dispatch_transmitted(wire_bytes);
                    }
                    umsh_mac::MacEventRef::Forwarded { .. } => {}
                    // Ticket resolution (failed + finished) is handled by
                    // dispatch_ticket_state above; there is no peer to
                    // notify since the frame never aired.
                    umsh_mac::MacEventRef::TxAbandoned { .. } => {}
                }
            })
            .await
            .map_err(HostError::Mac)?;

        let queued: Vec<(LocalIdentityId, umsh_core::PublicKey, OwnedMacCommand)> =
            pending_pfs.borrow_mut().drain(..).collect();
        for (identity_id, from, command) in queued {
            self.handle_pfs_command(identity_id, from, command).await;
        }

        let identity_replies: Vec<(LocalNode<M>, IdentityResponsePlan)> =
            pending_identity.borrow_mut().drain(..).collect();
        for (node, plan) in identity_replies {
            node.send_identity_response(plan).await;
        }

        self.service_protocol_timeouts().await;

        Ok(())
    }

    /// Service node-layer deadlines independently of MAC/radio wake events.
    ///
    /// A quiet radio is not a MAC wake source, so applications that multiplex
    /// [`pump_once`](Self::pump_once) with other work must also call this from
    /// their own timer. In particular, ping and PFS request timeouts must fire
    /// even when the remote peer sends no response at all.
    pub async fn service_protocol_timeouts(&mut self) {
        let now_ms = self.mac.now_ms().await;
        service_node_timeouts(&self.nodes, now_ms).await;
    }

    /// Detach an owned handle that services the same node-layer deadlines as
    /// [`service_protocol_timeouts`](Self::service_protocol_timeouts).
    ///
    /// Use this when the host pump runs as its own long-lived future (so the
    /// `Host` is exclusively borrowed) and a sibling timer task must still
    /// fire ping/PFS timeouts — e.g. while the pump is parked awaiting a slow
    /// physical transmit. The handle snapshots the current node set; nodes
    /// added afterwards are not covered by it.
    pub fn protocol_timeout_servicer(&self) -> ProtocolTimeoutServicer<M>
    where
        M: Clone,
    {
        ProtocolTimeoutServicer {
            mac: self.mac.clone(),
            nodes: self.nodes.clone(),
        }
    }

    /// Run the shared MAC/Host loop forever.
    ///
    /// This is the preferred long-lived driver for node-based applications. It keeps the
    /// runtime wake policy inside the MAC/Host stack rather than requiring callers to write
    /// poll/sleep loops themselves.
    pub async fn run(&mut self) -> Result<(), HostError<M::RunError>> {
        loop {
            self.pump_once().await?;
        }
    }

    async fn handle_pfs_command(
        &mut self,
        identity_id: LocalIdentityId,
        from: umsh_core::PublicKey,
        command: OwnedMacCommand,
    ) {
        let Some(node) = self.route_node(identity_id) else {
            return;
        };

        match node
            .handle_pfs_command(&from, &command, &self.pfs_control_options)
            .await
        {
            Ok(Some(PfsLifecycle::Established(peer))) => node.dispatch_pfs_established(peer),
            Ok(Some(PfsLifecycle::Ended(peer))) => node.dispatch_pfs_ended(peer),
            Ok(None) => {}
            // Previously the error was dropped here, so a failed negotiation
            // (e.g. no ephemeral identity slot) left both sides silently stuck.
            // Surface it so applications can report the failure.
            Err(err) => node.dispatch_pfs_failed(from, err.pfs_failure()),
        }
    }
}

/// Owned deadline-service handle detached from a [`Host`]; see
/// [`Host::protocol_timeout_servicer`].
pub struct ProtocolTimeoutServicer<M: MacBackend> {
    mac: M,
    nodes: Vec<(LocalIdentityId, LocalNode<M>)>,
}

impl<M: MacBackend> ProtocolTimeoutServicer<M> {
    /// Fire any elapsed ping and PFS deadlines on the snapshot's nodes.
    pub async fn service(&self) {
        let now_ms = self.mac.now_ms().await;
        service_node_timeouts(&self.nodes, now_ms).await;
    }
}

async fn service_node_timeouts<M: MacBackend>(
    nodes: &[(LocalIdentityId, LocalNode<M>)],
    now_ms: u64,
) {
    #[cfg(feature = "software-crypto")]
    for (_, node) in nodes {
        if let Ok(expired) = node.expire_pfs_sessions().await {
            for peer in expired {
                node.dispatch_pfs_ended(peer);
            }
        }
        // Abandon sent PFS requests that were never answered, so the
        // requester reports a timeout instead of sitting in `Requested`.
        for peer in node.expire_pfs_requests(now_ms) {
            node.dispatch_pfs_failed(peer, crate::node::PfsFailure::Timeout);
        }
    }

    for (_, node) in nodes {
        node.expire_pings(now_ms);
    }
}

fn route_node<M: MacBackend>(
    nodes: &[(LocalIdentityId, LocalNode<M>)],
    identity_id: LocalIdentityId,
) -> Option<LocalNode<M>> {
    nodes
        .iter()
        .find(|(id, _)| *id == identity_id)
        .map(|(_, node)| node.clone())
        .or_else(|| {
            nodes
                .iter()
                .find(|(_, node)| node.owns_ephemeral_identity(identity_id))
                .map(|(_, node)| node.clone())
        })
}

fn dispatch_payload_callbacks<M: MacBackend>(
    node: &LocalNode<M>,
    packet: &ReceivedPacketRef<'_>,
    from: umsh_core::PublicKey,
    pending_pfs: &Rc<RefCell<Vec<(LocalIdentityId, umsh_core::PublicKey, OwnedMacCommand)>>>,
    pending_identity: &Rc<RefCell<Vec<(LocalNode<M>, IdentityResponsePlan)>>>,
    now_ms: u64,
) {
    if packet.payload_type() == PayloadType::NodeIdentity {
        if let Ok(identity) = NodeIdentityPayload::from_bytes(packet.payload()) {
            node.dispatch_node_discovered(from, identity.name.as_deref());
        }
        return;
    }

    if packet.payload_type() == PayloadType::MacCommand {
        if let Ok(command) = mac_command::parse(packet.payload()) {
            // Broadcast admits MAC commands one at a time: only a command
            // whose definition permits broadcast carriage is acted on, and
            // today that is the Identity Request alone. Anything else riding
            // a broadcast is dropped here, before observers see it.
            if packet.packet_family() == umsh_mac::PacketFamily::Broadcast
                && !matches!(command, mac_command::MacCommand::IdentityRequest { .. })
            {
                return;
            }
            // Match borrowed variants before converting to owned.
            if let mac_command::MacCommand::EchoResponse { data } = command {
                node.match_pong(
                    from,
                    data,
                    packet,
                    packet.received_at_ms().unwrap_or(now_ms),
                );
            }
            // Identity Request: let the built-in responder (if enabled) build a
            // reply plan now, while the reception context is live; the async
            // pump sends it after this synchronous dispatch returns.
            if let mac_command::MacCommand::IdentityRequest { options } = command {
                if let Some(plan) = node.evaluate_identity_request(packet, from, options) {
                    pending_identity.borrow_mut().push((node.clone(), plan));
                }
            }
            let owned = OwnedMacCommand::from(command);
            node.dispatch_mac_command(from, &owned);
            if matches!(
                owned,
                OwnedMacCommand::PfsSessionRequest { .. }
                    | OwnedMacCommand::PfsSessionResponse { .. }
                    | OwnedMacCommand::EndPfsSession
            ) {
                pending_pfs
                    .borrow_mut()
                    .push((node.identity_id(), from, owned));
            }
        }
    }
}