unshell/unshell-protocol/src/endpoint/routing.rs

use crate::{Endpoint, EndpointError, Packet, RouteDirection};

impl Endpoint {
    /// Register an inbound packet and route it through the local endpoint state.
    ///
    /// Inbound transport data still uses the same local routing rules as packets
    /// generated by leaves: local destinations are delivered to `inbound`, and
    /// transit destinations are queued by their immediate next hop.
    pub fn add_inbound(&mut self, packet: Packet) -> Result<(), EndpointError> {
        self.route_packet(packet)
    }

    /// Register an outbound packet produced locally and route it to the next queue.
    ///
    /// This intentionally shares the same implementation as [`Self::add_inbound`]
    /// so local leaf output and received transport packets cannot drift into subtly
    /// different route semantics.
    pub fn add_outbound(&mut self, packet: Packet) -> Result<(), EndpointError> {
        self.route_packet(packet)
    }

    /// Route a packet by classifying its destination and mutating exactly one queue.
    ///
    /// Hook cleanup is deliberately last. A packet with `end_hook = true` should not
    /// tear down local hook state unless the packet has a valid route and is actually
    /// queued for forwarding. The route branches are kept inline rather than using
    /// an intermediate decision enum so size-focused builds have less structure to
    /// optimize away.
    fn route_packet(&mut self, packet: Packet) -> Result<(), EndpointError> {
        self.ensure_path_is_set()?;

        if packet.path == self.path {
            let local_id = self
                .path
                .last()
                .copied()
                .ok_or(EndpointError::EndpointPathUnset)?;

            self.inbound.entry(local_id).or_default().push_back(packet);
            return Ok(());
        }

        // Direction is derived from the local path. The packet never gets to declare
        // whether it is moving upward, because that would make the trust boundary spoofable.
        if packet.path.starts_with(&self.path) {
            let next_hop = packet
                .path
                .get(self.path.len())
                .copied()
                .ok_or(EndpointError::DestinationOutsideLocalTree)?;

            self.ensure_registered_connection(next_hop, RouteDirection::Downward)?;
            self.queue_outbound(packet, next_hop, RouteDirection::Downward);
            return Ok(());
        }

        if self.path.starts_with(&packet.path) {
            // Upward-routed packets must be tied to local hook state. Otherwise a
            // peer could forge a packet to an ancestor by choosing an older path.
            if !self.hooks.contains_key(&packet.hook_id) {
                return Err(EndpointError::UnknownHook {
                    hook_id: packet.hook_id,
                });
            }

            let parent_index = self
                .path
                .len()
                .checked_sub(2)
                .ok_or(EndpointError::MissingParentRoute)?;

            let next_hop = self.path[parent_index];
            self.ensure_registered_connection(next_hop, RouteDirection::Upward)?;
            self.queue_outbound(packet, next_hop, RouteDirection::Upward);
            return Ok(());
        }

        Err(EndpointError::DestinationOutsideLocalTree)
    }

    /// Reject routing before path-relative decisions when no absolute path is known.
    ///
    /// This preserves the current runtime sentinel where an empty path means the
    /// endpoint has not been attached to the tree yet.
    fn ensure_path_is_set(&self) -> Result<(), EndpointError> {
        if self.path.is_empty() {
            Err(EndpointError::EndpointPathUnset)
        } else {
            Ok(())
        }
    }

    /// Verify that the derived adjacent endpoint is registered in this direction.
    ///
    /// The current connection table stores direction as a boolean. Keeping the bool
    /// conversion here confines that legacy representation to one place in routing.
    fn ensure_registered_connection(
        &self,
        next_hop: u32,
        direction: RouteDirection,
    ) -> Result<(), EndpointError> {
        let is_upward = matches!(direction, RouteDirection::Upward);

        if self.connections.contains(&(next_hop, is_upward)) {
            Ok(())
        } else {
            Err(EndpointError::MissingConnection {
                next_hop,
                direction,
            })
        }
    }

    /// Queue `packet` after all route validation has already succeeded.
    ///
    /// `end_hook` closes local hook state only when hook traffic is moving upward
    /// toward the hook host. Downward calls may carry a response hook id, but that
    /// id is only a promise for future upward traffic and must not delete local
    /// state if it happens to collide with an existing hook id.
    fn queue_outbound(&mut self, packet: Packet, next_hop: u32, direction: RouteDirection) {
        if matches!(direction, RouteDirection::Upward) && packet.end_hook {
            self.hooks.remove(&packet.hook_id);
        }

        self.outbound.entry(next_hop).or_default().push_back(packet);
    }
}