unshell/unshell-protocol/src/protocol/tree/endpoint/hooks.rs

//! Hook-state transitions and route helpers.

use alloc::string::String;

use crate::protocol::{
    DataMessage, FaultMessage, PacketHeader, PacketType, ProtocolFault, encode_packet,
};

use super::super::{HookKey, RouteDecision};
use super::core::{EndpointError, EndpointOutcome, Ingress, LocalEvent, ProtocolEndpoint};

impl ProtocolEndpoint {
    /// Returns the route that would carry a locally generated hook fault for `hook_id`.
    ///
    /// The method does not mutate hook state. Runtime owners use it to preflight transport
    /// availability before calling [`fail_hook`](Self::fail_hook), which removes hook state when
    /// the fault is emitted.
    #[must_use]
    pub fn hook_fault_route(&self, hook_id: u64) -> Option<RouteDecision> {
        self.hooks
            .key_for_hook_id(hook_id)
            .map(|key| self.decide_route(&key.return_path))
    }

    /// Terminates a locally known hook with a protocol fault.
    ///
    /// Unknown hooks are treated as an intentional drop. Known hooks are removed before the fault
    /// is routed so no further local data can be emitted after the terminal fault.
    pub fn fail_hook(
        &mut self,
        hook_id: u64,
        fault: ProtocolFault,
    ) -> Result<EndpointOutcome, EndpointError> {
        let key = self.hooks.key_for_hook_id(hook_id);
        self.emit_fault_if_possible(key, fault)
    }

    pub(crate) fn emit_fault_if_possible(
        &mut self,
        key: Option<HookKey>,
        fault: ProtocolFault,
    ) -> Result<EndpointOutcome, EndpointError> {
        let Some(key) = key else {
            return Ok(EndpointOutcome::Dropped);
        };

        self.hooks.remove_pending(&key);
        self.hooks.remove_active(&key);

        let header = PacketHeader {
            packet_type: PacketType::Fault,
            src_path: self.path.clone(),
            dst_path: key.return_path.clone(),
            dst_leaf: None,
            hook_id: Some(key.hook_id),
        };
        let message = FaultMessage { fault };

        match self.decide_route(&key.return_path) {
            RouteDecision::Local => Ok(EndpointOutcome::Local(LocalEvent::Fault {
                header,
                message,
                hook_key: key,
            })),
            route => Ok(EndpointOutcome::Forward {
                route,
                frame: encode_packet(&header, &message)?,
            }),
        }
    }

    pub(crate) fn handle_local_data(
        &mut self,
        header: PacketHeader,
        message: DataMessage,
    ) -> Result<EndpointOutcome, EndpointError> {
        let hook_id = header.hook_id.expect("validated");
        let key = if let Some(key) =
            self.hooks
                .resolve_active_key(&self.path, hook_id, &header.src_path)
        {
            key
        } else {
            let pending_key = HookKey::new(self.path.clone(), hook_id);
            if self.hooks.pending(&pending_key).is_some_and(|pending| {
                pending.caller_src_path == header.src_path
                    && pending.procedure_id == message.procedure_id
            }) {
                self.hooks.activate_pending(&pending_key);
                pending_key
            } else {
                return Ok(EndpointOutcome::Dropped);
            }
        };

        let Some(active) = self.hooks.active(&key) else {
            return Ok(EndpointOutcome::Dropped);
        };

        if active.peer_path != header.src_path {
            // A reused hook id from the wrong peer is treated as terminal for this hook,
            // because the endpoint can no longer trust future traffic on it.
            self.hooks.remove_active(&key);
            return self.emit_fault_if_possible(Some(key), ProtocolFault::INVALID_HOOK_PEER);
        }

        if active.procedure_id != message.procedure_id {
            // Data frames stay bound to the procedure chosen by the original call.
            // A procedure mismatch is dropped rather than faulted because the wrong peer may be
            // replaying stale traffic, and converting that into a terminal hook fault would let a
            // stray packet tear down an otherwise valid stream.
            return Ok(EndpointOutcome::Dropped);
        }

        if message.end_hook && self.hooks.mark_peer_end(&key) {
            self.hooks.remove_active(&key);
        }

        Ok(EndpointOutcome::Local(LocalEvent::Data {
            header,
            message,
            hook_key: key,
        }))
    }

    pub(crate) fn handle_local_fault(
        &mut self,
        header: PacketHeader,
        message: FaultMessage,
    ) -> Result<EndpointOutcome, EndpointError> {
        let hook_id = header.hook_id.expect("validated");
        if let Some(key) = self
            .hooks
            .resolve_active_key(&self.path, hook_id, &header.src_path)
        {
            self.hooks.remove_active(&key);
            return Ok(EndpointOutcome::Local(LocalEvent::Fault {
                header,
                message,
                hook_key: key,
            }));
        }

        let pending_key = HookKey::new(self.path.clone(), hook_id);
        if self
            .hooks
            .pending(&pending_key)
            .is_some_and(|pending| pending.caller_src_path == header.src_path)
        {
            self.hooks.remove_pending(&pending_key);
            return Ok(EndpointOutcome::Local(LocalEvent::Fault {
                header,
                message,
                hook_key: pending_key,
            }));
        }

        Ok(EndpointOutcome::Dropped)
    }

    /// Returns the current route decision for an absolute destination path.
    ///
    /// Runtime owners use this to validate transport availability before invoking
    /// endpoint operations that also mutate hook state.
    #[must_use]
    pub fn route_decision(&self, dst_path: &[String]) -> RouteDecision {
        self.routing.route(dst_path)
    }

    pub(crate) fn decide_route(&self, dst_path: &[String]) -> RouteDecision {
        self.route_decision(dst_path)
    }

    /// Returns whether one `src_path` is topologically valid for the ingress side that delivered
    /// the frame.
    ///
    /// Parent ingress may carry packets from ancestors, siblings, or the endpoint itself, but not
    /// from descendants pretending to be upstream. Child ingress may only carry packets from that
    /// child subtree, and local ingress must exactly match the endpoint path.
    pub(crate) fn valid_source_for_ingress(&self, ingress: &Ingress, src_path: &[String]) -> bool {
        match ingress {
            Ingress::Parent => {
                // Parent ingress may carry packets from ancestors, siblings, or the endpoint
                // itself, but not from descendants pretending to be upstream.
                if src_path.len() < self.path.len() {
                    return true;
                }
                if src_path.len() == self.path.len() {
                    return src_path == self.path;
                }
                !src_path.starts_with(&self.path)
            }
            Ingress::Child(child_path) => src_path.starts_with(child_path),
            Ingress::Local => src_path == self.path,
        }
    }
}