From aa1e9be69663da7b2f0a52da9a60732b2dc1f783 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 09:54:37 -0600 Subject: [PATCH 01/12] Redesign interface event ownership. --- LEAF_MACRO_INTERFACE.md | 26 ++ src/crypto/feistel.rs | 6 +- src/crypto/feistel_state.rs | 57 +++- src/crypto/sha256.rs | 2 +- src/interface/mod.rs | 3 + src/interface/store.rs | 271 ++++++++++++------ src/interface/target.rs | 46 +++ src/protocol/runtime.rs | 200 +++++++++---- unshell-leaves/leaf-pty/src/constants.rs | 3 + unshell-leaves/leaf-pty/src/lib.rs | 1 + unshell-leaves/leaf-pty/src/procedure.rs | 19 ++ unshell-leaves/leaf-pty/src/state.rs | 6 +- .../leaf-pty/src/tests/interface.rs | 230 +++++++++++++++ unshell-leaves/leaf-pty/src/tests/mod.rs | 5 + .../src/{tests.rs => tests/session.rs} | 234 +-------------- unshell-leaves/leaf-pty/src/tests/support.rs | 141 +++++++++ 16 files changed, 882 insertions(+), 368 deletions(-) create mode 100644 src/interface/target.rs create mode 100644 unshell-leaves/leaf-pty/src/procedure.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/interface.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/mod.rs rename unshell-leaves/leaf-pty/src/{tests.rs => tests/session.rs} (54%) create mode 100644 unshell-leaves/leaf-pty/src/tests/support.rs diff --git a/LEAF_MACRO_INTERFACE.md b/LEAF_MACRO_INTERFACE.md index c8f3d87..117aa21 100644 --- a/LEAF_MACRO_INTERFACE.md +++ b/LEAF_MACRO_INTERFACE.md @@ -106,6 +106,32 @@ Generated leaves receive an optional mutable store during `update_interface`. Th helpers create and update the appropriate session/procedure views when packets are dispatched, sessions update, and outbound routes succeed or fail. +Internally, interface events are target-driven: + +```text +generated runtime + knows packet owner + | + v +InterfaceTarget::Session(SessionKey) +InterfaceTarget::Procedure(ProcedureKey) + | + v +InterfaceStore::record(...) + append InterfaceEvent + link event index to exactly one view + update SessionViewStatus when applicable +``` + +This is deliberately not inferred from `Packet`. A PTY session packet and a one-shot +procedure packet both have `procedure_id` and `hook_id`, but they should not both +create session views. The runtime already knows which dispatch branch handled the +packet, so that answer is carried into the store. + +Leaf-level retry queues also carry the same owner metadata. That matters because the +shared leaf outbox contains both rejected session-init responses and procedure +responses. Session-entry outboxes use their surrounding session key directly. + Time remains caller-supplied: ```rust diff --git a/src/crypto/feistel.rs b/src/crypto/feistel.rs index 28c3b74..bdf7a35 100644 --- a/src/crypto/feistel.rs +++ b/src/crypto/feistel.rs @@ -19,10 +19,10 @@ pub fn feistel_shuffle(index: u16, seed: u32) -> u16 { .rotate_left(rot_amount) .wrapping_add(round.wrapping_mul(0x9E3779B9)); - // Round function F: Simple multiplicative hash mixing R and sub_key - // We cast to u32 for multiplication to avoid overflow, then mask back to 8 bits + // Round function F: Simple multiplicative hash mixing R and sub_key. + // Casting to u8 keeps the low byte, which is the half-block width here. let r_u32 = r as u32; - let hash_val = ((r_u32.wrapping_mul(sub_key)) ^ (r_u32 >> 4)) as u8 & 0xFF; + let hash_val = ((r_u32.wrapping_mul(sub_key)) ^ (r_u32 >> 4)) as u8; // Feistel step: New L = Old R, New R = Old L XOR F(R, key) let temp = l; diff --git a/src/crypto/feistel_state.rs b/src/crypto/feistel_state.rs index 5dabb27..aa94424 100644 --- a/src/crypto/feistel_state.rs +++ b/src/crypto/feistel_state.rs @@ -1,12 +1,38 @@ use crate::crypto::feistel_shuffle; -#[cfg(feature = "counter_shuffle_none")] +/// Counter implementation selected by feature flags. +/// +/// Cargo's `--all-features` enables every counter strategy at once, so these cfgs are +/// intentionally priority-ordered instead of mutually exclusive aliases. The strongest +/// configured shuffle wins: Feistel+LCG, then Feistel, then the linear fallback. +#[cfg(all( + feature = "counter_shuffle_none", + not(any( + feature = "counter_shuffle_feistel", + feature = "counter_shuffle_feistel_lcg" + )) +))] pub type Counter = NoShuffle; -#[cfg(feature = "counter_shuffle_feistel")] + +/// Counter implementation selected when Feistel is enabled without Feistel+LCG. +#[cfg(all( + feature = "counter_shuffle_feistel", + not(feature = "counter_shuffle_feistel_lcg") +))] pub type Counter = FeistelShuffle; + +/// Default and strongest counter implementation. #[cfg(feature = "counter_shuffle_feistel_lcg")] pub type Counter = FeistelLCGShuffle; +/// Fallback used only when all counter shuffle features are disabled. +#[cfg(not(any( + feature = "counter_shuffle_none", + feature = "counter_shuffle_feistel", + feature = "counter_shuffle_feistel_lcg" +)))] +pub type Counter = NoShuffle; + const NONCE16_1: u16 = const_random::const_random!(u16); const NONCE16_2: u16 = const_random::const_random!(u16); const NONCE32: u32 = const_random::const_random!(u32); @@ -27,12 +53,21 @@ impl NoShuffle { Self(NONCE16_1) } + // This is an id generator API, not an iterator: callers need a bare `u16` and no + // exhaustion state because the counter intentionally wraps through the full space. + #[allow(clippy::should_implement_trait)] pub fn next(&mut self) -> u16 { self.0 = self.0.wrapping_add(1); self.0 } } +impl Default for NoShuffle { + fn default() -> Self { + Self::new() + } +} + /// Shuffle all 16 bit numbers, an actual shuffle /// But this still stores local values in a linear format pub struct FeistelShuffle(u16, u32); @@ -42,12 +77,21 @@ impl FeistelShuffle { Self(NONCE16_1, NONCE32) } + // This is an id generator API, not an iterator: callers need a bare `u16` and no + // exhaustion state because the counter intentionally wraps through the full space. + #[allow(clippy::should_implement_trait)] pub fn next(&mut self) -> u16 { self.0 = self.0.wrapping_add(FEISTEL_STEP); feistel_shuffle(self.0, self.1) } } +impl Default for FeistelShuffle { + fn default() -> Self { + Self::new() + } +} + /// Linear recursive shuffle, /// feeds back into itself and doesn't store the actual state. /// Harder to decompile @@ -65,6 +109,9 @@ impl FeistelLCGShuffle { Self { state: 0, a, c } } + // This is an id generator API, not an iterator: callers need a bare `u16` and no + // exhaustion state because the counter intentionally wraps through the full space. + #[allow(clippy::should_implement_trait)] pub fn next(&mut self) -> u16 { // 1. Advance state using LCG (Guarantees single cycle of 65536) self.state = self.state.wrapping_mul(self.a).wrapping_add(self.c); @@ -73,3 +120,9 @@ impl FeistelLCGShuffle { feistel_shuffle(self.state, self.a as u32) } } + +impl Default for FeistelLCGShuffle { + fn default() -> Self { + Self::new() + } +} diff --git a/src/crypto/sha256.rs b/src/crypto/sha256.rs index 91562d0..66826ab 100644 --- a/src/crypto/sha256.rs +++ b/src/crypto/sha256.rs @@ -107,7 +107,7 @@ const fn compress(state: &mut [u32; 8], block: &[u8; 64]) { /// Returns the SHA-256 digest of `input` as 32 raw bytes. pub const fn sha256(input: &[u8]) -> [u8; 32] { // Padded length is the next multiple of 64 that fits input + 1 (0x80) + 8 (length). - let padded_len = ((input.len() + 9 + 63) / 64) * 64; + let padded_len = (input.len() + 9).div_ceil(64) * 64; let mut state = H; let mut block_start = 0; diff --git a/src/interface/mod.rs b/src/interface/mod.rs index 923e381..513bcf9 100644 --- a/src/interface/mod.rs +++ b/src/interface/mod.rs @@ -6,9 +6,12 @@ mod event; mod key; mod store; +mod target; mod view; pub use event::{InterfaceEvent, InterfaceEventKind}; pub use key::{ProcedureKey, SessionKey}; pub use store::InterfaceStore; pub use view::{ProcedureView, SessionView, SessionViewStatus}; + +pub(crate) use target::InterfaceTarget; diff --git a/src/interface/store.rs b/src/interface/store.rs index 0318303..d979648 100644 --- a/src/interface/store.rs +++ b/src/interface/store.rs @@ -2,8 +2,8 @@ use alloc::{collections::BTreeMap, vec::Vec}; use crate::{ interface::{ - InterfaceEvent, InterfaceEventKind, ProcedureKey, ProcedureView, SessionKey, SessionView, - SessionViewStatus, + InterfaceEvent, InterfaceEventKind, InterfaceTarget, ProcedureKey, ProcedureView, + SessionKey, SessionView, SessionViewStatus, }, protocol::{EndpointError, HookID, Packet, SessionStatus}, }; @@ -15,7 +15,6 @@ use crate::{ /// itself stays with the renderer or application shell so protocol state remains /// headless and reusable. pub struct InterfaceStore { - next_sequence: u64, now_ns: Option, events: Vec, sessions: BTreeMap, @@ -26,7 +25,6 @@ impl InterfaceStore { /// Creates an empty caller-owned interface store. pub fn new() -> Self { Self { - next_sequence: 0, now_ns: None, events: Vec::new(), sessions: BTreeMap::new(), @@ -70,30 +68,32 @@ impl InterfaceStore { procedure_id: u32, hook_id: HookID, ) -> &mut SessionView { - self.sessions - .entry(SessionKey { - leaf_id, - procedure_id, - hook_id, - }) - .or_insert_with(SessionView::new) + self.session_view_for_key_mut(SessionKey { + leaf_id, + procedure_id, + hook_id, + }) } /// Returns or creates the view for a one-shot procedure family. pub fn procedure_view_mut(&mut self, leaf_id: u32, procedure_id: u32) -> &mut ProcedureView { - self.procedures - .entry(ProcedureKey { - leaf_id, - procedure_id, - }) - .or_insert_with(ProcedureView::new) + self.procedure_view_for_key_mut(ProcedureKey { + leaf_id, + procedure_id, + }) } /// Records a packet delivered to a generated leaf. pub fn record_inbound(&mut self, leaf_id: u32, packet: &Packet) { - self.push_packet_event( - leaf_id, - packet, + let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); + self.record_inbound_for(target, packet); + } + + /// Records a packet delivered to a target already known by generated runtime code. + pub(crate) fn record_inbound_for(&mut self, target: InterfaceTarget, packet: &Packet) { + self.record( + target, + None, InterfaceEventKind::Inbound { packet: packet.clone(), }, @@ -107,14 +107,25 @@ impl InterfaceStore { procedure_id: u32, hook_id: HookID, ) { - self.push_session_event( + self.record_session_packet_queued_for(InterfaceTarget::session( leaf_id, procedure_id, hook_id, + )); + } + + /// Records that a packet was queued for an existing session inbox. + pub(crate) fn record_session_packet_queued_for(&mut self, target: InterfaceTarget) { + let InterfaceTarget::Session(key) = target else { + return; + }; + + self.record( + target, None, InterfaceEventKind::SessionPacketQueued { - procedure_id, - hook_id, + procedure_id: key.procedure_id, + hook_id: key.hook_id, }, ); } @@ -127,14 +138,26 @@ impl InterfaceStore { hook_id: HookID, started_ns: Option, ) { - self.push_session_event( - leaf_id, - procedure_id, - hook_id, + let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); + self.record_session_created_for(target, started_ns); + } + + /// Records successful creation of a new session state for an explicit target. + pub(crate) fn record_session_created_for( + &mut self, + target: InterfaceTarget, + started_ns: Option, + ) { + let InterfaceTarget::Session(key) = target else { + return; + }; + + self.record( + target, Some(SessionViewStatus::Running), InterfaceEventKind::SessionCreated { - procedure_id, - hook_id, + procedure_id: key.procedure_id, + hook_id: key.hook_id, started_ns, finished_ns: self.now_ns, }, @@ -149,14 +172,26 @@ impl InterfaceStore { hook_id: HookID, started_ns: Option, ) { - self.push_session_event( - leaf_id, - procedure_id, - hook_id, + let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); + self.record_session_rejected_for(target, started_ns); + } + + /// Records rejection of a packet that could not create a session. + pub(crate) fn record_session_rejected_for( + &mut self, + target: InterfaceTarget, + started_ns: Option, + ) { + let InterfaceTarget::Session(key) = target else { + return; + }; + + self.record( + target, Some(SessionViewStatus::Rejected), InterfaceEventKind::SessionRejected { - procedure_id, - hook_id, + procedure_id: key.procedure_id, + hook_id: key.hook_id, started_ns, finished_ns: self.now_ns, }, @@ -172,14 +207,27 @@ impl InterfaceStore { status: SessionStatus, started_ns: Option, ) { - self.push_session_event( - leaf_id, - procedure_id, - hook_id, + let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); + self.record_session_update_for(target, status, started_ns); + } + + /// Records one session update tick for an explicit session target. + pub(crate) fn record_session_update_for( + &mut self, + target: InterfaceTarget, + status: SessionStatus, + started_ns: Option, + ) { + let InterfaceTarget::Session(key) = target else { + return; + }; + + self.record( + target, Some(SessionViewStatus::from_session_status(status)), InterfaceEventKind::SessionUpdated { - procedure_id, - hook_id, + procedure_id: key.procedure_id, + hook_id: key.hook_id, status, started_ns, finished_ns: self.now_ns, @@ -195,11 +243,29 @@ impl InterfaceStore { hook_id: HookID, started_ns: Option, ) { - self.push_procedure_event( - leaf_id, - procedure_id, + self.record_procedure_call_for( + InterfaceTarget::procedure(leaf_id, procedure_id), + hook_id, + started_ns, + ); + } + + /// Records one procedure call for an explicit procedure target. + pub(crate) fn record_procedure_call_for( + &mut self, + target: InterfaceTarget, + hook_id: HookID, + started_ns: Option, + ) { + let InterfaceTarget::Procedure(key) = target else { + return; + }; + + self.record( + target, + None, InterfaceEventKind::ProcedureCalled { - procedure_id, + procedure_id: key.procedure_id, hook_id, started_ns, finished_ns: self.now_ns, @@ -209,9 +275,15 @@ impl InterfaceStore { /// Records a packet emitted by leaf logic before route retry handling. pub fn record_outbound_queued(&mut self, leaf_id: u32, packet: &Packet) { - self.push_packet_event( - leaf_id, - packet, + let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); + self.record_outbound_queued_for(target, packet); + } + + /// Records a packet emitted by leaf logic before route retry handling. + pub(crate) fn record_outbound_queued_for(&mut self, target: InterfaceTarget, packet: &Packet) { + self.record( + target, + None, InterfaceEventKind::OutboundQueued { packet: packet.clone(), }, @@ -220,9 +292,15 @@ impl InterfaceStore { /// Records a route attempt for a queued outbound packet. pub fn record_route_attempt(&mut self, leaf_id: u32, packet: &Packet) { - self.push_packet_event( - leaf_id, - packet, + let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); + self.record_route_attempt_for(target, packet); + } + + /// Records a route attempt for a queued outbound packet. + pub(crate) fn record_route_attempt_for(&mut self, target: InterfaceTarget, packet: &Packet) { + self.record( + target, + None, InterfaceEventKind::RouteAttempt { packet: packet.clone(), }, @@ -231,9 +309,15 @@ impl InterfaceStore { /// Records a successful route attempt. pub fn record_route_success(&mut self, leaf_id: u32, packet: &Packet) { - self.push_packet_event( - leaf_id, - packet, + let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); + self.record_route_success_for(target, packet); + } + + /// Records a successful route attempt. + pub(crate) fn record_route_success_for(&mut self, target: InterfaceTarget, packet: &Packet) { + self.record( + target, + None, InterfaceEventKind::RouteSuccess { packet: packet.clone(), }, @@ -242,9 +326,20 @@ impl InterfaceStore { /// Records a failed route attempt without removing the packet from retry state. pub fn record_route_failure(&mut self, leaf_id: u32, packet: &Packet, error: EndpointError) { - self.push_packet_event( - leaf_id, - packet, + let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); + self.record_route_failure_for(target, packet, error); + } + + /// Records a failed route attempt without removing the packet from retry state. + pub(crate) fn record_route_failure_for( + &mut self, + target: InterfaceTarget, + packet: &Packet, + error: EndpointError, + ) { + self.record( + target, + None, InterfaceEventKind::RouteFailure { packet: packet.clone(), error, @@ -252,43 +347,43 @@ impl InterfaceStore { ); } - fn push_packet_event(&mut self, leaf_id: u32, packet: &Packet, kind: InterfaceEventKind) { - let index = self.push_event(leaf_id, kind); - self.link_packet_event(leaf_id, packet, index); - } - - fn push_session_event( + fn record( &mut self, - leaf_id: u32, - procedure_id: u32, - hook_id: HookID, + target: InterfaceTarget, status: Option, kind: InterfaceEventKind, ) { - let index = self.push_event(leaf_id, kind); - let view = self.session_view_mut(leaf_id, procedure_id, hook_id); - - if let Some(status) = status { - view.status = status; - } - - view.events.push(index); + let index = self.push_event(target.leaf_id(), kind); + self.link_event(target, status, index); } - fn push_procedure_event(&mut self, leaf_id: u32, procedure_id: u32, kind: InterfaceEventKind) { - let index = self.push_event(leaf_id, kind); - self.procedure_view_mut(leaf_id, procedure_id) - .events - .push(index); + fn link_event( + &mut self, + target: InterfaceTarget, + status: Option, + index: usize, + ) { + match target { + InterfaceTarget::Session(key) => { + let view = self.session_view_for_key_mut(key); + + if let Some(status) = status { + view.status = status; + } + + view.events.push(index); + } + InterfaceTarget::Procedure(key) => { + self.procedure_view_for_key_mut(key).events.push(index); + } + } } fn push_event(&mut self, leaf_id: u32, kind: InterfaceEventKind) -> usize { - let sequence = self.next_sequence; - self.next_sequence = self.next_sequence.wrapping_add(1); let index = self.events.len(); self.events.push(InterfaceEvent { - sequence, + sequence: index as u64, time_ns: self.now_ns, leaf_id, kind, @@ -297,10 +392,14 @@ impl InterfaceStore { index } - fn link_packet_event(&mut self, leaf_id: u32, packet: &Packet, index: usize) { - self.session_view_mut(leaf_id, packet.procedure_id, packet.hook_id) - .events - .push(index); + fn session_view_for_key_mut(&mut self, key: SessionKey) -> &mut SessionView { + self.sessions.entry(key).or_insert_with(SessionView::new) + } + + fn procedure_view_for_key_mut(&mut self, key: ProcedureKey) -> &mut ProcedureView { + self.procedures + .entry(key) + .or_insert_with(ProcedureView::new) } } diff --git a/src/interface/target.rs b/src/interface/target.rs new file mode 100644 index 0000000..bcedf82 --- /dev/null +++ b/src/interface/target.rs @@ -0,0 +1,46 @@ +use crate::{ + interface::{ProcedureKey, SessionKey}, + protocol::HookID, +}; + +/// Internal owner for one interface event. +/// +/// The runtime already knows whether a packet belongs to a hook-backed session or a +/// one-shot procedure. Keeping that answer explicit avoids reconstructing ownership +/// from packet fields later, which is what made procedure packet flow look like fake +/// session activity in the previous store implementation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum InterfaceTarget { + /// Event belongs to one hook-backed session instance. + Session(SessionKey), + + /// Event belongs to one one-shot procedure family. + Procedure(ProcedureKey), +} + +impl InterfaceTarget { + /// Builds a session target from the same pieces exposed by [`SessionKey`]. + pub(crate) fn session(leaf_id: u32, procedure_id: u32, hook_id: HookID) -> Self { + Self::Session(SessionKey { + leaf_id, + procedure_id, + hook_id, + }) + } + + /// Builds a procedure target from the same pieces exposed by [`ProcedureKey`]. + pub(crate) fn procedure(leaf_id: u32, procedure_id: u32) -> Self { + Self::Procedure(ProcedureKey { + leaf_id, + procedure_id, + }) + } + + /// Returns the leaf id used on the append-only event record. + pub(crate) fn leaf_id(self) -> u32 { + match self { + Self::Session(key) => key.leaf_id, + Self::Procedure(key) => key.leaf_id, + } + } +} diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index ca9cb1b..9f684d5 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -1,5 +1,7 @@ +use alloc::collections::VecDeque; + use crate::{ - interface::InterfaceStore, + interface::{InterfaceStore, InterfaceTarget}, protocol::{ Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionCtx, SessionEntry, SessionFamily, SessionInit, SessionInitResult, SessionStatus, @@ -12,25 +14,49 @@ use crate::{ /// session initialization responses and one-shot procedures, both of which need the /// same retry semantics as session output without becoming separate framework types. pub struct LeafOutbox { - packets: PacketQueue, + packets: VecDeque, +} + +/// One packet retained by a leaf-level retry queue. +/// +/// Session entry outboxes have an obvious owner from their surrounding session entry. +/// Leaf-level outboxes are mixed: rejected session initialization packets and one-shot +/// procedure responses both land here. Storing the owner beside the packet keeps route +/// logging precise without exposing another public queue type. +#[derive(Clone)] +struct LeafOutboxEntry { + packet: Packet, + target: LeafOutboxTarget, +} + +/// Interface owner attached to a leaf-level outbox entry. +#[derive(Clone, Copy)] +enum LeafOutboxTarget { + /// Compatibility path for packets queued through the public `push`/`extend` API. + InferFromPacket, + + /// Runtime-known session or procedure target. + Explicit(InterfaceTarget), } impl LeafOutbox { /// Creates an empty leaf-level outbox. pub fn new() -> Self { Self { - packets: PacketQueue::new(), + packets: VecDeque::new(), } } /// Adds one packet to the retry queue. pub fn push(&mut self, packet: Packet) { - self.packets.push_back(packet); + self.push_with_target(packet, LeafOutboxTarget::InferFromPacket); } /// Adds all packets from `packets` in FIFO order. pub fn extend(&mut self, packets: PacketQueue) { - self.packets.extend(packets); + for packet in packets { + self.push(packet); + } } /// Returns the number of queued packets. @@ -42,6 +68,22 @@ impl LeafOutbox { pub fn is_empty(&self) -> bool { self.packets.is_empty() } + + /// Adds one packet with a runtime-known interface target. + pub(crate) fn push_for_target(&mut self, packet: Packet, target: InterfaceTarget) { + self.push_with_target(packet, LeafOutboxTarget::Explicit(target)); + } + + /// Adds all packets with the same runtime-known interface target. + pub(crate) fn extend_for_target(&mut self, packets: PacketQueue, target: InterfaceTarget) { + for packet in packets { + self.push_for_target(packet, target); + } + } + + fn push_with_target(&mut self, packet: Packet, target: LeafOutboxTarget) { + self.packets.push_back(LeafOutboxEntry { packet, target }); + } } impl Default for LeafOutbox { @@ -67,9 +109,10 @@ pub fn dispatch_session( { let hook_id = packet.hook_id; let procedure_id = S::PROCEDURE_ID; + let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); if let Some(store) = interface.as_mut() { - store.record_inbound(leaf_id, &packet); + store.record_inbound_for(target, &packet); } if let Some(entry) = family @@ -80,7 +123,7 @@ pub fn dispatch_session( entry.inbox.push_back(packet); if let Some(store) = interface.as_mut() { - store.record_session_packet_queued(leaf_id, procedure_id, hook_id); + store.record_session_packet_queued_for(target); } return; @@ -95,21 +138,21 @@ pub fn dispatch_session( family.entries.push(SessionEntry::new(hook_id, state)); if let Some(store) = interface.as_mut() { - store.record_session_created(leaf_id, procedure_id, hook_id, started_ns); + store.record_session_created_for(target, started_ns); } } SessionInitResult::Rejected => { if let Some(store) = interface.as_mut() { - store.record_session_rejected(leaf_id, procedure_id, hook_id, started_ns); + store.record_session_rejected_for(target, started_ns); } } SessionInitResult::RejectedWith(packet) => { if let Some(store) = interface.as_mut() { - store.record_session_rejected(leaf_id, procedure_id, hook_id, started_ns); - store.record_outbound_queued(leaf_id, &packet); + store.record_session_rejected_for(target, started_ns); + store.record_outbound_queued_for(target, &packet); } - outbox.push(packet); + outbox.push_for_target(packet, target); } } } @@ -129,23 +172,26 @@ pub fn update_session_family( } let started_ns = interface.as_ref().and_then(|store| store.now_ns()); + let outbox_start = entry.outbox.len(); let reply_path = S::reply_path(&entry.state).to_vec(); - let mut ctx = SessionCtx::new( - entry.hook_id, - reply_path, - S::PROCEDURE_ID, - &mut entry.outbox, - ); - let status = S::update(leaf, &mut entry.state, &mut entry.inbox, &mut ctx); + let status = { + let mut ctx = SessionCtx::new( + entry.hook_id, + reply_path, + S::PROCEDURE_ID, + &mut entry.outbox, + ); + + S::update(leaf, &mut entry.state, &mut entry.inbox, &mut ctx) + }; + let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); if let Some(store) = interface.as_mut() { - store.record_session_update( - leaf_id, - S::PROCEDURE_ID, - entry.hook_id, - status, - started_ns, - ); + store.record_session_update_for(target, status, started_ns); + + for packet in entry.outbox.iter().skip(outbox_start) { + store.record_outbound_queued_for(target, packet); + } } if matches!(status, SessionStatus::Closed) { @@ -166,9 +212,10 @@ pub fn dispatch_procedure( P: Procedure, { let started_ns = interface.as_ref().and_then(|store| store.now_ns()); + let target = InterfaceTarget::procedure(leaf_id, P::PROCEDURE_ID); if let Some(store) = interface.as_mut() { - store.record_inbound(leaf_id, &packet); + store.record_inbound_for(target, &packet); } let hook_id = packet.hook_id; @@ -180,14 +227,14 @@ pub fn dispatch_procedure( let packets = procedure_out.into_packets(); if let Some(store) = interface.as_mut() { - store.record_procedure_call(leaf_id, P::PROCEDURE_ID, hook_id, started_ns); + store.record_procedure_call_for(target, hook_id, started_ns); for packet in &packets { - store.record_outbound_queued(leaf_id, packet); + store.record_outbound_queued_for(target, packet); } } - outbox.extend(packets); + outbox.extend_for_target(packets, target); } /// Flushes a generated leaf-level outbox through endpoint routing. @@ -197,7 +244,17 @@ pub fn flush_leaf_outbox( outbox: &mut LeafOutbox, interface: &mut Option<&mut InterfaceStore>, ) -> bool { - flush_packet_queue_with_interface(endpoint, leaf_id, &mut outbox.packets, interface) + while let Some(entry) = outbox.packets.front().cloned() { + let target = resolve_leaf_outbox_target(leaf_id, &entry); + + if !flush_packet_with_target(endpoint, target, &entry.packet, interface) { + return false; + } + + outbox.packets.pop_front(); + } + + true } /// Flushes and retains one generated session family. @@ -210,7 +267,8 @@ pub fn flush_session_family( S: Session, { for entry in &mut family.entries { - flush_packet_queue_with_interface(endpoint, leaf_id, &mut entry.outbox, interface); + let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); + flush_packet_queue_with_target(endpoint, target, &mut entry.outbox, interface); } family @@ -230,31 +288,73 @@ pub fn flush_packet_queue_with_interface( interface: &mut Option<&mut InterfaceStore>, ) -> bool { while let Some(packet) = outbox.front().cloned() { - if let Some(store) = interface.as_mut() { - store.record_route_attempt(leaf_id, &packet); + let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); + + if !flush_packet_with_target(endpoint, target, &packet, interface) { + return false; } - match endpoint.add_outbound(packet.clone()) { - Ok(()) => { - if let Some(store) = interface.as_mut() { - store.record_route_success(leaf_id, &packet); - } - - outbox.pop_front(); - } - Err(error) => { - if let Some(store) = interface.as_mut() { - store.record_route_failure(leaf_id, &packet, error); - } - - return false; - } - } + outbox.pop_front(); } true } +/// Flushes a packet queue whose owner is already known by the generated runtime. +fn flush_packet_queue_with_target( + endpoint: &mut Endpoint, + target: InterfaceTarget, + outbox: &mut PacketQueue, + interface: &mut Option<&mut InterfaceStore>, +) -> bool { + while let Some(packet) = outbox.front().cloned() { + if !flush_packet_with_target(endpoint, target, &packet, interface) { + return false; + } + + outbox.pop_front(); + } + + true +} + +fn flush_packet_with_target( + endpoint: &mut Endpoint, + target: InterfaceTarget, + packet: &Packet, + interface: &mut Option<&mut InterfaceStore>, +) -> bool { + if let Some(store) = interface.as_mut() { + store.record_route_attempt_for(target, packet); + } + + match endpoint.add_outbound(packet.clone()) { + Ok(()) => { + if let Some(store) = interface.as_mut() { + store.record_route_success_for(target, packet); + } + + true + } + Err(error) => { + if let Some(store) = interface.as_mut() { + store.record_route_failure_for(target, packet, error); + } + + false + } + } +} + +fn resolve_leaf_outbox_target(leaf_id: u32, entry: &LeafOutboxEntry) -> InterfaceTarget { + match entry.target { + LeafOutboxTarget::InferFromPacket => { + InterfaceTarget::session(leaf_id, entry.packet.procedure_id, entry.packet.hook_id) + } + LeafOutboxTarget::Explicit(target) => target, + } +} + /// Returns the path used by generated procedure responses. fn parent_reply_path(endpoint: &Endpoint) -> alloc::vec::Vec { if endpoint.path.len() > 1 { diff --git a/unshell-leaves/leaf-pty/src/constants.rs b/unshell-leaves/leaf-pty/src/constants.rs index cfce7fd..71b2009 100644 --- a/unshell-leaves/leaf-pty/src/constants.rs +++ b/unshell-leaves/leaf-pty/src/constants.rs @@ -6,6 +6,9 @@ pub const LEAF_FAKE_PTY: u32 = hash_32!("dev.unshell.v1.pty"); /// Outer procedure id used by all fake PTY session packets. pub const PROC_PTY: u32 = hash_32!("dev.unshell.v1.pty.pty"); +/// One-shot procedure id used by tests to prove procedure interface ownership. +pub(crate) const PROC_PING: u32 = hash_32!("dev.unshell.v1.pty.ping"); + /// Downward opcode that opens one PTY session. pub const OP_OPEN: u8 = 0; diff --git a/unshell-leaves/leaf-pty/src/lib.rs b/unshell-leaves/leaf-pty/src/lib.rs index ab00097..d040da0 100644 --- a/unshell-leaves/leaf-pty/src/lib.rs +++ b/unshell-leaves/leaf-pty/src/lib.rs @@ -11,6 +11,7 @@ extern crate alloc; mod codec; mod constants; +mod procedure; mod session; mod state; diff --git a/unshell-leaves/leaf-pty/src/procedure.rs b/unshell-leaves/leaf-pty/src/procedure.rs new file mode 100644 index 0000000..7de6f19 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/procedure.rs @@ -0,0 +1,19 @@ +use unshell::protocol::{Endpoint, Packet, Procedure, ProcedureOut}; + +use crate::{constants::PROC_PING, state::FakePtyState}; + +/// One-shot echo procedure used to exercise generated procedure dispatch. +/// +/// The fake PTY leaf is primarily session-oriented, so this deliberately small +/// procedure gives tests a non-session packet family. That keeps interface logging +/// honest: procedure packets should populate [`unshell::interface::ProcedureView`] +/// instead of being inferred as hook-backed sessions. +pub(crate) struct PingProcedure; + +impl Procedure for PingProcedure { + const PROCEDURE_ID: u32 = PROC_PING; + + fn handle(_: &mut FakePtyState, _: &mut Endpoint, packet: Packet, out: &mut ProcedureOut) { + out.send_final(&packet.data); + } +} diff --git a/unshell-leaves/leaf-pty/src/state.rs b/unshell-leaves/leaf-pty/src/state.rs index f73cfd7..92cdf53 100644 --- a/unshell-leaves/leaf-pty/src/state.rs +++ b/unshell-leaves/leaf-pty/src/state.rs @@ -1,6 +1,6 @@ use unshell::protocol::{HookID, unshell_leaf}; -use crate::{constants::LEAF_FAKE_PTY, session::PtySession}; +use crate::{constants::LEAF_FAKE_PTY, procedure::PingProcedure, session::PtySession}; /// User-owned state for the generated fake PTY leaf. /// @@ -47,6 +47,8 @@ unshell_leaf! { sessions { pty: PtySession, } - procedures {} + procedures { + ping: PingProcedure, + } } } diff --git a/unshell-leaves/leaf-pty/src/tests/interface.rs b/unshell-leaves/leaf-pty/src/tests/interface.rs new file mode 100644 index 0000000..a3977c4 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/interface.rs @@ -0,0 +1,230 @@ +use alloc::vec; + +use unshell::{ + interface::{InterfaceEventKind, InterfaceStore, ProcedureKey, SessionKey, SessionViewStatus}, + protocol::{Leaf, Packet}, +}; + +use crate::{ + FakePtyLeaf, FakePtyState, OP_EXIT, OP_OPENED, OP_TERMINATE, PROC_PTY, constants::PROC_PING, + frame_opcode, pty_open_packet, +}; + +use super::support::{ + ENDPOINT_A, ENDPOINT_B, assert_frame, drain_parent_packets, drain_parent_pty_packets, + pty_endpoints, send_downward_frame, transfer_packets, +}; + +fn view_has_event(interface: &InterfaceStore, event_indexes: &[usize], mut predicate: F) -> bool +where + F: FnMut(&InterfaceEventKind) -> bool, +{ + event_indexes + .iter() + .any(|index| predicate(&interface.events()[*index].kind)) +} + +fn send_downward_ping( + endpoint_a: &mut unshell::protocol::Endpoint, + endpoint_b: &mut unshell::protocol::Endpoint, + hook_id: u16, + payload: &[u8], +) { + endpoint_a + .add_outbound(Packet { + hook_id, + end_hook: false, + path: vec![ENDPOINT_A, ENDPOINT_B], + procedure_id: PROC_PING, + data: payload.to_vec(), + }) + .unwrap(); + + transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); +} + +#[test] +fn interface_update_records_session_flow() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let mut interface = InterfaceStore::new(); + let hook_id = endpoint_a.get_hook_id(); + + endpoint_a + .add_outbound(pty_open_packet( + vec![ENDPOINT_A, ENDPOINT_B], + hook_id, + &[ENDPOINT_A], + )) + .unwrap(); + transfer_packets(&mut endpoint_a, &mut endpoint_b, ENDPOINT_B, ENDPOINT_A); + + leaf.update_interface(&mut endpoint_b, &mut interface); + + let session_key = SessionKey { + leaf_id: leaf.get_id(), + procedure_id: PROC_PTY, + hook_id, + }; + let session_view = interface.session_views().get(&session_key).unwrap(); + + assert_eq!(leaf.active_session_count(), 1); + assert!(view_has_event( + &interface, + &session_view.events, + |event| matches!( + event, + InterfaceEventKind::SessionCreated { hook_id: recorded_hook, .. } + if *recorded_hook == hook_id + ), + )); + assert!(view_has_event( + &interface, + &session_view.events, + |event| matches!( + event, + InterfaceEventKind::OutboundQueued { packet } + if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_OPENED) + ), + )); + assert!(view_has_event( + &interface, + &session_view.events, + |event| matches!( + event, + InterfaceEventKind::RouteSuccess { packet } + if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_OPENED) + ), + )); +} + +#[test] +fn interface_update_records_failed_final_route_without_dropping_session() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let mut interface = InterfaceStore::new(); + let hook_id = endpoint_a.get_hook_id(); + + endpoint_a + .add_outbound(pty_open_packet( + vec![ENDPOINT_A, ENDPOINT_B], + hook_id, + &[ENDPOINT_A], + )) + .unwrap(); + transfer_packets(&mut endpoint_a, &mut endpoint_b, ENDPOINT_B, ENDPOINT_A); + leaf.update_interface(&mut endpoint_b, &mut interface); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_TERMINATE, + &[], + false, + ); + endpoint_b.connections.remove(&(ENDPOINT_A, true)); + leaf.update_interface(&mut endpoint_b, &mut interface); + + let session_key = SessionKey { + leaf_id: leaf.get_id(), + procedure_id: PROC_PTY, + hook_id, + }; + let session_view = interface.session_views().get(&session_key).unwrap(); + + assert_eq!(leaf.active_session_count(), 1); + assert_eq!(leaf.pending_packet_count(), 1); + assert_eq!(session_view.status, SessionViewStatus::Closed); + assert!(view_has_event( + &interface, + &session_view.events, + |event| matches!( + event, + InterfaceEventKind::RouteFailure { packet, .. } + if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_EXIT) + ), + )); + + endpoint_b.connections.insert((ENDPOINT_A, true)); + leaf.update_interface(&mut endpoint_b, &mut interface); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + let session_view = interface.session_views().get(&session_key).unwrap(); + + assert_eq!(leaf.active_session_count(), 0); + assert_eq!(packets.len(), 1); + assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); + assert!(view_has_event( + &interface, + &session_view.events, + |event| matches!( + event, + InterfaceEventKind::RouteSuccess { packet } + if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_EXIT) + ), + )); +} + +#[test] +fn interface_update_records_procedure_flow_without_session_view() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let mut interface = InterfaceStore::new(); + let hook_id = endpoint_a.get_hook_id(); + + send_downward_ping(&mut endpoint_a, &mut endpoint_b, hook_id, b"ping"); + leaf.update_interface(&mut endpoint_b, &mut interface); + + let leaf_id = leaf.get_id(); + let procedure_key = ProcedureKey { + leaf_id, + procedure_id: PROC_PING, + }; + let session_key = SessionKey { + leaf_id, + procedure_id: PROC_PING, + hook_id, + }; + let procedure_view = interface.procedure_views().get(&procedure_key).unwrap(); + + assert!(!interface.session_views().contains_key(&session_key)); + assert!(view_has_event( + &interface, + &procedure_view.events, + |event| matches!( + event, + InterfaceEventKind::Inbound { packet } + if packet.hook_id == hook_id && packet.procedure_id == PROC_PING + ), + )); + assert!(view_has_event( + &interface, + &procedure_view.events, + |event| matches!( + event, + InterfaceEventKind::ProcedureCalled { procedure_id, hook_id: recorded_hook, .. } + if *procedure_id == PROC_PING && *recorded_hook == hook_id + ), + )); + assert!(view_has_event( + &interface, + &procedure_view.events, + |event| matches!( + event, + InterfaceEventKind::RouteSuccess { packet } + if packet.hook_id == hook_id && packet.procedure_id == PROC_PING + ), + )); + + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_packets(&mut endpoint_a, PROC_PING); + + assert_eq!(packets.len(), 1); + assert_eq!(packets[0].hook_id, hook_id); + assert!(packets[0].end_hook); + assert_eq!(packets[0].data, b"ping".to_vec()); +} diff --git a/unshell-leaves/leaf-pty/src/tests/mod.rs b/unshell-leaves/leaf-pty/src/tests/mod.rs new file mode 100644 index 0000000..554fc56 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/mod.rs @@ -0,0 +1,5 @@ +mod session; +mod support; + +#[cfg(feature = "interface")] +mod interface; diff --git a/unshell-leaves/leaf-pty/src/tests.rs b/unshell-leaves/leaf-pty/src/tests/session.rs similarity index 54% rename from unshell-leaves/leaf-pty/src/tests.rs rename to unshell-leaves/leaf-pty/src/tests/session.rs index 8e5d62a..23b53ad 100644 --- a/unshell-leaves/leaf-pty/src/tests.rs +++ b/unshell-leaves/leaf-pty/src/tests/session.rs @@ -1,127 +1,17 @@ use alloc::{vec, vec::Vec}; -use unshell::protocol::{Endpoint, Leaf, Packet}; +use unshell::protocol::{Leaf, Packet}; -#[cfg(feature = "interface")] -use unshell::interface::{InterfaceEventKind, InterfaceStore, SessionKey, SessionViewStatus}; - -use super::{ - FakePtyLeaf, FakePtyState, OP_ABORT, OP_ERROR, OP_EXIT, OP_INPUT, OP_OPENED, OP_OUTPUT, - OP_STDIN_EOF, OP_TERMINATE, PROC_PTY, frame_opcode, frame_payload, pty_open_packet, pty_packet, +use crate::{ + FakePtyLeaf, FakePtyState, OP_ABORT, OP_ERROR, OP_EXIT, OP_INPUT, OP_OUTPUT, OP_STDIN_EOF, + OP_TERMINATE, pty_open_packet, }; -const ENDPOINT_A: u32 = 0; -const ENDPOINT_B: u32 = 1; -const PROC_OTHER: u32 = 31; - -/// Creates a bare endpoint at a known absolute path. -fn endpoint_at(id: u32, path: Vec) -> Endpoint { - let mut endpoint = Endpoint::new(id, vec![]); - endpoint.path = path; - endpoint -} - -/// Creates the parent/child endpoint pair used by PTY session tests. -fn pty_endpoints() -> (Endpoint, Endpoint) { - let mut endpoint_a = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); - let mut endpoint_b = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - - endpoint_a.connections.insert((ENDPOINT_B, false)); - endpoint_b.connections.insert((ENDPOINT_A, true)); - - (endpoint_a, endpoint_b) -} - -/// Transfers every queued packet for `next_hop` into `receiver` as `remote_id` traffic. -fn transfer_packets(sender: &mut Endpoint, receiver: &mut Endpoint, next_hop: u32, remote_id: u32) { - let mut packets = Vec::new(); - sender.take_outbound_clear(next_hop, |packet| packets.push(packet.clone())); - - for packet in packets { - receiver.add_inbound_from(remote_id, packet).unwrap(); - } -} - -/// Sends one downward PTY frame from endpoint A to endpoint B. -fn send_downward_frame( - endpoint_a: &mut Endpoint, - endpoint_b: &mut Endpoint, - hook_id: u16, - opcode: u8, - payload: &[u8], - end_hook: bool, -) { - endpoint_a - .add_outbound(pty_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - end_hook, - opcode, - payload, - )) - .unwrap(); - transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); -} - -/// Opens a fake PTY session and delivers the `Opened` response to endpoint A. -fn open_pty_session( - endpoint_a: &mut Endpoint, - endpoint_b: &mut Endpoint, - leaf: &mut FakePtyLeaf, -) -> u16 { - let hook_id = endpoint_a.get_hook_id(); - endpoint_a - .add_outbound(pty_open_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - &[ENDPOINT_A], - )) - .unwrap(); - - transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); - leaf.update(endpoint_b); - transfer_packets(endpoint_b, endpoint_a, ENDPOINT_A, ENDPOINT_B); - - hook_id -} - -/// Drains PTY packets delivered to endpoint A. -fn drain_parent_pty_packets(endpoint: &mut Endpoint) -> Vec { - let mut packets = Vec::new(); - endpoint.take_inbound_matching( - ENDPOINT_A, - |packet| packet.procedure_id == PROC_PTY, - |packet| packets.push(packet), - ); - packets -} - -/// Asserts that local hook state still contains `hook_id`. -fn assert_hook_present(endpoint: &Endpoint, hook_id: u16) { - assert!(endpoint.has_hook(hook_id)); -} - -/// Asserts that local hook state no longer contains `hook_id`. -fn assert_hook_removed(endpoint: &Endpoint, hook_id: u16) { - assert!(!endpoint.has_hook(hook_id)); -} - -/// Asserts that `packet` carries the expected PTY frame. -fn assert_frame(packet: &Packet, hook_id: u16, opcode: u8, end_hook: bool, payload: &[u8]) { - assert_eq!(packet.hook_id, hook_id); - assert_eq!(packet.end_hook, end_hook); - assert_eq!(frame_opcode(packet), Some(opcode)); - assert_eq!(frame_payload(packet), payload); -} - -/// Returns true when `packets` contains the requested frame. -fn has_frame(packets: &[Packet], hook_id: u16, opcode: u8, payload: &[u8]) -> bool { - packets.iter().any(|packet| { - packet.hook_id == hook_id - && frame_opcode(packet) == Some(opcode) - && frame_payload(packet) == payload - }) -} +use super::support::{ + ENDPOINT_A, ENDPOINT_B, PROC_OTHER, assert_frame, assert_hook_present, assert_hook_removed, + assert_opened, drain_parent_pty_packets, endpoint_at, has_frame, open_pty_session, + pty_endpoints, send_downward_frame, transfer_packets, +}; #[test] fn open_pty_paves_hook_and_creates_session() { @@ -137,7 +27,7 @@ fn open_pty_paves_hook_and_creates_session() { assert_hook_present(&endpoint_a, hook_id); assert_hook_present(&endpoint_b, hook_id); assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_OPENED, false, &[]); + assert_opened(&packets[0], hook_id); } #[test] @@ -394,107 +284,3 @@ fn pty_leaf_does_not_consume_other_leaf_packets() { assert_eq!(other_packets[0].procedure_id, PROC_OTHER); assert_eq!(other_packets[0].data, b"leave-me".to_vec()); } - -#[cfg(feature = "interface")] -#[test] -fn interface_update_records_session_flow() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let mut interface = InterfaceStore::new(); - let hook_id = endpoint_a.get_hook_id(); - - endpoint_a - .add_outbound(pty_open_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - &[ENDPOINT_A], - )) - .unwrap(); - transfer_packets(&mut endpoint_a, &mut endpoint_b, ENDPOINT_B, ENDPOINT_A); - - leaf.update_interface(&mut endpoint_b, &mut interface); - - assert_eq!(leaf.active_session_count(), 1); - assert!(interface.events().iter().any(|event| { - matches!( - &event.kind, - InterfaceEventKind::SessionCreated { hook_id: recorded_hook, .. } - if *recorded_hook == hook_id - ) - })); - assert!(interface.events().iter().any(|event| { - matches!( - &event.kind, - InterfaceEventKind::RouteSuccess { packet } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_OPENED) - ) - })); -} - -#[cfg(feature = "interface")] -#[test] -fn interface_update_records_failed_final_route_without_dropping_session() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let mut interface = InterfaceStore::new(); - let hook_id = endpoint_a.get_hook_id(); - - endpoint_a - .add_outbound(pty_open_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - &[ENDPOINT_A], - )) - .unwrap(); - transfer_packets(&mut endpoint_a, &mut endpoint_b, ENDPOINT_B, ENDPOINT_A); - leaf.update_interface(&mut endpoint_b, &mut interface); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_TERMINATE, - &[], - false, - ); - endpoint_b.connections.remove(&(ENDPOINT_A, true)); - leaf.update_interface(&mut endpoint_b, &mut interface); - - let session_key = SessionKey { - leaf_id: leaf.get_id(), - procedure_id: PROC_PTY, - hook_id, - }; - - assert_eq!(leaf.active_session_count(), 1); - assert_eq!(leaf.pending_packet_count(), 1); - assert_eq!( - interface.session_views().get(&session_key).unwrap().status, - SessionViewStatus::Closed - ); - assert!(interface.events().iter().any(|event| { - matches!( - &event.kind, - InterfaceEventKind::RouteFailure { packet, .. } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_EXIT) - ) - })); - - endpoint_b.connections.insert((ENDPOINT_A, true)); - leaf.update_interface(&mut endpoint_b, &mut interface); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(leaf.active_session_count(), 0); - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); - assert!(interface.events().iter().any(|event| { - matches!( - &event.kind, - InterfaceEventKind::RouteSuccess { packet } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_EXIT) - ) - })); -} diff --git a/unshell-leaves/leaf-pty/src/tests/support.rs b/unshell-leaves/leaf-pty/src/tests/support.rs new file mode 100644 index 0000000..efee503 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/support.rs @@ -0,0 +1,141 @@ +use alloc::{vec, vec::Vec}; + +use unshell::protocol::{Endpoint, Leaf, Packet}; + +use crate::{ + FakePtyLeaf, OP_OPENED, PROC_PTY, frame_opcode, frame_payload, pty_open_packet, pty_packet, +}; + +pub(super) const ENDPOINT_A: u32 = 0; +pub(super) const ENDPOINT_B: u32 = 1; +pub(super) const PROC_OTHER: u32 = 31; + +/// Creates a bare endpoint at a known absolute path. +pub(super) fn endpoint_at(id: u32, path: Vec) -> Endpoint { + let mut endpoint = Endpoint::new(id, vec![]); + endpoint.path = path; + endpoint +} + +/// Creates the parent/child endpoint pair used by PTY session tests. +pub(super) fn pty_endpoints() -> (Endpoint, Endpoint) { + let mut endpoint_a = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); + let mut endpoint_b = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + + endpoint_a.connections.insert((ENDPOINT_B, false)); + endpoint_b.connections.insert((ENDPOINT_A, true)); + + (endpoint_a, endpoint_b) +} + +/// Transfers every queued packet for `next_hop` into `receiver` as `remote_id` traffic. +pub(super) fn transfer_packets( + sender: &mut Endpoint, + receiver: &mut Endpoint, + next_hop: u32, + remote_id: u32, +) { + let mut packets = Vec::new(); + sender.take_outbound_clear(next_hop, |packet| packets.push(packet.clone())); + + for packet in packets { + receiver.add_inbound_from(remote_id, packet).unwrap(); + } +} + +/// Sends one downward PTY frame from endpoint A to endpoint B. +pub(super) fn send_downward_frame( + endpoint_a: &mut Endpoint, + endpoint_b: &mut Endpoint, + hook_id: u16, + opcode: u8, + payload: &[u8], + end_hook: bool, +) { + endpoint_a + .add_outbound(pty_packet( + vec![ENDPOINT_A, ENDPOINT_B], + hook_id, + end_hook, + opcode, + payload, + )) + .unwrap(); + transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); +} + +/// Opens a fake PTY session and delivers the `Opened` response to endpoint A. +pub(super) fn open_pty_session( + endpoint_a: &mut Endpoint, + endpoint_b: &mut Endpoint, + leaf: &mut FakePtyLeaf, +) -> u16 { + let hook_id = endpoint_a.get_hook_id(); + endpoint_a + .add_outbound(pty_open_packet( + vec![ENDPOINT_A, ENDPOINT_B], + hook_id, + &[ENDPOINT_A], + )) + .unwrap(); + + transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); + leaf.update(endpoint_b); + transfer_packets(endpoint_b, endpoint_a, ENDPOINT_A, ENDPOINT_B); + + hook_id +} + +/// Drains packets for `procedure_id` delivered to endpoint A. +pub(super) fn drain_parent_packets(endpoint: &mut Endpoint, procedure_id: u32) -> Vec { + let mut packets = Vec::new(); + endpoint.take_inbound_matching( + ENDPOINT_A, + |packet| packet.procedure_id == procedure_id, + |packet| packets.push(packet), + ); + packets +} + +/// Drains PTY packets delivered to endpoint A. +pub(super) fn drain_parent_pty_packets(endpoint: &mut Endpoint) -> Vec { + drain_parent_packets(endpoint, PROC_PTY) +} + +/// Asserts that local hook state still contains `hook_id`. +pub(super) fn assert_hook_present(endpoint: &Endpoint, hook_id: u16) { + assert!(endpoint.has_hook(hook_id)); +} + +/// Asserts that local hook state no longer contains `hook_id`. +pub(super) fn assert_hook_removed(endpoint: &Endpoint, hook_id: u16) { + assert!(!endpoint.has_hook(hook_id)); +} + +/// Asserts that `packet` carries the expected PTY frame. +pub(super) fn assert_frame( + packet: &Packet, + hook_id: u16, + opcode: u8, + end_hook: bool, + payload: &[u8], +) { + assert_eq!(packet.hook_id, hook_id); + assert_eq!(packet.end_hook, end_hook); + assert_eq!(frame_opcode(packet), Some(opcode)); + assert_eq!(frame_payload(packet), payload); +} + +/// Returns true when `packets` contains the requested frame. +pub(super) fn has_frame(packets: &[Packet], hook_id: u16, opcode: u8, payload: &[u8]) -> bool { + packets.iter().any(|packet| { + packet.hook_id == hook_id + && frame_opcode(packet) == Some(opcode) + && frame_payload(packet) == payload + }) +} + +/// Asserts that a packet is the fake PTY open acknowledgement. +pub(super) fn assert_opened(packet: &Packet, hook_id: u16) { + assert_frame(packet, hook_id, OP_OPENED, false, &[]); +} From 8a817cb5ebacce4b53f57a93cdbac862dc7b91f2 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 10:27:59 -0600 Subject: [PATCH 02/12] Simplify interface event recording. --- src/interface/mod.rs | 3 +- src/interface/store.rs | 235 ++++++++++++++-------------------------- src/interface/target.rs | 46 -------- src/protocol/runtime.rs | 199 +++++++++++++++++++++++----------- 4 files changed, 220 insertions(+), 263 deletions(-) delete mode 100644 src/interface/target.rs diff --git a/src/interface/mod.rs b/src/interface/mod.rs index 513bcf9..65bce89 100644 --- a/src/interface/mod.rs +++ b/src/interface/mod.rs @@ -6,7 +6,6 @@ mod event; mod key; mod store; -mod target; mod view; pub use event::{InterfaceEvent, InterfaceEventKind}; @@ -14,4 +13,4 @@ pub use key::{ProcedureKey, SessionKey}; pub use store::InterfaceStore; pub use view::{ProcedureView, SessionView, SessionViewStatus}; -pub(crate) use target::InterfaceTarget; +pub(crate) use store::InterfaceTarget; diff --git a/src/interface/store.rs b/src/interface/store.rs index d979648..e210d23 100644 --- a/src/interface/store.rs +++ b/src/interface/store.rs @@ -2,12 +2,54 @@ use alloc::{collections::BTreeMap, vec::Vec}; use crate::{ interface::{ - InterfaceEvent, InterfaceEventKind, InterfaceTarget, ProcedureKey, ProcedureView, - SessionKey, SessionView, SessionViewStatus, + InterfaceEvent, InterfaceEventKind, ProcedureKey, ProcedureView, SessionKey, SessionView, + SessionViewStatus, }, protocol::{EndpointError, HookID, Packet, SessionStatus}, }; +/// Internal owner for one interface event. +/// +/// The runtime already knows whether a packet belongs to a hook-backed session or a +/// one-shot procedure. Keeping that answer explicit avoids reconstructing ownership +/// from packet fields later, which is what made procedure packet flow look like fake +/// session activity in the previous store implementation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum InterfaceTarget { + /// Event belongs to one hook-backed session instance. + Session(SessionKey), + + /// Event belongs to one one-shot procedure family. + Procedure(ProcedureKey), +} + +impl InterfaceTarget { + /// Builds a session target from the same pieces exposed by [`SessionKey`]. + pub(crate) fn session(leaf_id: u32, procedure_id: u32, hook_id: HookID) -> Self { + Self::Session(SessionKey { + leaf_id, + procedure_id, + hook_id, + }) + } + + /// Builds a procedure target from the same pieces exposed by [`ProcedureKey`]. + pub(crate) fn procedure(leaf_id: u32, procedure_id: u32) -> Self { + Self::Procedure(ProcedureKey { + leaf_id, + procedure_id, + }) + } + + /// Returns the leaf id used on the append-only event record. + pub(crate) fn leaf_id(self) -> u32 { + match self { + Self::Session(key) => key.leaf_id, + Self::Procedure(key) => key.leaf_id, + } + } +} + /// Caller-owned view and packet-flow store for interface frontends. /// /// Generated leaves receive a mutable reference to this store during interface-aware @@ -86,14 +128,8 @@ impl InterfaceStore { /// Records a packet delivered to a generated leaf. pub fn record_inbound(&mut self, leaf_id: u32, packet: &Packet) { let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); - self.record_inbound_for(target, packet); - } - - /// Records a packet delivered to a target already known by generated runtime code. - pub(crate) fn record_inbound_for(&mut self, target: InterfaceTarget, packet: &Packet) { - self.record( + self.record_for( target, - None, InterfaceEventKind::Inbound { packet: packet.clone(), }, @@ -107,25 +143,11 @@ impl InterfaceStore { procedure_id: u32, hook_id: HookID, ) { - self.record_session_packet_queued_for(InterfaceTarget::session( - leaf_id, - procedure_id, - hook_id, - )); - } - - /// Records that a packet was queued for an existing session inbox. - pub(crate) fn record_session_packet_queued_for(&mut self, target: InterfaceTarget) { - let InterfaceTarget::Session(key) = target else { - return; - }; - - self.record( - target, - None, + self.record_for( + InterfaceTarget::session(leaf_id, procedure_id, hook_id), InterfaceEventKind::SessionPacketQueued { - procedure_id: key.procedure_id, - hook_id: key.hook_id, + procedure_id, + hook_id, }, ); } @@ -138,26 +160,11 @@ impl InterfaceStore { hook_id: HookID, started_ns: Option, ) { - let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); - self.record_session_created_for(target, started_ns); - } - - /// Records successful creation of a new session state for an explicit target. - pub(crate) fn record_session_created_for( - &mut self, - target: InterfaceTarget, - started_ns: Option, - ) { - let InterfaceTarget::Session(key) = target else { - return; - }; - - self.record( - target, - Some(SessionViewStatus::Running), + self.record_for( + InterfaceTarget::session(leaf_id, procedure_id, hook_id), InterfaceEventKind::SessionCreated { - procedure_id: key.procedure_id, - hook_id: key.hook_id, + procedure_id, + hook_id, started_ns, finished_ns: self.now_ns, }, @@ -172,26 +179,11 @@ impl InterfaceStore { hook_id: HookID, started_ns: Option, ) { - let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); - self.record_session_rejected_for(target, started_ns); - } - - /// Records rejection of a packet that could not create a session. - pub(crate) fn record_session_rejected_for( - &mut self, - target: InterfaceTarget, - started_ns: Option, - ) { - let InterfaceTarget::Session(key) = target else { - return; - }; - - self.record( - target, - Some(SessionViewStatus::Rejected), + self.record_for( + InterfaceTarget::session(leaf_id, procedure_id, hook_id), InterfaceEventKind::SessionRejected { - procedure_id: key.procedure_id, - hook_id: key.hook_id, + procedure_id, + hook_id, started_ns, finished_ns: self.now_ns, }, @@ -207,27 +199,11 @@ impl InterfaceStore { status: SessionStatus, started_ns: Option, ) { - let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); - self.record_session_update_for(target, status, started_ns); - } - - /// Records one session update tick for an explicit session target. - pub(crate) fn record_session_update_for( - &mut self, - target: InterfaceTarget, - status: SessionStatus, - started_ns: Option, - ) { - let InterfaceTarget::Session(key) = target else { - return; - }; - - self.record( - target, - Some(SessionViewStatus::from_session_status(status)), + self.record_for( + InterfaceTarget::session(leaf_id, procedure_id, hook_id), InterfaceEventKind::SessionUpdated { - procedure_id: key.procedure_id, - hook_id: key.hook_id, + procedure_id, + hook_id, status, started_ns, finished_ns: self.now_ns, @@ -243,29 +219,10 @@ impl InterfaceStore { hook_id: HookID, started_ns: Option, ) { - self.record_procedure_call_for( + self.record_for( InterfaceTarget::procedure(leaf_id, procedure_id), - hook_id, - started_ns, - ); - } - - /// Records one procedure call for an explicit procedure target. - pub(crate) fn record_procedure_call_for( - &mut self, - target: InterfaceTarget, - hook_id: HookID, - started_ns: Option, - ) { - let InterfaceTarget::Procedure(key) = target else { - return; - }; - - self.record( - target, - None, InterfaceEventKind::ProcedureCalled { - procedure_id: key.procedure_id, + procedure_id, hook_id, started_ns, finished_ns: self.now_ns, @@ -276,14 +233,8 @@ impl InterfaceStore { /// Records a packet emitted by leaf logic before route retry handling. pub fn record_outbound_queued(&mut self, leaf_id: u32, packet: &Packet) { let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); - self.record_outbound_queued_for(target, packet); - } - - /// Records a packet emitted by leaf logic before route retry handling. - pub(crate) fn record_outbound_queued_for(&mut self, target: InterfaceTarget, packet: &Packet) { - self.record( + self.record_for( target, - None, InterfaceEventKind::OutboundQueued { packet: packet.clone(), }, @@ -293,14 +244,8 @@ impl InterfaceStore { /// Records a route attempt for a queued outbound packet. pub fn record_route_attempt(&mut self, leaf_id: u32, packet: &Packet) { let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); - self.record_route_attempt_for(target, packet); - } - - /// Records a route attempt for a queued outbound packet. - pub(crate) fn record_route_attempt_for(&mut self, target: InterfaceTarget, packet: &Packet) { - self.record( + self.record_for( target, - None, InterfaceEventKind::RouteAttempt { packet: packet.clone(), }, @@ -310,14 +255,8 @@ impl InterfaceStore { /// Records a successful route attempt. pub fn record_route_success(&mut self, leaf_id: u32, packet: &Packet) { let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); - self.record_route_success_for(target, packet); - } - - /// Records a successful route attempt. - pub(crate) fn record_route_success_for(&mut self, target: InterfaceTarget, packet: &Packet) { - self.record( + self.record_for( target, - None, InterfaceEventKind::RouteSuccess { packet: packet.clone(), }, @@ -327,19 +266,8 @@ impl InterfaceStore { /// Records a failed route attempt without removing the packet from retry state. pub fn record_route_failure(&mut self, leaf_id: u32, packet: &Packet, error: EndpointError) { let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); - self.record_route_failure_for(target, packet, error); - } - - /// Records a failed route attempt without removing the packet from retry state. - pub(crate) fn record_route_failure_for( - &mut self, - target: InterfaceTarget, - packet: &Packet, - error: EndpointError, - ) { - self.record( + self.record_for( target, - None, InterfaceEventKind::RouteFailure { packet: packet.clone(), error, @@ -347,22 +275,14 @@ impl InterfaceStore { ); } - fn record( - &mut self, - target: InterfaceTarget, - status: Option, - kind: InterfaceEventKind, - ) { + pub(crate) fn record_for(&mut self, target: InterfaceTarget, kind: InterfaceEventKind) { let index = self.push_event(target.leaf_id(), kind); - self.link_event(target, status, index); + self.link_event(target, index); } - fn link_event( - &mut self, - target: InterfaceTarget, - status: Option, - index: usize, - ) { + fn link_event(&mut self, target: InterfaceTarget, index: usize) { + let status = Self::status_for_event(&self.events[index].kind); + match target { InterfaceTarget::Session(key) => { let view = self.session_view_for_key_mut(key); @@ -379,6 +299,17 @@ impl InterfaceStore { } } + fn status_for_event(kind: &InterfaceEventKind) -> Option { + match kind { + InterfaceEventKind::SessionCreated { .. } => Some(SessionViewStatus::Running), + InterfaceEventKind::SessionRejected { .. } => Some(SessionViewStatus::Rejected), + InterfaceEventKind::SessionUpdated { status, .. } => { + Some(SessionViewStatus::from_session_status(*status)) + } + _ => None, + } + } + fn push_event(&mut self, leaf_id: u32, kind: InterfaceEventKind) -> usize { let index = self.events.len(); diff --git a/src/interface/target.rs b/src/interface/target.rs deleted file mode 100644 index bcedf82..0000000 --- a/src/interface/target.rs +++ /dev/null @@ -1,46 +0,0 @@ -use crate::{ - interface::{ProcedureKey, SessionKey}, - protocol::HookID, -}; - -/// Internal owner for one interface event. -/// -/// The runtime already knows whether a packet belongs to a hook-backed session or a -/// one-shot procedure. Keeping that answer explicit avoids reconstructing ownership -/// from packet fields later, which is what made procedure packet flow look like fake -/// session activity in the previous store implementation. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum InterfaceTarget { - /// Event belongs to one hook-backed session instance. - Session(SessionKey), - - /// Event belongs to one one-shot procedure family. - Procedure(ProcedureKey), -} - -impl InterfaceTarget { - /// Builds a session target from the same pieces exposed by [`SessionKey`]. - pub(crate) fn session(leaf_id: u32, procedure_id: u32, hook_id: HookID) -> Self { - Self::Session(SessionKey { - leaf_id, - procedure_id, - hook_id, - }) - } - - /// Builds a procedure target from the same pieces exposed by [`ProcedureKey`]. - pub(crate) fn procedure(leaf_id: u32, procedure_id: u32) -> Self { - Self::Procedure(ProcedureKey { - leaf_id, - procedure_id, - }) - } - - /// Returns the leaf id used on the append-only event record. - pub(crate) fn leaf_id(self) -> u32 { - match self { - Self::Session(key) => key.leaf_id, - Self::Procedure(key) => key.leaf_id, - } - } -} diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index 9f684d5..ecf1347 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -1,7 +1,7 @@ use alloc::collections::VecDeque; use crate::{ - interface::{InterfaceStore, InterfaceTarget}, + interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}, protocol::{ Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionCtx, SessionEntry, SessionFamily, SessionInit, SessionInitResult, SessionStatus, @@ -26,17 +26,7 @@ pub struct LeafOutbox { #[derive(Clone)] struct LeafOutboxEntry { packet: Packet, - target: LeafOutboxTarget, -} - -/// Interface owner attached to a leaf-level outbox entry. -#[derive(Clone, Copy)] -enum LeafOutboxTarget { - /// Compatibility path for packets queued through the public `push`/`extend` API. - InferFromPacket, - - /// Runtime-known session or procedure target. - Explicit(InterfaceTarget), + target: Option, } impl LeafOutbox { @@ -49,7 +39,7 @@ impl LeafOutbox { /// Adds one packet to the retry queue. pub fn push(&mut self, packet: Packet) { - self.push_with_target(packet, LeafOutboxTarget::InferFromPacket); + self.push_with_target(packet, None); } /// Adds all packets from `packets` in FIFO order. @@ -71,7 +61,11 @@ impl LeafOutbox { /// Adds one packet with a runtime-known interface target. pub(crate) fn push_for_target(&mut self, packet: Packet, target: InterfaceTarget) { - self.push_with_target(packet, LeafOutboxTarget::Explicit(target)); + self.push_with_target(packet, Some(target)); + } + + fn push_with_target(&mut self, packet: Packet, target: Option) { + self.packets.push_back(LeafOutboxEntry { packet, target }); } /// Adds all packets with the same runtime-known interface target. @@ -80,10 +74,6 @@ impl LeafOutbox { self.push_for_target(packet, target); } } - - fn push_with_target(&mut self, packet: Packet, target: LeafOutboxTarget) { - self.packets.push_back(LeafOutboxEntry { packet, target }); - } } impl Default for LeafOutbox { @@ -112,7 +102,12 @@ pub fn dispatch_session( let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); if let Some(store) = interface.as_mut() { - store.record_inbound_for(target, &packet); + store.record_for( + target, + InterfaceEventKind::Inbound { + packet: packet.clone(), + }, + ); } if let Some(entry) = family @@ -123,7 +118,13 @@ pub fn dispatch_session( entry.inbox.push_back(packet); if let Some(store) = interface.as_mut() { - store.record_session_packet_queued_for(target); + store.record_for( + target, + InterfaceEventKind::SessionPacketQueued { + procedure_id, + hook_id, + }, + ); } return; @@ -138,18 +139,47 @@ pub fn dispatch_session( family.entries.push(SessionEntry::new(hook_id, state)); if let Some(store) = interface.as_mut() { - store.record_session_created_for(target, started_ns); + store.record_for( + target, + InterfaceEventKind::SessionCreated { + procedure_id, + hook_id, + started_ns, + finished_ns: store.now_ns(), + }, + ); } } SessionInitResult::Rejected => { if let Some(store) = interface.as_mut() { - store.record_session_rejected_for(target, started_ns); + store.record_for( + target, + InterfaceEventKind::SessionRejected { + procedure_id, + hook_id, + started_ns, + finished_ns: store.now_ns(), + }, + ); } } SessionInitResult::RejectedWith(packet) => { if let Some(store) = interface.as_mut() { - store.record_session_rejected_for(target, started_ns); - store.record_outbound_queued_for(target, &packet); + store.record_for( + target, + InterfaceEventKind::SessionRejected { + procedure_id, + hook_id, + started_ns, + finished_ns: store.now_ns(), + }, + ); + store.record_for( + target, + InterfaceEventKind::OutboundQueued { + packet: packet.clone(), + }, + ); } outbox.push_for_target(packet, target); @@ -187,10 +217,24 @@ pub fn update_session_family( let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); if let Some(store) = interface.as_mut() { - store.record_session_update_for(target, status, started_ns); + store.record_for( + target, + InterfaceEventKind::SessionUpdated { + procedure_id: S::PROCEDURE_ID, + hook_id: entry.hook_id, + status, + started_ns, + finished_ns: store.now_ns(), + }, + ); for packet in entry.outbox.iter().skip(outbox_start) { - store.record_outbound_queued_for(target, packet); + store.record_for( + target, + InterfaceEventKind::OutboundQueued { + packet: packet.clone(), + }, + ); } } @@ -215,7 +259,12 @@ pub fn dispatch_procedure( let target = InterfaceTarget::procedure(leaf_id, P::PROCEDURE_ID); if let Some(store) = interface.as_mut() { - store.record_inbound_for(target, &packet); + store.record_for( + target, + InterfaceEventKind::Inbound { + packet: packet.clone(), + }, + ); } let hook_id = packet.hook_id; @@ -227,10 +276,23 @@ pub fn dispatch_procedure( let packets = procedure_out.into_packets(); if let Some(store) = interface.as_mut() { - store.record_procedure_call_for(target, hook_id, started_ns); + store.record_for( + target, + InterfaceEventKind::ProcedureCalled { + procedure_id: P::PROCEDURE_ID, + hook_id, + started_ns, + finished_ns: store.now_ns(), + }, + ); for packet in &packets { - store.record_outbound_queued_for(target, packet); + store.record_for( + target, + InterfaceEventKind::OutboundQueued { + packet: packet.clone(), + }, + ); } } @@ -244,17 +306,13 @@ pub fn flush_leaf_outbox( outbox: &mut LeafOutbox, interface: &mut Option<&mut InterfaceStore>, ) -> bool { - while let Some(entry) = outbox.packets.front().cloned() { - let target = resolve_leaf_outbox_target(leaf_id, &entry); + flush_outbox(endpoint, &mut outbox.packets, interface, |entry| { + let target = entry.target.unwrap_or_else(|| { + InterfaceTarget::session(leaf_id, entry.packet.procedure_id, entry.packet.hook_id) + }); - if !flush_packet_with_target(endpoint, target, &entry.packet, interface) { - return false; - } - - outbox.packets.pop_front(); - } - - true + (target, entry.packet.clone()) + }) } /// Flushes and retains one generated session family. @@ -287,17 +345,12 @@ pub fn flush_packet_queue_with_interface( outbox: &mut PacketQueue, interface: &mut Option<&mut InterfaceStore>, ) -> bool { - while let Some(packet) = outbox.front().cloned() { - let target = InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id); - - if !flush_packet_with_target(endpoint, target, &packet, interface) { - return false; - } - - outbox.pop_front(); - } - - true + flush_outbox(endpoint, outbox, interface, |packet| { + ( + InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id), + packet.clone(), + ) + }) } /// Flushes a packet queue whose owner is already known by the generated runtime. @@ -307,7 +360,20 @@ fn flush_packet_queue_with_target( outbox: &mut PacketQueue, interface: &mut Option<&mut InterfaceStore>, ) -> bool { - while let Some(packet) = outbox.front().cloned() { + flush_outbox(endpoint, outbox, interface, |packet| { + (target, packet.clone()) + }) +} + +fn flush_outbox( + endpoint: &mut Endpoint, + outbox: &mut VecDeque, + interface: &mut Option<&mut InterfaceStore>, + mut packet_for: impl FnMut(&T) -> (InterfaceTarget, Packet), +) -> bool { + while let Some(item) = outbox.front() { + let (target, packet) = packet_for(item); + if !flush_packet_with_target(endpoint, target, &packet, interface) { return false; } @@ -325,20 +391,36 @@ fn flush_packet_with_target( interface: &mut Option<&mut InterfaceStore>, ) -> bool { if let Some(store) = interface.as_mut() { - store.record_route_attempt_for(target, packet); + store.record_for( + target, + InterfaceEventKind::RouteAttempt { + packet: packet.clone(), + }, + ); } match endpoint.add_outbound(packet.clone()) { Ok(()) => { if let Some(store) = interface.as_mut() { - store.record_route_success_for(target, packet); + store.record_for( + target, + InterfaceEventKind::RouteSuccess { + packet: packet.clone(), + }, + ); } true } Err(error) => { if let Some(store) = interface.as_mut() { - store.record_route_failure_for(target, packet, error); + store.record_for( + target, + InterfaceEventKind::RouteFailure { + packet: packet.clone(), + error, + }, + ); } false @@ -346,15 +428,6 @@ fn flush_packet_with_target( } } -fn resolve_leaf_outbox_target(leaf_id: u32, entry: &LeafOutboxEntry) -> InterfaceTarget { - match entry.target { - LeafOutboxTarget::InferFromPacket => { - InterfaceTarget::session(leaf_id, entry.packet.procedure_id, entry.packet.hook_id) - } - LeafOutboxTarget::Explicit(target) => target, - } -} - /// Returns the path used by generated procedure responses. fn parent_reply_path(endpoint: &Endpoint) -> alloc::vec::Vec { if endpoint.path.len() > 1 { From 8ab72d35b0948d1505965303d9d15a46d7e361d0 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:18:16 -0600 Subject: [PATCH 03/12] Derive session routing from hooks --- Cargo.lock | 7 +++ Cargo.toml | 2 +- src/crypto/mod.rs | 16 +++--- src/protocol/endpoint/hooks.rs | 26 ++++++++++ src/protocol/leaf_template.rs | 1 + src/protocol/runtime.rs | 43 +++++++++++----- src/protocol/session.rs | 93 +++++++++++++++++++++------------- 7 files changed, 130 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cb7df14..decd17c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -829,6 +829,13 @@ dependencies = [ "unshell", ] +[[package]] +name = "leaf-shell" +version = "0.1.0" +dependencies = [ + "unshell", +] + [[package]] name = "leb128fmt" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index ef2c4bf..5452c78 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ members = [ "ush-obfuscate", "base62", - "unshell-leaves/leaf-pty", + "unshell-leaves/leaf-pty", "unshell-leaves/leaf-shell", ] resolver = "2" diff --git a/src/crypto/mod.rs b/src/crypto/mod.rs index beb35ca..684af3f 100644 --- a/src/crypto/mod.rs +++ b/src/crypto/mod.rs @@ -1,5 +1,3 @@ -use alloc::string::String; - // TODO: Make this seed dependent on env var; pub const GLOBAL_SEED: u32 = 0xDEAFBEEF; // pub const GLOBAL_NONCE: u32 = { @@ -55,17 +53,17 @@ macro_rules! hash_32 { }}; } -pub fn hash_string_32(input: String) -> u32 { +// pub const fn hash_string_32(input: String) -> u32 { +// let hash: [u8; 32] = sha256(input.as_bytes()); +// u32::from_be_bytes([hash[0], hash[8], hash[16], hash[24]]) +// } + +pub const fn hash_str_32(input: &str) -> u32 { let hash: [u8; 32] = sha256(input.as_bytes()); u32::from_be_bytes([hash[0], hash[8], hash[16], hash[24]]) } -pub fn hash_str_32(input: &str) -> u32 { - let hash: [u8; 32] = sha256(input.as_bytes()); - u32::from_be_bytes([hash[0], hash[8], hash[16], hash[24]]) -} - -pub fn hash_32(input: u32) -> u32 { +pub const fn hash_32(input: u32) -> u32 { let hash: [u8; 32] = sha256(&input.to_be_bytes()); u32::from_be_bytes([hash[0], hash[8], hash[16], hash[24]]) } diff --git a/src/protocol/endpoint/hooks.rs b/src/protocol/endpoint/hooks.rs index 0b4f977..d2e1b8a 100644 --- a/src/protocol/endpoint/hooks.rs +++ b/src/protocol/endpoint/hooks.rs @@ -1,3 +1,5 @@ +use alloc::vec::Vec; + use crate::protocol::{Endpoint, EndpointError, EndpointName}; /// Compact identifier for one routed return channel. @@ -79,6 +81,30 @@ impl Endpoint { self.close_hook(hook_id) } + /// Returns the destination path for packets sent back over `hook_id`. + /// + /// Hooks record the adjacent peer that paved the return channel. This helper turns + /// that peer into the packet path required by the current router: parent peers map + /// to the parent path, and child peers map to the direct child path. Session logic + /// should not store this path itself. + pub(crate) fn hook_path(&self, hook_id: HookID) -> Result, EndpointError> { + let peer = self + .hook_peer(hook_id) + .ok_or(EndpointError::UnknownHook { hook_id })?; + + if self.path.is_empty() { + return Err(EndpointError::EndpointPathUnset); + } + + if self.path.len() > 1 && self.path[self.path.len() - 2] == peer { + Ok(self.path[..self.path.len() - 1].to_vec()) + } else { + let mut path = self.path.clone(); + path.push(peer); + Ok(path) + } + } + /// Validates that `actual_peer` is the peer allowed to use `hook_id`. pub(crate) fn ensure_hook_peer( &self, diff --git a/src/protocol/leaf_template.rs b/src/protocol/leaf_template.rs index cd54fde..339065b 100644 --- a/src/protocol/leaf_template.rs +++ b/src/protocol/leaf_template.rs @@ -123,6 +123,7 @@ macro_rules! unshell_leaf { == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID { $crate::protocol::dispatch_session::<$State, $Session>( + endpoint, leaf_id, &mut self.state, &mut self.$session_field, diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index ecf1347..c16f5d6 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -4,7 +4,7 @@ use crate::{ interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}, protocol::{ Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionCtx, SessionEntry, - SessionFamily, SessionInit, SessionInitResult, SessionStatus, + SessionFamily, SessionInit, SessionInitError, SessionStatus, }, }; @@ -88,6 +88,7 @@ impl Default for LeafOutbox { /// find the hook, initialize missing sessions, queue rejected responses, and update /// interface state when a caller supplied one. pub fn dispatch_session( + endpoint: &Endpoint, leaf_id: u32, leaf: &mut L, family: &mut SessionFamily, @@ -131,12 +132,27 @@ pub fn dispatch_session( } let started_ns = interface.as_ref().and_then(|store| store.now_ns()); + let Ok(path) = endpoint.hook_path(hook_id) else { + if let Some(store) = interface.as_mut() { + store.record_for( + target, + InterfaceEventKind::SessionRejected { + procedure_id, + hook_id, + started_ns, + finished_ns: store.now_ns(), + }, + ); + } + + return; + }; let packet_path = packet.path.clone(); let mut init = SessionInit::new(hook_id, packet_path); match S::init(leaf, packet, &mut init) { - SessionInitResult::Created(state) => { - family.entries.push(SessionEntry::new(hook_id, state)); + Ok(state) => { + family.entries.push(SessionEntry::new(hook_id, path, state)); if let Some(store) = interface.as_mut() { store.record_for( @@ -150,7 +166,7 @@ pub fn dispatch_session( ); } } - SessionInitResult::Rejected => { + Err(SessionInitError::Rejected) => { if let Some(store) = interface.as_mut() { store.record_for( target, @@ -163,7 +179,15 @@ pub fn dispatch_session( ); } } - SessionInitResult::RejectedWith(packet) => { + Err(SessionInitError::Response { data, end_hook }) => { + let packet = Packet { + hook_id, + end_hook, + path, + procedure_id, + data, + }; + if let Some(store) = interface.as_mut() { store.record_for( target, @@ -203,14 +227,9 @@ pub fn update_session_family( let started_ns = interface.as_ref().and_then(|store| store.now_ns()); let outbox_start = entry.outbox.len(); - let reply_path = S::reply_path(&entry.state).to_vec(); + let path = entry.path.clone(); let status = { - let mut ctx = SessionCtx::new( - entry.hook_id, - reply_path, - S::PROCEDURE_ID, - &mut entry.outbox, - ); + let mut ctx = SessionCtx::new(entry.hook_id, path, S::PROCEDURE_ID, &mut entry.outbox); S::update(leaf, &mut entry.state, &mut entry.inbox, &mut ctx) }; diff --git a/src/protocol/session.rs b/src/protocol/session.rs index eefb3bb..0889226 100644 --- a/src/protocol/session.rs +++ b/src/protocol/session.rs @@ -18,16 +18,12 @@ use crate::interface::SessionView; /// const PROCEDURE_ID: u32 = 7; /// type State = MySessionState; /// -/// fn reply_path(state: &Self::State) -> &[u32] { -/// &state.reply_path -/// } -/// /// fn init( /// leaf: &mut MyLeafState, /// packet: Packet, /// ctx: &mut SessionInit, -/// ) -> SessionInitResult { -/// SessionInitResult::Created(MySessionState::from_open(leaf, packet, ctx)) +/// ) -> Result { +/// Ok(MySessionState::from_open(leaf, packet, ctx)) /// } /// /// fn update( @@ -50,20 +46,16 @@ pub trait Session { /// Application state stored for one live hook. type State; - /// Returns the destination path for responses emitted by this session. - /// - /// `Packet` currently carries only a destination path, so protocols that need to - /// reply to a caller should capture a reply path during [`Self::init`]. The - /// generated leaf clones this path into [`SessionCtx`] before calling update so - /// session code can mutably borrow its state while emitting frames. - fn reply_path(session: &Self::State) -> &[u32]; - /// Creates one session state from a packet whose hook has no active session. /// - /// Returning [`SessionInitResult::RejectedWith`] lets the generated leaf route a - /// protocol-level failure response with the same retry guarantees as normal - /// output. Returning [`SessionInitResult::Rejected`] silently consumes the packet. - fn init(leaf: &mut L, packet: Packet, ctx: &mut SessionInit) -> SessionInitResult; + /// The generated runtime derives all response routing from hook state. Session + /// initialization therefore returns only application state or a protocol-level + /// rejection; it never stores or receives a caller reply path. + fn init( + leaf: &mut L, + packet: Packet, + ctx: &mut SessionInit, + ) -> Result; /// Advances one active hook session. /// @@ -119,16 +111,42 @@ impl SessionInit { } } -/// Result of trying to create a session from a packet without an active hook entry. -pub enum SessionInitResult { - /// A new session was created and should be stored by the generated leaf. - Created(S), - - /// The packet was intentionally consumed without creating state or a response. +/// Error returned when a packet cannot create a new session. +pub enum SessionInitError { + /// The packet was intentionally consumed without creating state or sending output. Rejected, - /// The packet was rejected with a response that the generated leaf must route. - RejectedWith(Packet), + /// The packet was rejected with response data that should be sent on the same hook. + Response { + /// Raw `Packet::data` for the response frame. + data: Vec, + + /// Whether the response should close the hook after successful routing. + end_hook: bool, + }, +} + +impl SessionInitError { + /// Creates a silent session rejection. + pub fn rejected() -> Self { + Self::Rejected + } + + /// Creates a non-final response for a rejected session open. + pub fn response(data: Vec) -> Self { + Self::Response { + data, + end_hook: false, + } + } + + /// Creates a final response for a rejected session open. + pub fn response_final(data: Vec) -> Self { + Self::Response { + data, + end_hook: true, + } + } } /// Session lifecycle status returned from [`Session::update`]. @@ -153,7 +171,7 @@ pub enum SessionStatus { /// routing in generated code is what makes final-frame retries reliable. pub struct SessionCtx<'a> { hook_id: HookID, - reply_path: Vec, + path: Vec, procedure_id: u32, outbox: &'a mut PacketQueue, } @@ -162,13 +180,13 @@ impl<'a> SessionCtx<'a> { /// Creates a context for one session update call. pub fn new( hook_id: HookID, - reply_path: Vec, + path: Vec, procedure_id: u32, outbox: &'a mut PacketQueue, ) -> Self { Self { hook_id, - reply_path, + path, procedure_id, outbox, } @@ -179,11 +197,6 @@ impl<'a> SessionCtx<'a> { self.hook_id } - /// Returns the destination path used for packets emitted through this context. - pub fn reply_path(&self) -> &[u32] { - &self.reply_path - } - /// Queues a one-byte-opcode frame without closing the hook. pub fn send(&mut self, opcode: u8, data: &[u8]) { self.send_frame(opcode, data, false); @@ -233,7 +246,7 @@ impl<'a> SessionCtx<'a> { self.outbox.push_back(Packet { hook_id: self.hook_id, end_hook, - path: self.reply_path.clone(), + path: self.path.clone(), procedure_id: self.procedure_id, data, }); @@ -249,6 +262,13 @@ pub struct SessionEntry { /// Hook id associated with this live session. pub hook_id: HookID, + /// Destination path for packets emitted on this hook. + /// + /// This is generated runtime state, not user session state. It is captured from + /// endpoint hook routing when the session is created so leaf sessions never have + /// to carry or understand a reply path. + pub path: Vec, + /// Application-owned session state. pub state: S, @@ -300,9 +320,10 @@ impl Default for SessionFamily { impl SessionEntry { /// Creates one active session entry for `hook_id`. - pub fn new(hook_id: HookID, state: S) -> Self { + pub fn new(hook_id: HookID, path: Vec, state: S) -> Self { Self { hook_id, + path, state, inbox: PacketQueue::new(), outbox: PacketQueue::new(), From f2bc2d912d45ad6fd43c70a0d0f15a0f6c4b905c Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:30:25 -0600 Subject: [PATCH 04/12] Store session state directly --- LEAF_MACRO_INTERFACE.md | 13 ++++++++----- src/protocol/leaf_template.rs | 4 +--- src/protocol/runtime.rs | 6 +++--- src/protocol/session.rs | 26 +++++++++----------------- 4 files changed, 21 insertions(+), 28 deletions(-) diff --git a/LEAF_MACRO_INTERFACE.md b/LEAF_MACRO_INTERFACE.md index 117aa21..869288d 100644 --- a/LEAF_MACRO_INTERFACE.md +++ b/LEAF_MACRO_INTERFACE.md @@ -41,7 +41,7 @@ unshell_leaf! { authors: unshell::alloc::vec!["ASTATIN3"], }, sessions { - pty: PtySession, + pty: PtySessionState, } procedures {} } @@ -59,10 +59,14 @@ The example above expands to the equivalent of: pub struct FakePtyLeaf { state: FakePtyState, outbox: LeafOutbox, - pty: SessionFamily<>::State>, + pty: SessionFamily, } ``` +Session types are the per-hook state values themselves. There is no separate +zero-sized handler struct; a type like `PtySessionState` implements `Session` and is +stored directly in the generated `SessionFamily`. + The wrapper implements: - `new(state)` @@ -149,13 +153,12 @@ Ratatui rendering is a plain feature-gated pass: leaf.render_ratatui(frame, area, &mut interface); ``` -Session rendering is an associated function because session families are type-level -contracts, not stored objects: +Session rendering is an associated function on the stored session state type: ```rust fn render_ratatui( leaf: &LeafState, - session: &Self::State, + session: &Self, view: &mut SessionView, frame: &mut ratatui::Frame<'_>, area: ratatui::layout::Rect, diff --git a/src/protocol/leaf_template.rs b/src/protocol/leaf_template.rs index 339065b..cdbe86c 100644 --- a/src/protocol/leaf_template.rs +++ b/src/protocol/leaf_template.rs @@ -17,9 +17,7 @@ macro_rules! unshell_leaf { state: $State, outbox: $crate::protocol::LeafOutbox, $( - $session_field: $crate::protocol::SessionFamily< - <$Session as $crate::protocol::Session<$State>>::State, - >, + $session_field: $crate::protocol::SessionFamily<$Session>, )* } diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index c16f5d6..78a319e 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -91,7 +91,7 @@ pub fn dispatch_session( endpoint: &Endpoint, leaf_id: u32, leaf: &mut L, - family: &mut SessionFamily, + family: &mut SessionFamily, packet: Packet, outbox: &mut LeafOutbox, interface: &mut Option<&mut InterfaceStore>, @@ -215,7 +215,7 @@ pub fn dispatch_session( pub fn update_session_family( leaf_id: u32, leaf: &mut L, - family: &mut SessionFamily, + family: &mut SessionFamily, interface: &mut Option<&mut InterfaceStore>, ) where S: Session, @@ -338,7 +338,7 @@ pub fn flush_leaf_outbox( pub fn flush_session_family( endpoint: &mut Endpoint, leaf_id: u32, - family: &mut SessionFamily, + family: &mut SessionFamily, interface: &mut Option<&mut InterfaceStore>, ) where S: Session, diff --git a/src/protocol/session.rs b/src/protocol/session.rs index 0889226..e0e9f22 100644 --- a/src/protocol/session.rs +++ b/src/protocol/session.rs @@ -9,26 +9,25 @@ use crate::interface::SessionView; /// /// A session family maps one outer `procedure_id` to many live hook instances. The /// generated leaf owns packet grouping, retry-safe output flushing, and final cleanup; -/// the session implementation owns only application behavior. +/// the session value owns one hook's application behavior and mutable state. /// /// # Example /// /// ```rust,ignore -/// impl Session for MySession { +/// impl Session for MySessionState { /// const PROCEDURE_ID: u32 = 7; -/// type State = MySessionState; /// /// fn init( /// leaf: &mut MyLeafState, /// packet: Packet, /// ctx: &mut SessionInit, -/// ) -> Result { +/// ) -> Result { /// Ok(MySessionState::from_open(leaf, packet, ctx)) /// } /// /// fn update( /// leaf: &mut MyLeafState, -/// session: &mut Self::State, +/// session: &mut Self, /// incoming: &mut PacketQueue, /// ctx: &mut SessionCtx<'_>, /// ) -> SessionStatus { @@ -39,23 +38,16 @@ use crate::interface::SessionView; /// } /// } /// ``` -pub trait Session { +pub trait Session: Sized { /// Outer packet procedure id used by every packet in this session family. const PROCEDURE_ID: u32; - /// Application state stored for one live hook. - type State; - - /// Creates one session state from a packet whose hook has no active session. + /// Creates one session value from a packet whose hook has no active session. /// /// The generated runtime derives all response routing from hook state. Session /// initialization therefore returns only application state or a protocol-level /// rejection; it never stores or receives a caller reply path. - fn init( - leaf: &mut L, - packet: Packet, - ctx: &mut SessionInit, - ) -> Result; + fn init(leaf: &mut L, packet: Packet, ctx: &mut SessionInit) -> Result; /// Advances one active hook session. /// @@ -65,7 +57,7 @@ pub trait Session { /// generated retry rules. fn update( leaf: &mut L, - session: &mut Self::State, + session: &mut Self, incoming: &mut PacketQueue, ctx: &mut SessionCtx<'_>, ) -> SessionStatus; @@ -73,7 +65,7 @@ pub trait Session { #[cfg(feature = "interface_ratatui")] fn render_ratatui( _: &L, - _: &Self::State, + _: &Self, _: &mut SessionView, _: &mut ratatui::Frame<'_>, _: ratatui::layout::Rect, From 64e53c8cfe68878101d0fb1b4ae8fa07dd77e113 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:33:17 -0600 Subject: [PATCH 05/12] Remove session init context --- src/protocol/runtime.rs | 7 ++----- src/protocol/session.rs | 35 ++--------------------------------- 2 files changed, 4 insertions(+), 38 deletions(-) diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index 78a319e..15406b7 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -4,7 +4,7 @@ use crate::{ interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}, protocol::{ Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionCtx, SessionEntry, - SessionFamily, SessionInit, SessionInitError, SessionStatus, + SessionFamily, SessionInitError, SessionStatus, }, }; @@ -147,10 +147,7 @@ pub fn dispatch_session( return; }; - let packet_path = packet.path.clone(); - let mut init = SessionInit::new(hook_id, packet_path); - - match S::init(leaf, packet, &mut init) { + match S::init(leaf, packet) { Ok(state) => { family.entries.push(SessionEntry::new(hook_id, path, state)); diff --git a/src/protocol/session.rs b/src/protocol/session.rs index e0e9f22..a6a5941 100644 --- a/src/protocol/session.rs +++ b/src/protocol/session.rs @@ -20,9 +20,8 @@ use crate::interface::SessionView; /// fn init( /// leaf: &mut MyLeafState, /// packet: Packet, -/// ctx: &mut SessionInit, /// ) -> Result { -/// Ok(MySessionState::from_open(leaf, packet, ctx)) +/// Ok(MySessionState::from_open(leaf, packet)) /// } /// /// fn update( @@ -47,7 +46,7 @@ pub trait Session: Sized { /// The generated runtime derives all response routing from hook state. Session /// initialization therefore returns only application state or a protocol-level /// rejection; it never stores or receives a caller reply path. - fn init(leaf: &mut L, packet: Packet, ctx: &mut SessionInit) -> Result; + fn init(leaf: &mut L, packet: Packet) -> Result; /// Advances one active hook session. /// @@ -73,36 +72,6 @@ pub trait Session: Sized { } } -/// Context passed to [`Session::init`]. -/// -/// This carries routing metadata that the generated leaf already knows before the -/// session state exists. Protocols that need source paths should encode them in the -/// packet payload; `packet_path` is the destination path that routed the packet here. -pub struct SessionInit { - hook_id: HookID, - packet_path: Vec, -} - -impl SessionInit { - /// Creates initialization metadata for a delivered packet. - pub fn new(hook_id: HookID, packet_path: Vec) -> Self { - Self { - hook_id, - packet_path, - } - } - - /// Returns the hook id that will identify the new session. - pub fn hook_id(&self) -> HookID { - self.hook_id - } - - /// Returns the destination path from the packet that reached this leaf. - pub fn packet_path(&self) -> &[u32] { - &self.packet_path - } -} - /// Error returned when a packet cannot create a new session. pub enum SessionInitError { /// The packet was intentionally consumed without creating state or sending output. From 4cd496ed2b4411ff0dc7f27d3054d6ae768e9d5a Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 12:11:01 -0600 Subject: [PATCH 06/12] Simplify session routing path --- LEAF_MACRO_INTERFACE.md | 14 +- src/protocol/endpoint/hooks.rs | 48 +++++- src/protocol/leaf_template.rs | 11 +- src/protocol/runtime.rs | 102 +++---------- src/protocol/session.rs | 140 ++---------------- unshell-leaves/leaf-pty/src/codec.rs | 60 +------- unshell-leaves/leaf-pty/src/lib.rs | 5 +- unshell-leaves/leaf-pty/src/session.rs | 92 ++++++------ unshell-leaves/leaf-pty/src/state.rs | 4 +- .../leaf-pty/src/tests/interface.rs | 62 ++------ unshell-leaves/leaf-pty/src/tests/session.rs | 18 +-- unshell-leaves/leaf-pty/src/tests/support.rs | 6 +- 12 files changed, 166 insertions(+), 396 deletions(-) diff --git a/LEAF_MACRO_INTERFACE.md b/LEAF_MACRO_INTERFACE.md index 869288d..73912d9 100644 --- a/LEAF_MACRO_INTERFACE.md +++ b/LEAF_MACRO_INTERFACE.md @@ -88,11 +88,11 @@ The macro delegates behavior to small helpers: - `update_session_family` - `dispatch_procedure` - `flush_leaf_outbox` -- `flush_session_family` -- `flush_packet_queue_with_interface` This keeps the macro readable. The helper functions own the mechanics of session -lookup, initialization, retry-safe flushing, and optional interface logging. +lookup, initialization, procedure response flushing, and optional interface logging. +Sessions route their own output immediately through `Endpoint` helpers to avoid a +per-session output context and retry queue in small implant builds. ## Interface Store @@ -108,7 +108,7 @@ InterfaceStore Generated leaves receive an optional mutable store during `update_interface`. The helpers create and update the appropriate session/procedure views when packets are -dispatched, sessions update, and outbound routes succeed or fail. +dispatched, sessions update, and queued procedure outbound routes succeed or fail. Internally, interface events are target-driven: @@ -132,9 +132,9 @@ procedure packet both have `procedure_id` and `hook_id`, but they should not bot create session views. The runtime already knows which dispatch branch handled the packet, so that answer is carried into the store. -Leaf-level retry queues also carry the same owner metadata. That matters because the -shared leaf outbox contains both rejected session-init responses and procedure -responses. Session-entry outboxes use their surrounding session key directly. +Leaf-level retry queues carry the same owner metadata for procedure responses. +Session responses bypass this queue and use `Endpoint::send_hook_raw` or +`Endpoint::send_hook_frame` directly. Time remains caller-supplied: diff --git a/src/protocol/endpoint/hooks.rs b/src/protocol/endpoint/hooks.rs index d2e1b8a..097d548 100644 --- a/src/protocol/endpoint/hooks.rs +++ b/src/protocol/endpoint/hooks.rs @@ -1,6 +1,6 @@ use alloc::vec::Vec; -use crate::protocol::{Endpoint, EndpointError, EndpointName}; +use crate::protocol::{Endpoint, EndpointError, EndpointName, Packet}; /// Compact identifier for one routed return channel. /// @@ -105,6 +105,52 @@ impl Endpoint { } } + /// Routes raw response data over an existing hook immediately. + /// + /// This is the compact session-output path: it avoids an intermediate context and + /// retry queue. If a final packet cannot route, the local hook is still removed so + /// an implant does not retain dead hook state forever. + pub fn send_hook_raw( + &mut self, + hook_id: HookID, + procedure_id: u32, + data: Vec, + end_hook: bool, + ) -> Result<(), EndpointError> { + let path = self.hook_path(hook_id)?; + let packet = Packet { + hook_id, + end_hook, + path, + procedure_id, + data, + }; + + let result = self.add_outbound(packet); + + if result.is_err() && end_hook { + self.close_hook(hook_id); + } + + result + } + + /// Routes a one-byte-opcode response frame over an existing hook immediately. + pub fn send_hook_frame( + &mut self, + hook_id: HookID, + procedure_id: u32, + opcode: u8, + payload: &[u8], + end_hook: bool, + ) -> Result<(), EndpointError> { + let mut data = Vec::with_capacity(payload.len() + 1); + data.push(opcode); + data.extend_from_slice(payload); + + self.send_hook_raw(hook_id, procedure_id, data, end_hook) + } + /// Validates that `actual_peer` is the peer allowed to use `hook_id`. pub(crate) fn ensure_hook_peer( &self, diff --git a/src/protocol/leaf_template.rs b/src/protocol/leaf_template.rs index cdbe86c..0040fc8 100644 --- a/src/protocol/leaf_template.rs +++ b/src/protocol/leaf_template.rs @@ -98,6 +98,7 @@ macro_rules! unshell_leaf { $( $crate::protocol::update_session_family::<$State, $Session>( + endpoint, leaf_id, &mut self.state, &mut self.$session_field, @@ -126,7 +127,6 @@ macro_rules! unshell_leaf { &mut self.state, &mut self.$session_field, packet, - &mut self.outbox, interface, ); return; @@ -167,15 +167,6 @@ macro_rules! unshell_leaf { &mut self.outbox, interface, ); - - $( - $crate::protocol::flush_session_family::<$State, $Session>( - endpoint, - leaf_id, - &mut self.$session_field, - interface, - ); - )* } } diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index 15406b7..bb20582 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -3,26 +3,27 @@ use alloc::collections::VecDeque; use crate::{ interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}, protocol::{ - Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionCtx, SessionEntry, + Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionEntry, SessionFamily, SessionInitError, SessionStatus, }, }; /// Retry queue shared by generated leaves. /// -/// Sessions already own per-hook outboxes. This leaf-level queue is for rejected -/// session initialization responses and one-shot procedures, both of which need the -/// same retry semantics as session output without becoming separate framework types. +/// Leaf-level retry queue shared by generated leaves. +/// +/// Sessions route directly through `Endpoint` to keep their runtime shape small. This +/// queue remains only for one-shot procedures, whose handlers still use `ProcedureOut` +/// and should not route while the procedure is borrowing leaf state. pub struct LeafOutbox { packets: VecDeque, } /// One packet retained by a leaf-level retry queue. /// -/// Session entry outboxes have an obvious owner from their surrounding session entry. -/// Leaf-level outboxes are mixed: rejected session initialization packets and one-shot -/// procedure responses both land here. Storing the owner beside the packet keeps route -/// logging precise without exposing another public queue type. +/// Procedure responses from different generated branches share one queue. Storing the +/// owner beside the packet keeps route logging precise without exposing another public +/// queue type. #[derive(Clone)] struct LeafOutboxEntry { packet: Packet, @@ -85,15 +86,14 @@ impl Default for LeafOutbox { /// Dispatches one packet into a generated session family. /// /// The macro picks `S` and the family field. This helper owns the boring details: -/// find the hook, initialize missing sessions, queue rejected responses, and update +/// find the hook, initialize missing sessions, route rejected responses, and update /// interface state when a caller supplied one. pub fn dispatch_session( - endpoint: &Endpoint, + endpoint: &mut Endpoint, leaf_id: u32, leaf: &mut L, family: &mut SessionFamily, packet: Packet, - outbox: &mut LeafOutbox, interface: &mut Option<&mut InterfaceStore>, ) where S: Session, @@ -149,7 +149,7 @@ pub fn dispatch_session( }; match S::init(leaf, packet) { Ok(state) => { - family.entries.push(SessionEntry::new(hook_id, path, state)); + family.entries.push(SessionEntry::new(hook_id, state)); if let Some(store) = interface.as_mut() { store.record_for( @@ -195,21 +195,16 @@ pub fn dispatch_session( finished_ns: store.now_ns(), }, ); - store.record_for( - target, - InterfaceEventKind::OutboundQueued { - packet: packet.clone(), - }, - ); } - outbox.push_for_target(packet, target); + let _ = flush_packet_with_target(endpoint, target, &packet, interface); } } } /// Updates every live session in one generated session family. pub fn update_session_family( + endpoint: &mut Endpoint, leaf_id: u32, leaf: &mut L, family: &mut SessionFamily, @@ -223,13 +218,7 @@ pub fn update_session_family( } let started_ns = interface.as_ref().and_then(|store| store.now_ns()); - let outbox_start = entry.outbox.len(); - let path = entry.path.clone(); - let status = { - let mut ctx = SessionCtx::new(entry.hook_id, path, S::PROCEDURE_ID, &mut entry.outbox); - - S::update(leaf, &mut entry.state, &mut entry.inbox, &mut ctx) - }; + let status = S::update(leaf, &mut entry.state, &mut entry.inbox, endpoint); let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); if let Some(store) = interface.as_mut() { @@ -243,21 +232,14 @@ pub fn update_session_family( finished_ns: store.now_ns(), }, ); - - for packet in entry.outbox.iter().skip(outbox_start) { - store.record_for( - target, - InterfaceEventKind::OutboundQueued { - packet: packet.clone(), - }, - ); - } } if matches!(status, SessionStatus::Closed) { entry.closed = true; } } + + family.entries.retain(|entry| !entry.closed); } /// Dispatches one packet into a generated one-shot procedure. @@ -331,56 +313,6 @@ pub fn flush_leaf_outbox( }) } -/// Flushes and retains one generated session family. -pub fn flush_session_family( - endpoint: &mut Endpoint, - leaf_id: u32, - family: &mut SessionFamily, - interface: &mut Option<&mut InterfaceStore>, -) where - S: Session, -{ - for entry in &mut family.entries { - let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); - flush_packet_queue_with_target(endpoint, target, &mut entry.outbox, interface); - } - - family - .entries - .retain(|entry| !entry.closed || !entry.outbox.is_empty()); -} - -/// Flushes a retry queue through [`Endpoint::add_outbound`]. -/// -/// This is the interface-aware version of [`crate::protocol::flush_packet_queue`]. It -/// logs route attempts before trying them, then logs either success or the route error -/// without dropping the packet on failure. -pub fn flush_packet_queue_with_interface( - endpoint: &mut Endpoint, - leaf_id: u32, - outbox: &mut PacketQueue, - interface: &mut Option<&mut InterfaceStore>, -) -> bool { - flush_outbox(endpoint, outbox, interface, |packet| { - ( - InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id), - packet.clone(), - ) - }) -} - -/// Flushes a packet queue whose owner is already known by the generated runtime. -fn flush_packet_queue_with_target( - endpoint: &mut Endpoint, - target: InterfaceTarget, - outbox: &mut PacketQueue, - interface: &mut Option<&mut InterfaceStore>, -) -> bool { - flush_outbox(endpoint, outbox, interface, |packet| { - (target, packet.clone()) - }) -} - fn flush_outbox( endpoint: &mut Endpoint, outbox: &mut VecDeque, diff --git a/src/protocol/session.rs b/src/protocol/session.rs index a6a5941..405b990 100644 --- a/src/protocol/session.rs +++ b/src/protocol/session.rs @@ -28,10 +28,10 @@ use crate::interface::SessionView; /// leaf: &mut MyLeafState, /// session: &mut Self, /// incoming: &mut PacketQueue, -/// ctx: &mut SessionCtx<'_>, +/// endpoint: &mut Endpoint, /// ) -> SessionStatus { /// while let Some(packet) = incoming.pop_front() { -/// session.apply(leaf, packet, ctx); +/// session.apply(leaf, packet, endpoint); /// } /// SessionStatus::Running /// } @@ -51,14 +51,15 @@ pub trait Session: Sized { /// Advances one active hook session. /// /// The generated leaf calls this for every live session on each update tick so - /// sessions can poll external workers even when no new packet arrived. Outbound - /// packets must be queued through `ctx`; direct endpoint routing would bypass the - /// generated retry rules. + /// sessions can poll external workers even when no new packet arrived. Session + /// output is routed immediately through `endpoint`; callers that need retry + /// semantics should keep their own compact application state and retry on a later + /// tick. fn update( leaf: &mut L, session: &mut Self, incoming: &mut PacketQueue, - ctx: &mut SessionCtx<'_>, + endpoint: &mut Endpoint, ) -> SessionStatus; #[cfg(feature = "interface_ratatui")] @@ -121,99 +122,11 @@ pub enum SessionStatus { /// The session has finished application work. /// - /// The generated leaf still retains the entry until every queued packet routes - /// successfully, which prevents a failed final frame from losing session cleanup. + /// The generated leaf removes the entry after the update tick. Final packets are + /// routed immediately by the session before returning this status. Closed, } -/// Mutable output context passed to [`Session::update`]. -/// -/// The context queues packets only; it never routes them immediately. Centralizing -/// routing in generated code is what makes final-frame retries reliable. -pub struct SessionCtx<'a> { - hook_id: HookID, - path: Vec, - procedure_id: u32, - outbox: &'a mut PacketQueue, -} - -impl<'a> SessionCtx<'a> { - /// Creates a context for one session update call. - pub fn new( - hook_id: HookID, - path: Vec, - procedure_id: u32, - outbox: &'a mut PacketQueue, - ) -> Self { - Self { - hook_id, - path, - procedure_id, - outbox, - } - } - - /// Returns the hook id used for packets emitted through this context. - pub fn hook_id(&self) -> HookID { - self.hook_id - } - - /// Queues a one-byte-opcode frame without closing the hook. - pub fn send(&mut self, opcode: u8, data: &[u8]) { - self.send_frame(opcode, data, false); - } - - /// Queues a one-byte-opcode frame that closes the hook after successful routing. - pub fn send_final(&mut self, opcode: u8, data: &[u8]) { - self.send_frame(opcode, data, true); - } - - /// Queues a protocol-specific error frame without closing the hook. - /// - /// The `code` is used as the frame opcode because the protocol layer does not - /// reserve a universal error opcode. Leaves that have a dedicated error opcode can - /// pass that value here or call [`Self::send`] directly. - pub fn error(&mut self, code: u8, data: &[u8]) { - self.send(code, data); - } - - /// Queues a protocol-specific error frame that closes the hook after routing. - pub fn error_final(&mut self, code: u8, data: &[u8]) { - self.send_final(code, data); - } - - /// Queues raw packet data without adding an opcode byte. - pub fn send_raw(&mut self, data: &[u8]) { - self.send_raw_with_end(data, false); - } - - /// Queues raw packet data and closes the hook after successful routing. - pub fn send_raw_final(&mut self, data: &[u8]) { - self.send_raw_with_end(data, true); - } - - fn send_frame(&mut self, opcode: u8, data: &[u8], end_hook: bool) { - let mut frame = Vec::with_capacity(data.len() + 1); - frame.push(opcode); - frame.extend_from_slice(data); - self.enqueue_data(frame, end_hook); - } - - fn send_raw_with_end(&mut self, data: &[u8], end_hook: bool) { - self.enqueue_data(data.to_vec(), end_hook); - } - - fn enqueue_data(&mut self, data: Vec, end_hook: bool) { - self.outbox.push_back(Packet { - hook_id: self.hook_id, - end_hook, - path: self.path.clone(), - procedure_id: self.procedure_id, - data, - }); - } -} - /// Storage entry used by macro-generated session stores. /// /// The fields are public so generated code in downstream crates can keep the update @@ -223,23 +136,13 @@ pub struct SessionEntry { /// Hook id associated with this live session. pub hook_id: HookID, - /// Destination path for packets emitted on this hook. - /// - /// This is generated runtime state, not user session state. It is captured from - /// endpoint hook routing when the session is created so leaf sessions never have - /// to carry or understand a reply path. - pub path: Vec, - /// Application-owned session state. pub state: S, /// Packets delivered for this hook but not yet consumed by the session. pub inbox: PacketQueue, - /// Packets emitted by the session but not yet accepted by endpoint routing. - pub outbox: PacketQueue, - - /// Whether application logic has finished and only retry flushing may remain. + /// Whether application logic has finished and should be removed after update. pub closed: bool, } @@ -266,7 +169,7 @@ impl SessionFamily { let mut count = 0usize; for entry in &self.entries { - count += entry.inbox.len() + entry.outbox.len(); + count += entry.inbox.len(); } count @@ -281,31 +184,12 @@ impl Default for SessionFamily { impl SessionEntry { /// Creates one active session entry for `hook_id`. - pub fn new(hook_id: HookID, path: Vec, state: S) -> Self { + pub fn new(hook_id: HookID, state: S) -> Self { Self { hook_id, - path, state, inbox: PacketQueue::new(), - outbox: PacketQueue::new(), closed: false, } } } - -/// Flushes a retry queue through [`Endpoint::add_outbound`]. -/// -/// The packet at the front is cloned for each attempt and removed only after routing -/// succeeds. This preserves final frames when a route is temporarily unavailable. -/// The return value is true when the queue was fully drained. -pub fn flush_packet_queue(endpoint: &mut Endpoint, outbox: &mut PacketQueue) -> bool { - while let Some(packet) = outbox.front().cloned() { - if endpoint.add_outbound(packet).is_err() { - return false; - } - - outbox.pop_front(); - } - - true -} diff --git a/unshell-leaves/leaf-pty/src/codec.rs b/unshell-leaves/leaf-pty/src/codec.rs index f421385..50409a3 100644 --- a/unshell-leaves/leaf-pty/src/codec.rs +++ b/unshell-leaves/leaf-pty/src/codec.rs @@ -2,7 +2,7 @@ use alloc::vec::Vec; use unshell::protocol::{HookID, Packet}; -use crate::{OP_ERROR, OP_OPEN, PROC_PTY}; +use crate::{OP_OPEN, PROC_PTY}; /// Encodes a tiny PTY frame into `Packet::data`. pub fn encode_frame(opcode: u8, payload: &[u8]) -> Vec { @@ -12,35 +12,9 @@ pub fn encode_frame(opcode: u8, payload: &[u8]) -> Vec { data } -/// Encodes an `Open` payload with the caller's reply path. -pub fn encode_open(reply_path: &[u32]) -> Vec { - let mut data = Vec::with_capacity(2 + reply_path.len() * 4); - data.push(OP_OPEN); - data.push(reply_path.len() as u8); - - for segment in reply_path { - data.extend_from_slice(&segment.to_le_bytes()); - } - - data -} - -/// Decodes the reply path embedded in an `Open` payload after the opcode byte. -pub fn decode_open_reply_path(payload: &[u8]) -> Option> { - let path_len = usize::from(*payload.first()?); - let path_bytes = path_len.checked_mul(4)?; - let expected_len = 1usize.checked_add(path_bytes)?; - - if payload.len() != expected_len { - return None; - } - - let mut path = Vec::with_capacity(path_len); - for chunk in payload[1..].chunks_exact(4) { - path.push(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); - } - - Some(path) +/// Encodes an `Open` frame. +pub fn encode_open() -> Vec { + alloc::vec![OP_OPEN] } /// Returns the opcode byte from a PTY packet, if present. @@ -74,33 +48,13 @@ pub fn pty_packet( } } -/// Builds an outer PTY open packet with the specialized open payload shape. -pub fn pty_open_packet(path: Vec, hook_id: HookID, reply_path: &[u32]) -> Packet { +/// Builds an outer PTY open packet. +pub fn pty_open_packet(path: Vec, hook_id: HookID) -> Packet { Packet { hook_id, end_hook: false, path, procedure_id: PROC_PTY, - data: encode_open(reply_path), - } -} - -/// Builds a final error packet for session initialization failures. -pub(crate) fn error_packet(hook_id: HookID, reply_path: Vec, payload: &[u8]) -> Packet { - Packet { - hook_id, - end_hook: true, - path: reply_path, - procedure_id: PROC_PTY, - data: encode_frame(OP_ERROR, payload), - } -} - -/// Infers the caller reply path from a locally delivered destination path. -pub(crate) fn reply_path_from_destination(destination: &[u32]) -> Vec { - if destination.len() > 1 { - destination[..destination.len() - 1].to_vec() - } else { - destination.to_vec() + data: encode_open(), } } diff --git a/unshell-leaves/leaf-pty/src/lib.rs b/unshell-leaves/leaf-pty/src/lib.rs index d040da0..5b42b6f 100644 --- a/unshell-leaves/leaf-pty/src/lib.rs +++ b/unshell-leaves/leaf-pty/src/lib.rs @@ -16,11 +16,10 @@ mod session; mod state; pub use codec::{ - decode_open_reply_path, encode_frame, encode_open, frame_opcode, frame_payload, - pty_open_packet, pty_packet, + encode_frame, encode_open, frame_opcode, frame_payload, pty_open_packet, pty_packet, }; pub use constants::*; -pub use session::{PtySession, PtySessionState}; +pub use session::PtySessionState; pub use state::{FakePtyLeaf, FakePtyState}; #[cfg(test)] diff --git a/unshell-leaves/leaf-pty/src/session.rs b/unshell-leaves/leaf-pty/src/session.rs index eccd085..66a5e75 100644 --- a/unshell-leaves/leaf-pty/src/session.rs +++ b/unshell-leaves/leaf-pty/src/session.rs @@ -1,14 +1,9 @@ -use alloc::vec::Vec; - use unshell::protocol::{ - HookID, Packet, PacketQueue, Session, SessionCtx, SessionInit, SessionInitResult, SessionStatus, + Endpoint, HookID, Packet, PacketQueue, Session, SessionInitError, SessionStatus, }; use crate::{ - codec::{ - decode_open_reply_path, error_packet, frame_opcode, frame_payload, - reply_path_from_destination, - }, + codec::{encode_frame, frame_opcode, frame_payload}, constants::{ OP_ABORT, OP_ERROR, OP_EXIT, OP_INPUT, OP_OPEN, OP_OPENED, OP_OUTPUT, OP_STDIN_EOF, OP_TERMINATE, PROC_PTY, @@ -16,51 +11,32 @@ use crate::{ state::FakePtyState, }; -/// Session contract for one hook-backed fake PTY. -pub struct PtySession; - /// Per-hook fake PTY session state. /// -/// A real PTY leaf will replace the pending flags with a worker handle. The reply path -/// and hook lifecycle behavior should stay the same. +/// A real PTY leaf will replace the pending flags with a worker handle. Hook routing +/// is owned by the generated runtime, so this state only tracks PTY behavior. pub struct PtySessionState { hook_id: HookID, - reply_path: Vec, opened_pending: bool, stdin_closed: bool, } -impl Session for PtySession { +impl Session for PtySessionState { const PROCEDURE_ID: u32 = PROC_PTY; - type State = PtySessionState; - - fn reply_path(session: &Self::State) -> &[u32] { - &session.reply_path - } - - fn init( - leaf: &mut FakePtyState, - packet: Packet, - ctx: &mut SessionInit, - ) -> SessionInitResult { + fn init(leaf: &mut FakePtyState, packet: Packet) -> Result { if frame_opcode(&packet) != Some(OP_OPEN) { - return SessionInitResult::RejectedWith(error_packet( - ctx.hook_id(), - reply_path_from_destination(ctx.packet_path()), + return Err(SessionInitError::response_final(encode_frame( + OP_ERROR, b"unknown-session", - )); + ))); } - let reply_path = decode_open_reply_path(frame_payload(&packet)) - .unwrap_or_else(|| reply_path_from_destination(ctx.packet_path())); - leaf.active_count += 1; leaf.total_opened += 1; - SessionInitResult::Created(PtySessionState { - hook_id: ctx.hook_id(), - reply_path, + Ok(Self { + hook_id: packet.hook_id, opened_pending: true, stdin_closed: false, }) @@ -68,24 +44,44 @@ impl Session for PtySession { fn update( leaf: &mut FakePtyState, - session: &mut Self::State, + session: &mut Self, incoming: &mut PacketQueue, - ctx: &mut SessionCtx<'_>, + endpoint: &mut Endpoint, ) -> SessionStatus { if session.opened_pending { - ctx.send(OP_OPENED, &[]); + let _ = endpoint.send_hook_frame( + session.hook_id, + Self::PROCEDURE_ID, + OP_OPENED, + &[], + false, + ); session.opened_pending = false; } while let Some(packet) = incoming.pop_front() { match frame_opcode(&packet) { - Some(OP_INPUT) => ctx.send(OP_OUTPUT, frame_payload(&packet)), + Some(OP_INPUT) => { + let _ = endpoint.send_hook_frame( + session.hook_id, + Self::PROCEDURE_ID, + OP_OUTPUT, + frame_payload(&packet), + false, + ); + } Some(OP_STDIN_EOF) => { session.stdin_closed = true; leaf.last_stdin_eof_hook = Some(session.hook_id); } Some(OP_TERMINATE) => { - ctx.send_final(OP_EXIT, &[0]); + let _ = endpoint.send_hook_frame( + session.hook_id, + Self::PROCEDURE_ID, + OP_EXIT, + &[0], + true, + ); close_session(leaf); return SessionStatus::Closed; } @@ -94,12 +90,24 @@ impl Session for PtySession { return SessionStatus::Closed; } Some(OP_OPEN) => { - ctx.send_final(OP_ERROR, b"duplicate-open"); + let _ = endpoint.send_hook_frame( + session.hook_id, + Self::PROCEDURE_ID, + OP_ERROR, + b"duplicate-open", + true, + ); close_session(leaf); return SessionStatus::Closed; } _ => { - ctx.send_final(OP_ERROR, b"unknown-opcode"); + let _ = endpoint.send_hook_frame( + session.hook_id, + Self::PROCEDURE_ID, + OP_ERROR, + b"unknown-opcode", + true, + ); close_session(leaf); return SessionStatus::Closed; } diff --git a/unshell-leaves/leaf-pty/src/state.rs b/unshell-leaves/leaf-pty/src/state.rs index 92cdf53..bdedfe6 100644 --- a/unshell-leaves/leaf-pty/src/state.rs +++ b/unshell-leaves/leaf-pty/src/state.rs @@ -1,6 +1,6 @@ use unshell::protocol::{HookID, unshell_leaf}; -use crate::{constants::LEAF_FAKE_PTY, procedure::PingProcedure, session::PtySession}; +use crate::{constants::LEAF_FAKE_PTY, procedure::PingProcedure, session::PtySessionState}; /// User-owned state for the generated fake PTY leaf. /// @@ -45,7 +45,7 @@ unshell_leaf! { authors: unshell::alloc::vec!["ASTATIN3"], }, sessions { - pty: PtySession, + pty: PtySessionState, } procedures { ping: PingProcedure, diff --git a/unshell-leaves/leaf-pty/src/tests/interface.rs b/unshell-leaves/leaf-pty/src/tests/interface.rs index a3977c4..e2929f0 100644 --- a/unshell-leaves/leaf-pty/src/tests/interface.rs +++ b/unshell-leaves/leaf-pty/src/tests/interface.rs @@ -2,17 +2,16 @@ use alloc::vec; use unshell::{ interface::{InterfaceEventKind, InterfaceStore, ProcedureKey, SessionKey, SessionViewStatus}, - protocol::{Leaf, Packet}, + protocol::{Leaf, Packet, SessionStatus}, }; use crate::{ - FakePtyLeaf, FakePtyState, OP_EXIT, OP_OPENED, OP_TERMINATE, PROC_PTY, constants::PROC_PING, - frame_opcode, pty_open_packet, + FakePtyLeaf, FakePtyState, OP_TERMINATE, PROC_PTY, constants::PROC_PING, pty_open_packet, }; use super::support::{ - ENDPOINT_A, ENDPOINT_B, assert_frame, drain_parent_packets, drain_parent_pty_packets, - pty_endpoints, send_downward_frame, transfer_packets, + ENDPOINT_A, ENDPOINT_B, drain_parent_packets, drain_parent_pty_packets, pty_endpoints, + send_downward_frame, transfer_packets, }; fn view_has_event(interface: &InterfaceStore, event_indexes: &[usize], mut predicate: F) -> bool @@ -51,11 +50,7 @@ fn interface_update_records_session_flow() { let hook_id = endpoint_a.get_hook_id(); endpoint_a - .add_outbound(pty_open_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - &[ENDPOINT_A], - )) + .add_outbound(pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) .unwrap(); transfer_packets(&mut endpoint_a, &mut endpoint_b, ENDPOINT_B, ENDPOINT_A); @@ -83,34 +78,21 @@ fn interface_update_records_session_flow() { &session_view.events, |event| matches!( event, - InterfaceEventKind::OutboundQueued { packet } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_OPENED) - ), - )); - assert!(view_has_event( - &interface, - &session_view.events, - |event| matches!( - event, - InterfaceEventKind::RouteSuccess { packet } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_OPENED) + InterfaceEventKind::SessionUpdated { hook_id: recorded_hook, status, .. } + if *recorded_hook == hook_id && *status == SessionStatus::Running ), )); } #[test] -fn interface_update_records_failed_final_route_without_dropping_session() { +fn interface_update_records_failed_direct_route_without_retry() { let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); let mut leaf = FakePtyLeaf::new(FakePtyState::new()); let mut interface = InterfaceStore::new(); let hook_id = endpoint_a.get_hook_id(); endpoint_a - .add_outbound(pty_open_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - &[ENDPOINT_A], - )) + .add_outbound(pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) .unwrap(); transfer_packets(&mut endpoint_a, &mut endpoint_b, ENDPOINT_B, ENDPOINT_A); leaf.update_interface(&mut endpoint_b, &mut interface); @@ -135,18 +117,9 @@ fn interface_update_records_failed_final_route_without_dropping_session() { }; let session_view = interface.session_views().get(&session_key).unwrap(); - assert_eq!(leaf.active_session_count(), 1); - assert_eq!(leaf.pending_packet_count(), 1); + assert_eq!(leaf.active_session_count(), 0); + assert_eq!(leaf.pending_packet_count(), 0); assert_eq!(session_view.status, SessionViewStatus::Closed); - assert!(view_has_event( - &interface, - &session_view.events, - |event| matches!( - event, - InterfaceEventKind::RouteFailure { packet, .. } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_EXIT) - ), - )); endpoint_b.connections.insert((ENDPOINT_A, true)); leaf.update_interface(&mut endpoint_b, &mut interface); @@ -156,17 +129,8 @@ fn interface_update_records_failed_final_route_without_dropping_session() { let session_view = interface.session_views().get(&session_key).unwrap(); assert_eq!(leaf.active_session_count(), 0); - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); - assert!(view_has_event( - &interface, - &session_view.events, - |event| matches!( - event, - InterfaceEventKind::RouteSuccess { packet } - if packet.hook_id == hook_id && frame_opcode(packet) == Some(OP_EXIT) - ), - )); + assert!(packets.is_empty()); + assert_eq!(session_view.status, SessionViewStatus::Closed); } #[test] diff --git a/unshell-leaves/leaf-pty/src/tests/session.rs b/unshell-leaves/leaf-pty/src/tests/session.rs index 23b53ad..285cf83 100644 --- a/unshell-leaves/leaf-pty/src/tests/session.rs +++ b/unshell-leaves/leaf-pty/src/tests/session.rs @@ -124,7 +124,7 @@ fn exit_end_hook_cleans_route_and_session() { } #[test] -fn failed_final_exit_route_retries_without_losing_session() { +fn failed_final_exit_route_closes_session_without_retry() { let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); let mut leaf = FakePtyLeaf::new(FakePtyState::new()); let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); @@ -141,19 +141,18 @@ fn failed_final_exit_route_retries_without_losing_session() { endpoint_b.connections.remove(&(ENDPOINT_A, true)); leaf.update(&mut endpoint_b); - assert_eq!(leaf.active_session_count(), 1); - assert_eq!(leaf.pending_packet_count(), 1); - assert_hook_present(&endpoint_b, hook_id); + assert_eq!(leaf.active_session_count(), 0); + assert_eq!(leaf.pending_packet_count(), 0); + assert_hook_removed(&endpoint_b, hook_id); endpoint_b.connections.insert((ENDPOINT_A, true)); leaf.update(&mut endpoint_b); transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); let packets = drain_parent_pty_packets(&mut endpoint_a); - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); + assert!(packets.is_empty()); assert_eq!(leaf.active_session_count(), 0); - assert_hook_removed(&endpoint_a, hook_id); + assert_hook_present(&endpoint_a, hook_id); assert_hook_removed(&endpoint_b, hook_id); } @@ -252,10 +251,7 @@ fn pty_leaf_does_not_consume_other_leaf_packets() { endpoint.connections.insert((ENDPOINT_A, true)); endpoint - .add_inbound_from( - ENDPOINT_A, - pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], 7, &[ENDPOINT_A]), - ) + .add_inbound_from(ENDPOINT_A, pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], 7)) .unwrap(); endpoint .add_inbound_from( diff --git a/unshell-leaves/leaf-pty/src/tests/support.rs b/unshell-leaves/leaf-pty/src/tests/support.rs index efee503..4298db7 100644 --- a/unshell-leaves/leaf-pty/src/tests/support.rs +++ b/unshell-leaves/leaf-pty/src/tests/support.rs @@ -72,11 +72,7 @@ pub(super) fn open_pty_session( ) -> u16 { let hook_id = endpoint_a.get_hook_id(); endpoint_a - .add_outbound(pty_open_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - &[ENDPOINT_A], - )) + .add_outbound(pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) .unwrap(); transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); From 7749f62629148c664b9acc645c31d7538767e52e Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:08:26 -0600 Subject: [PATCH 07/12] Shrink endpoint runtime footprint --- Cargo.lock | 8 + Cargo.toml | 2 +- build.sh | 5 +- examples/endpoint_test/Cargo.toml | 17 + examples/endpoint_test/src/main.rs | 19 + src/interface/mod.rs | 1 + src/protocol/endpoint/hooks.rs | 17 +- src/protocol/endpoint/mod.rs | 162 +++++-- src/protocol/endpoint/routing.rs | 12 +- src/protocol/leaf_template.rs | 142 +++++- src/protocol/leaf_template/no_procedures.rs | 240 ++++++++++ src/protocol/mod.rs | 11 +- src/protocol/runtime.rs | 409 +++++++++++------- src/protocol/tests/merkle_sync/harness.rs | 54 ++- src/protocol/tests/merkle_sync/leaves.rs | 4 +- src/protocol/tests/merkle_sync/state.rs | 4 +- src/protocol/tests/oneshot/mod.rs | 224 +++++----- src/protocol/tests/oneshot/streams.rs | 199 ++++----- src/protocol/tests/oneshot/support.rs | 14 +- .../leaf-pty/src/tests/interface.rs | 4 +- unshell-leaves/leaf-pty/src/tests/session.rs | 6 +- unshell-leaves/leaf-pty/src/tests/support.rs | 6 +- unshell-leaves/leaf-shell/Cargo.toml | 28 ++ unshell-leaves/leaf-shell/src/lib.rs | 3 + unshell-leaves/leaf-shell/src/shell/mod.rs | 143 ++++++ 25 files changed, 1245 insertions(+), 489 deletions(-) create mode 100644 examples/endpoint_test/Cargo.toml create mode 100644 examples/endpoint_test/src/main.rs create mode 100644 src/protocol/leaf_template/no_procedures.rs create mode 100644 unshell-leaves/leaf-shell/Cargo.toml create mode 100644 unshell-leaves/leaf-shell/src/lib.rs create mode 100644 unshell-leaves/leaf-shell/src/shell/mod.rs diff --git a/Cargo.lock b/Cargo.lock index decd17c..6e9a5a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -512,6 +512,14 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +[[package]] +name = "endpoint_test" +version = "0.1.0" +dependencies = [ + "leaf-shell", + "unshell", +] + [[package]] name = "equivalent" version = "1.0.2" diff --git a/Cargo.toml b/Cargo.toml index 5452c78..f0abdc6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ members = [ "ush-obfuscate", "base62", - "unshell-leaves/leaf-pty", "unshell-leaves/leaf-shell", + "unshell-leaves/leaf-pty", "unshell-leaves/leaf-shell", "examples/endpoint_test", ] resolver = "2" diff --git a/build.sh b/build.sh index 86ccad0..0474926 100755 --- a/build.sh +++ b/build.sh @@ -8,9 +8,10 @@ set -e OBFUSCATION_KEY=kjwerkwerkjbwejehrwhje \ -cargo build --profile minimize -p treetest $@ +# RUSTFLAGS="-Zlocation-detail=none -Zfmt-debug=none" \ +cargo build --profile minimize -p endpoint_test $@ -export BINARY=./target/minimize/treetest +export BINARY=./target/minimize/endpoint_test declare -a headers=( ".gnu_debuglink" # - Debug information link diff --git a/examples/endpoint_test/Cargo.toml b/examples/endpoint_test/Cargo.toml new file mode 100644 index 0000000..b6a7a76 --- /dev/null +++ b/examples/endpoint_test/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "endpoint_test" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +include.workspace = true + +[dependencies] +unshell = { workspace = true } +leaf-shell = { path = "../../unshell-leaves/leaf-shell" } + +[[bin]] +name = "endpoint_test" +path = "src/main.rs" +test = false diff --git a/examples/endpoint_test/src/main.rs b/examples/endpoint_test/src/main.rs new file mode 100644 index 0000000..7d4f826 --- /dev/null +++ b/examples/endpoint_test/src/main.rs @@ -0,0 +1,19 @@ +#![no_std] +#![no_main] + +extern crate alloc; + +use leaf_shell::{ShellLeaf, ShellState}; +use unshell::protocol::{Endpoint, Leaf}; + +const ID: u32 = 0x12345678; + +#[unsafe(no_mangle)] +pub fn main(_argc: i32, _argv: *const *const u8) { + let mut endpoint = Endpoint::new(ID); + let mut shell = ShellLeaf::new(ShellState::new()); + + loop { + shell.update(&mut endpoint); + } +} diff --git a/src/interface/mod.rs b/src/interface/mod.rs index 65bce89..c1372cc 100644 --- a/src/interface/mod.rs +++ b/src/interface/mod.rs @@ -13,4 +13,5 @@ pub use key::{ProcedureKey, SessionKey}; pub use store::InterfaceStore; pub use view::{ProcedureView, SessionView, SessionViewStatus}; +#[cfg(feature = "interface")] pub(crate) use store::InterfaceTarget; diff --git a/src/protocol/endpoint/hooks.rs b/src/protocol/endpoint/hooks.rs index 097d548..1cf7a03 100644 --- a/src/protocol/endpoint/hooks.rs +++ b/src/protocol/endpoint/hooks.rs @@ -24,7 +24,7 @@ impl Endpoint { for _ in 0..=HookID::MAX { let candidate = self.last_hook.next(); - if !self.hooks.contains_key(&candidate) { + if !self.has_hook(candidate) { return candidate; } } @@ -49,12 +49,14 @@ impl Endpoint { /// tests; ordinary leaf procedures should usually let packet routing pave hooks /// instead of mutating hook state by hand. pub fn accept_hook(&mut self, hook_id: HookID, peer: u32) -> Option { - self.hooks.insert(hook_id, peer) + self.hook_insert(hook_id, peer) } /// Returns true when `hook_id` is currently active. pub fn has_hook(&self, hook_id: HookID) -> bool { - self.hooks.contains_key(&hook_id) + self.hooks + .iter() + .any(|(existing_hook, _)| *existing_hook == hook_id) } /// Returns the adjacent peer currently associated with `hook_id`. @@ -63,7 +65,10 @@ impl Endpoint { /// a child for downward calls that will reply upward, or a parent for a local /// callee that will emit an upward response. pub fn hook_peer(&self, hook_id: HookID) -> Option { - self.hooks.get(&hook_id).copied() + self.hooks + .iter() + .find(|(existing_hook, _)| *existing_hook == hook_id) + .map(|(_, peer)| *peer) } /// Returns the number of active hooks on this endpoint. @@ -174,11 +179,11 @@ impl Endpoint { /// Opens or refreshes `hook_id` for the adjacent `peer` after downward routing succeeds. pub(crate) fn open_hook(&mut self, hook_id: HookID, peer: EndpointName) { - self.hooks.insert(hook_id, peer); + self.hook_insert(hook_id, peer); } /// Removes `hook_id` and reports whether it existed. pub(crate) fn close_hook(&mut self, hook_id: HookID) -> bool { - self.hooks.remove(&hook_id).is_some() + self.hook_remove(hook_id).is_some() } } diff --git a/src/protocol/endpoint/mod.rs b/src/protocol/endpoint/mod.rs index 10ac013..f2f11f2 100644 --- a/src/protocol/endpoint/mod.rs +++ b/src/protocol/endpoint/mod.rs @@ -3,11 +3,11 @@ mod routing; pub use hooks::HookID; -use alloc::{boxed::Box, vec::Vec}; +use alloc::vec::Vec; use crate::{ crypto::Counter, - protocol::{ConnectionSet, HookMap, Leaf, Packet, Path, RouteMap}, + protocol::{ConnectionSet, EndpointName, HookMap, Packet, PacketQueue, Path, RouteMap}, }; pub struct Endpoint { @@ -19,7 +19,6 @@ pub struct Endpoint { // Absolute path for this node. Must be set by some leaf pub path: Path, - pub leaves: Vec>, // Map of connections so that we can know what is connected // and which endpoints are authorities @@ -34,7 +33,13 @@ pub struct Endpoint { } impl Endpoint { - pub fn new(id: u32, leaves: Vec>) -> Self { + /// Creates endpoint routing state for one protocol node. + /// + /// Leaves are intentionally owned by the caller instead of stored behind + /// endpoint-local trait objects. That keeps minimized binaries from pulling in + /// dynamic dispatch and allocation paths when a firmware-style application uses a + /// fixed set of concrete leaves. + pub fn new(id: u32) -> Self { Self { id, // Init the hook at 0, which will increment @@ -42,25 +47,47 @@ impl Endpoint { // Set the current path as an empty vec path: Vec::new(), - leaves, - hooks: HookMap::new(), - connections: ConnectionSet::new(), - inbound: RouteMap::new(), - outbound: RouteMap::new(), + hooks: Vec::new(), + connections: Vec::new(), + inbound: Vec::new(), + outbound: Vec::new(), } } - /// Pass the endpoint state into all of the leaves - pub fn update(&mut self) { - // Grab the leaf vec temporarily so that we can iter over self - // Apparently this only swaps out pointers - let mut leaves = core::mem::take(&mut self.leaves); + /// Registers an adjacent endpoint and returns whether this is a new edge. + /// + /// Endpoint routing tables are intentionally tiny in the minimized firmware + /// profile. A linear vector keeps that profile from linking tree-map machinery + /// while preserving the old set semantics: duplicate connection registrations do + /// not create duplicate route entries. + pub fn add_connection(&mut self, remote_id: EndpointName, is_authority: bool) -> bool { + let connection = (remote_id, is_authority); - for leaf in leaves.iter_mut() { - leaf.update(self); + if self.connection_contains(remote_id, is_authority) { + false + } else { + self.connections.push(connection); + true } + } - self.leaves = leaves; + /// Removes an adjacent endpoint registration and reports whether it existed. + pub fn remove_connection(&mut self, remote_id: EndpointName, is_authority: bool) -> bool { + let Some(index) = self + .connections + .iter() + .position(|connection| *connection == (remote_id, is_authority)) + else { + return false; + }; + + self.connections.remove(index); + true + } + + /// Returns whether an adjacent endpoint is registered in the requested direction. + pub fn connection_contains(&self, remote_id: EndpointName, is_authority: bool) -> bool { + self.connections.contains(&(remote_id, is_authority)) } /// Run a function over all inbound packets with some ID then clear it. @@ -83,7 +110,7 @@ impl Endpoint { P: FnMut(&Packet) -> bool, F: FnMut(Packet), { - let Some(mut queue) = self.inbound.remove(&path) else { + let Some(mut queue) = Self::route_remove(path, &mut self.inbound) else { return; }; @@ -98,7 +125,7 @@ impl Endpoint { } if !unmatched.is_empty() { - self.inbound.entry(path).or_default().extend(unmatched); + Self::route_queue_mut(path, &mut self.inbound).extend(unmatched); } } @@ -114,7 +141,7 @@ impl Endpoint { where F: FnMut(&Packet), { - if let Some(queue) = queue.get_mut(&path) { + if let Some(queue) = Self::route_queue_mut_existing(path, queue) { for packet in queue.iter() { f(packet); } @@ -123,10 +150,95 @@ impl Endpoint { } } - pub fn iter_leaves(&mut self) -> core::slice::IterMut<'_, Box> - where - F: FnMut(&Packet), - { - self.leaves.iter_mut() + /// Appends a packet to the route queue for `endpoint`. + pub(crate) fn route_push(endpoint: EndpointName, packet: Packet, routes: &mut RouteMap) { + Self::route_queue_mut(endpoint, routes).push_back(packet); + } + + /// Returns the route queue for `endpoint` if one exists. + #[cfg(test)] + pub(crate) fn route_get(endpoint: EndpointName, routes: &RouteMap) -> Option<&PacketQueue> { + routes + .iter() + .find(|(queued_endpoint, _)| *queued_endpoint == endpoint) + .map(|(_, queue)| queue) + } + + /// Removes and returns the queue for `endpoint`. + pub(crate) fn route_remove( + endpoint: EndpointName, + routes: &mut RouteMap, + ) -> Option { + let index = routes + .iter() + .position(|(queued_endpoint, _)| *queued_endpoint == endpoint)?; + + Some(routes.remove(index).1) + } + + /// Returns whether a route queue exists for `endpoint`. + #[cfg(test)] + pub(crate) fn route_contains(endpoint: EndpointName, routes: &RouteMap) -> bool { + Self::route_get(endpoint, routes).is_some() + } + + /// Returns whether no route queues are present. + #[cfg(test)] + pub(crate) fn routes_is_empty(routes: &RouteMap) -> bool { + routes.is_empty() + } + + /// Returns the route queue for `endpoint`, creating it on first use. + fn route_queue_mut(endpoint: EndpointName, routes: &mut RouteMap) -> &mut PacketQueue { + if let Some(index) = routes + .iter() + .position(|(queued_endpoint, _)| *queued_endpoint == endpoint) + { + &mut routes[index].1 + } else { + routes.push((endpoint, PacketQueue::new())); + &mut routes.last_mut().unwrap().1 + } + } + + /// Returns the existing route queue for `endpoint` without allocating a new one. + fn route_queue_mut_existing( + endpoint: EndpointName, + routes: &mut RouteMap, + ) -> Option<&mut PacketQueue> { + routes + .iter_mut() + .find(|(queued_endpoint, _)| *queued_endpoint == endpoint) + .map(|(_, queue)| queue) + } + + /// Inserts or updates a hook and returns the previously associated peer. + pub(crate) fn hook_insert( + &mut self, + hook_id: HookID, + peer: EndpointName, + ) -> Option { + if let Some((_, existing_peer)) = self + .hooks + .iter_mut() + .find(|(existing_hook, _)| *existing_hook == hook_id) + { + let previous = *existing_peer; + *existing_peer = peer; + Some(previous) + } else { + self.hooks.push((hook_id, peer)); + None + } + } + + /// Removes a hook and returns the peer it pointed at. + pub(crate) fn hook_remove(&mut self, hook_id: HookID) -> Option { + let index = self + .hooks + .iter() + .position(|(existing_hook, _)| *existing_hook == hook_id)?; + + Some(self.hooks.remove(index).1) } } diff --git a/src/protocol/endpoint/routing.rs b/src/protocol/endpoint/routing.rs index 01af5c9..5be49b6 100644 --- a/src/protocol/endpoint/routing.rs +++ b/src/protocol/endpoint/routing.rs @@ -96,7 +96,7 @@ impl Endpoint { /// Delivers a packet to local leaves without changing hook state. fn deliver_local(&mut self, packet: Packet) -> Result<(), EndpointError> { let local_id = self.local_id()?; - self.inbound.entry(local_id).or_default().push_back(packet); + Self::route_push(local_id, packet, &mut self.inbound); Ok(()) } @@ -127,7 +127,7 @@ impl Endpoint { let end_hook = packet.end_hook; self.ensure_registered_connection(next_hop, RouteDirection::Downward)?; - self.outbound.entry(next_hop).or_default().push_back(packet); + Self::route_push(next_hop, packet, &mut self.outbound); self.apply_downward_hook_lifecycle(hook_id, end_hook, next_hop); Ok(()) } @@ -148,7 +148,7 @@ impl Endpoint { self.ensure_upward_hook_peer(hook_id, actual_peer)?; self.ensure_registered_connection(next_hop, RouteDirection::Upward)?; - self.outbound.entry(next_hop).or_default().push_back(packet); + Self::route_push(next_hop, packet, &mut self.outbound); self.apply_upward_hook_lifecycle(hook_id, end_hook); Ok(()) } @@ -195,8 +195,8 @@ impl Endpoint { /// Derives packet direction from a registered inbound adjacent peer. fn inbound_direction_from_peer(&self, remote_id: u32) -> Result { - let is_upstream = self.connections.contains(&(remote_id, true)); - let is_downstream = self.connections.contains(&(remote_id, false)); + let is_upstream = self.connection_contains(remote_id, true); + let is_downstream = self.connection_contains(remote_id, false); match (is_upstream, is_downstream) { (true, false) => Ok(RouteDirection::Downward), @@ -235,7 +235,7 @@ impl Endpoint { ) -> Result<(), EndpointError> { let is_upward = matches!(direction, RouteDirection::Upward); - if self.connections.contains(&(next_hop, is_upward)) { + if self.connection_contains(next_hop, is_upward) { Ok(()) } else { Err(EndpointError::MissingConnection { diff --git a/src/protocol/leaf_template.rs b/src/protocol/leaf_template.rs index 0040fc8..2f0da51 100644 --- a/src/protocol/leaf_template.rs +++ b/src/protocol/leaf_template.rs @@ -1,3 +1,5 @@ +mod no_procedures; + /// Declares a generated leaf wrapper using a small template-like syntax. /// /// The macro deliberately requires callers to name every generated session field. It @@ -5,6 +7,23 @@ /// macro. All real dispatch and retry behavior lives in normal Rust helpers. #[macro_export] macro_rules! unshell_leaf { + ( + $vis:vis leaf $Leaf:ident for $State:ty { + id: $id:expr, + meta: $meta:expr, + sessions { $( $session_field:ident : $Session:ty ),* $(,)? } + procedures {} + } + ) => { + $crate::__unshell_leaf_no_procedures! { + $vis leaf $Leaf for $State { + id: $id, + meta: $meta, + sessions { $( $session_field : $Session ),* } + } + } + }; + ( $vis:vis leaf $Leaf:ident for $State:ty { id: $id:expr, @@ -72,10 +91,9 @@ macro_rules! unshell_leaf { fn __unshell_update_inner( &mut self, endpoint: &mut $crate::protocol::Endpoint, - mut interface: Option<&mut $crate::interface::InterfaceStore>, ) { let leaf_id = $id; - self.__unshell_flush_all(endpoint, &mut interface); + self.__unshell_flush_all(endpoint); let Some(local_id) = endpoint.path.last().copied() else { return; @@ -89,31 +107,104 @@ macro_rules! unshell_leaf { ); for packet in packets { - self.__unshell_dispatch_packet( - endpoint, - packet, - &mut interface, - ); + self.__unshell_dispatch_packet(endpoint, packet); } $( $crate::protocol::update_session_family::<$State, $Session>( endpoint, - leaf_id, &mut self.state, &mut self.$session_field, - &mut interface, ); )* - self.__unshell_flush_all(endpoint, &mut interface); + self.__unshell_flush_all(endpoint); + } + + #[cfg(feature = "interface")] + fn __unshell_update_interface_inner( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + interface: &mut $crate::interface::InterfaceStore, + ) { + let leaf_id = $id; + self.__unshell_flush_all_interface(endpoint, interface); + + let Some(local_id) = endpoint.path.last().copied() else { + return; + }; + + let mut packets = $crate::alloc::vec::Vec::new(); + endpoint.take_inbound_matching( + local_id, + Self::__unshell_packet_is_owned, + |packet| packets.push(packet), + ); + + for packet in packets { + self.__unshell_dispatch_packet_interface(endpoint, packet, interface); + } + + $( + $crate::protocol::update_session_family_interface::<$State, $Session>( + endpoint, + leaf_id, + &mut self.state, + &mut self.$session_field, + interface, + ); + )* + + self.__unshell_flush_all_interface(endpoint, interface); } fn __unshell_dispatch_packet( &mut self, endpoint: &mut $crate::protocol::Endpoint, packet: $crate::protocol::Packet, - interface: &mut Option<&mut $crate::interface::InterfaceStore>, + ) { + let leaf_id = $id; + let _ = leaf_id; + + $( + if packet.procedure_id + == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID + { + $crate::protocol::dispatch_session::<$State, $Session>( + endpoint, + &mut self.state, + &mut self.$session_field, + packet, + ); + return; + } + )* + + $( + if packet.procedure_id + == <$Procedure as $crate::protocol::Procedure<$State>>::PROCEDURE_ID + { + let _ = stringify!($procedure_field); + $crate::protocol::dispatch_procedure::<$State, $Procedure>( + &mut self.state, + endpoint, + packet, + &mut self.outbox, + ); + return; + } + )* + + let _ = endpoint; + let _ = packet; + } + + #[cfg(feature = "interface")] + fn __unshell_dispatch_packet_interface( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + packet: $crate::protocol::Packet, + interface: &mut $crate::interface::InterfaceStore, ) { let leaf_id = $id; @@ -121,7 +212,7 @@ macro_rules! unshell_leaf { if packet.procedure_id == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID { - $crate::protocol::dispatch_session::<$State, $Session>( + $crate::protocol::dispatch_session_interface::<$State, $Session>( endpoint, leaf_id, &mut self.state, @@ -138,7 +229,7 @@ macro_rules! unshell_leaf { == <$Procedure as $crate::protocol::Procedure<$State>>::PROCEDURE_ID { let _ = stringify!($procedure_field); - $crate::protocol::dispatch_procedure::<$State, $Procedure>( + $crate::protocol::dispatch_procedure_interface::<$State, $Procedure>( leaf_id, &mut self.state, endpoint, @@ -152,16 +243,31 @@ macro_rules! unshell_leaf { let _ = endpoint; let _ = packet; + let _ = interface; } fn __unshell_flush_all( &mut self, endpoint: &mut $crate::protocol::Endpoint, - interface: &mut Option<&mut $crate::interface::InterfaceStore>, + ) { + let leaf_id = $id; + let _ = leaf_id; + + $crate::protocol::flush_leaf_outbox( + endpoint, + &mut self.outbox, + ); + } + + #[cfg(feature = "interface")] + fn __unshell_flush_all_interface( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + interface: &mut $crate::interface::InterfaceStore, ) { let leaf_id = $id; - $crate::protocol::flush_leaf_outbox( + $crate::protocol::flush_leaf_outbox_interface( endpoint, leaf_id, &mut self.outbox, @@ -175,17 +281,19 @@ macro_rules! unshell_leaf { $id } + #[inline(never)] fn update(&mut self, endpoint: &mut $crate::protocol::Endpoint) { - self.__unshell_update_inner(endpoint, None); + self.__unshell_update_inner(endpoint); } #[cfg(feature = "interface")] + #[inline(never)] fn update_interface( &mut self, endpoint: &mut $crate::protocol::Endpoint, interface: &mut $crate::interface::InterfaceStore, ) { - self.__unshell_update_inner(endpoint, Some(interface)); + self.__unshell_update_interface_inner(endpoint, interface); } #[cfg(feature = "interface")] diff --git a/src/protocol/leaf_template/no_procedures.rs b/src/protocol/leaf_template/no_procedures.rs new file mode 100644 index 0000000..ad10a1b --- /dev/null +++ b/src/protocol/leaf_template/no_procedures.rs @@ -0,0 +1,240 @@ +/// Expands the `unshell_leaf!` specialization for leaves without one-shot procedures. +/// +/// This helper stays separate from the public macro because the no-procedure shape is +/// intentionally different: it does not allocate a `LeafOutbox`, so tiny leaves such as +/// the shell leaf avoid carrying unused procedure retry machinery in the optimized +/// endpoint binary. +#[doc(hidden)] +#[macro_export] +macro_rules! __unshell_leaf_no_procedures { + ( + $vis:vis leaf $Leaf:ident for $State:ty { + id: $id:expr, + meta: $meta:expr, + sessions { $( $session_field:ident : $Session:ty ),* $(,)? } + } + ) => { + $vis struct $Leaf { + state: $State, + $( + $session_field: $crate::protocol::SessionFamily<$Session>, + )* + } + + impl $Leaf { + /// Creates the generated leaf wrapper around user-owned state. + pub fn new(state: $State) -> Self { + Self { + state, + $( + $session_field: $crate::protocol::SessionFamily::new(), + )* + } + } + + /// Returns immutable access to the user-owned leaf state. + pub fn state(&self) -> &$State { + &self.state + } + + /// Returns mutable access to the user-owned leaf state. + pub fn state_mut(&mut self) -> &mut $State { + &mut self.state + } + + /// Returns the number of active session entries across all families. + pub fn active_session_count(&self) -> usize { + 0usize $(+ self.$session_field.entries.len())* + } + + /// Returns queued packets owned by this generated leaf. + pub fn pending_packet_count(&self) -> usize { + 0usize $(+ self.$session_field.pending_packet_count())* + } + + fn __unshell_packet_is_owned(packet: &$crate::protocol::Packet) -> bool { + false + $( + || packet.procedure_id + == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID + )* + } + + fn __unshell_update_inner( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + ) { + let leaf_id = $id; + let _ = leaf_id; + + let Some(local_id) = endpoint.path.last().copied() else { + return; + }; + + let mut packets = $crate::alloc::vec::Vec::new(); + endpoint.take_inbound_matching( + local_id, + Self::__unshell_packet_is_owned, + |packet| packets.push(packet), + ); + + for packet in packets { + self.__unshell_dispatch_packet(endpoint, packet); + } + + $( + $crate::protocol::update_session_family::<$State, $Session>( + endpoint, + &mut self.state, + &mut self.$session_field, + ); + )* + } + + #[cfg(feature = "interface")] + fn __unshell_update_interface_inner( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + interface: &mut $crate::interface::InterfaceStore, + ) { + let leaf_id = $id; + let _ = leaf_id; + + let Some(local_id) = endpoint.path.last().copied() else { + return; + }; + + let mut packets = $crate::alloc::vec::Vec::new(); + endpoint.take_inbound_matching( + local_id, + Self::__unshell_packet_is_owned, + |packet| packets.push(packet), + ); + + for packet in packets { + self.__unshell_dispatch_packet_interface(endpoint, packet, interface); + } + + $( + $crate::protocol::update_session_family_interface::<$State, $Session>( + endpoint, + leaf_id, + &mut self.state, + &mut self.$session_field, + interface, + ); + )* + } + + fn __unshell_dispatch_packet( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + packet: $crate::protocol::Packet, + ) { + let leaf_id = $id; + let _ = leaf_id; + + $( + if packet.procedure_id + == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID + { + $crate::protocol::dispatch_session::<$State, $Session>( + endpoint, + &mut self.state, + &mut self.$session_field, + packet, + ); + return; + } + )* + + let _ = endpoint; + let _ = packet; + } + + #[cfg(feature = "interface")] + fn __unshell_dispatch_packet_interface( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + packet: $crate::protocol::Packet, + interface: &mut $crate::interface::InterfaceStore, + ) { + let leaf_id = $id; + + $( + if packet.procedure_id + == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID + { + $crate::protocol::dispatch_session_interface::<$State, $Session>( + endpoint, + leaf_id, + &mut self.state, + &mut self.$session_field, + packet, + interface, + ); + return; + } + )* + + let _ = endpoint; + let _ = packet; + let _ = interface; + } + } + + impl $crate::protocol::Leaf for $Leaf { + fn get_id(&self) -> u32 { + $id + } + + #[inline(never)] + fn update(&mut self, endpoint: &mut $crate::protocol::Endpoint) { + self.__unshell_update_inner(endpoint); + } + + #[cfg(feature = "interface")] + #[inline(never)] + fn update_interface( + &mut self, + endpoint: &mut $crate::protocol::Endpoint, + interface: &mut $crate::interface::InterfaceStore, + ) { + self.__unshell_update_interface_inner(endpoint, interface); + } + + #[cfg(feature = "interface")] + fn get_meta(&self) -> $crate::protocol::LeafMeta { + $meta + } + + #[cfg(feature = "interface_ratatui")] + fn render_ratatui( + &mut self, + frame: &mut $crate::protocol::ratatui::Frame<'_>, + area: $crate::protocol::ratatui::layout::Rect, + interface: &mut $crate::interface::InterfaceStore, + ) { + let leaf_id = $id; + let _ = leaf_id; + + $( + for entry in &mut self.$session_field.entries { + let view = interface.session_view_mut( + leaf_id, + <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID, + entry.hook_id, + ); + <$Session as $crate::protocol::Session<$State>>::render_ratatui( + &self.state, + &entry.state, + view, + frame, + area, + ); + } + )* + } + } + }; +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index c794b58..3c030a8 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -22,17 +22,14 @@ pub use session::*; pub use ratatui; // Various named types used for brevity -use alloc::{ - collections::{btree_map::BTreeMap, btree_set::BTreeSet, vec_deque::VecDeque}, - vec::Vec, -}; +use alloc::{collections::vec_deque::VecDeque, vec::Vec}; type Path = Vec; type EndpointName = u32; -type ConnectionSet = BTreeSet<(EndpointName, bool)>; -type HookMap = BTreeMap; +type ConnectionSet = Vec<(EndpointName, bool)>; +type HookMap = Vec<(HookID, EndpointName)>; pub type PacketQueue = VecDeque; -type RouteMap = BTreeMap; +type RouteMap = Vec<(EndpointName, PacketQueue)>; #[cfg(test)] mod tests { diff --git a/src/protocol/runtime.rs b/src/protocol/runtime.rs index bb20582..1448226 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime.rs @@ -1,11 +1,10 @@ use alloc::collections::VecDeque; -use crate::{ - interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}, - protocol::{ - Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionEntry, - SessionFamily, SessionInitError, SessionStatus, - }, +#[cfg(feature = "interface")] +use crate::interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}; +use crate::protocol::{ + Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionEntry, SessionFamily, + SessionInitError, SessionStatus, }; /// Retry queue shared by generated leaves. @@ -27,6 +26,7 @@ pub struct LeafOutbox { #[derive(Clone)] struct LeafOutboxEntry { packet: Packet, + #[cfg(feature = "interface")] target: Option, } @@ -40,7 +40,11 @@ impl LeafOutbox { /// Adds one packet to the retry queue. pub fn push(&mut self, packet: Packet) { - self.push_with_target(packet, None); + self.packets.push_back(LeafOutboxEntry { + packet, + #[cfg(feature = "interface")] + target: None, + }); } /// Adds all packets from `packets` in FIFO order. @@ -61,15 +65,16 @@ impl LeafOutbox { } /// Adds one packet with a runtime-known interface target. + #[cfg(feature = "interface")] pub(crate) fn push_for_target(&mut self, packet: Packet, target: InterfaceTarget) { - self.push_with_target(packet, Some(target)); - } - - fn push_with_target(&mut self, packet: Packet, target: Option) { - self.packets.push_back(LeafOutboxEntry { packet, target }); + self.packets.push_back(LeafOutboxEntry { + packet, + target: Some(target), + }); } /// Adds all packets with the same runtime-known interface target. + #[cfg(feature = "interface")] pub(crate) fn extend_for_target(&mut self, packets: PacketQueue, target: InterfaceTarget) { for packet in packets { self.push_for_target(packet, target); @@ -86,96 +91,36 @@ impl Default for LeafOutbox { /// Dispatches one packet into a generated session family. /// /// The macro picks `S` and the family field. This helper owns the boring details: -/// find the hook, initialize missing sessions, route rejected responses, and update -/// interface state when a caller supplied one. +/// find the hook, initialize missing sessions, and route rejected responses. The +/// interface build uses the sibling logging helper so the smallest endpoint binary +/// does not mention the interface logging types on its hot update path. pub fn dispatch_session( endpoint: &mut Endpoint, - leaf_id: u32, leaf: &mut L, family: &mut SessionFamily, packet: Packet, - interface: &mut Option<&mut InterfaceStore>, ) where S: Session, { let hook_id = packet.hook_id; let procedure_id = S::PROCEDURE_ID; - let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); - - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::Inbound { - packet: packet.clone(), - }, - ); - } - if let Some(entry) = family .entries .iter_mut() .find(|entry| entry.hook_id == hook_id) { entry.inbox.push_back(packet); - - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::SessionPacketQueued { - procedure_id, - hook_id, - }, - ); - } - return; } - let started_ns = interface.as_ref().and_then(|store| store.now_ns()); let Ok(path) = endpoint.hook_path(hook_id) else { - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::SessionRejected { - procedure_id, - hook_id, - started_ns, - finished_ns: store.now_ns(), - }, - ); - } - return; }; match S::init(leaf, packet) { Ok(state) => { family.entries.push(SessionEntry::new(hook_id, state)); - - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::SessionCreated { - procedure_id, - hook_id, - started_ns, - finished_ns: store.now_ns(), - }, - ); - } - } - Err(SessionInitError::Rejected) => { - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::SessionRejected { - procedure_id, - hook_id, - started_ns, - finished_ns: store.now_ns(), - }, - ); - } } + Err(SessionInitError::Rejected) => {} Err(SessionInitError::Response { data, end_hook }) => { let packet = Packet { hook_id, @@ -185,19 +130,7 @@ pub fn dispatch_session( data, }; - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::SessionRejected { - procedure_id, - hook_id, - started_ns, - finished_ns: store.now_ns(), - }, - ); - } - - let _ = flush_packet_with_target(endpoint, target, &packet, interface); + let _ = endpoint.add_outbound(packet); } } } @@ -205,10 +138,8 @@ pub fn dispatch_session( /// Updates every live session in one generated session family. pub fn update_session_family( endpoint: &mut Endpoint, - leaf_id: u32, leaf: &mut L, family: &mut SessionFamily, - interface: &mut Option<&mut InterfaceStore>, ) where S: Session, { @@ -217,22 +148,7 @@ pub fn update_session_family( continue; } - let started_ns = interface.as_ref().and_then(|store| store.now_ns()); let status = S::update(leaf, &mut entry.state, &mut entry.inbox, endpoint); - let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); - - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::SessionUpdated { - procedure_id: S::PROCEDURE_ID, - hook_id: entry.hook_id, - status, - started_ns, - finished_ns: store.now_ns(), - }, - ); - } if matches!(status, SessionStatus::Closed) { entry.closed = true; @@ -244,26 +160,200 @@ pub fn update_session_family( /// Dispatches one packet into a generated one-shot procedure. pub fn dispatch_procedure( + leaf: &mut L, + endpoint: &mut Endpoint, + packet: Packet, + outbox: &mut LeafOutbox, +) where + P: Procedure, +{ + let hook_id = packet.hook_id; + let mut procedure_out = + ProcedureOut::new(hook_id, parent_reply_path(endpoint), P::PROCEDURE_ID); + + P::handle(leaf, endpoint, packet, &mut procedure_out); + + let packets = procedure_out.into_packets(); + outbox.extend(packets); +} + +/// Flushes a generated leaf-level outbox through endpoint routing. +pub fn flush_leaf_outbox(endpoint: &mut Endpoint, outbox: &mut LeafOutbox) -> bool { + while let Some(entry) = outbox.packets.front() { + if endpoint.add_outbound(entry.packet.clone()).is_err() { + return false; + } + + outbox.packets.pop_front(); + } + + true +} + +/// Dispatches one packet into a generated session family with interface logging. +#[cfg(feature = "interface")] +pub fn dispatch_session_interface( + endpoint: &mut Endpoint, + leaf_id: u32, + leaf: &mut L, + family: &mut SessionFamily, + packet: Packet, + interface: &mut InterfaceStore, +) where + S: Session, +{ + let hook_id = packet.hook_id; + let procedure_id = S::PROCEDURE_ID; + let target = InterfaceTarget::session(leaf_id, procedure_id, hook_id); + + interface.record_for( + target, + InterfaceEventKind::Inbound { + packet: packet.clone(), + }, + ); + + if let Some(entry) = family + .entries + .iter_mut() + .find(|entry| entry.hook_id == hook_id) + { + entry.inbox.push_back(packet); + + interface.record_for( + target, + InterfaceEventKind::SessionPacketQueued { + procedure_id, + hook_id, + }, + ); + + return; + } + + let started_ns = interface.now_ns(); + let Ok(path) = endpoint.hook_path(hook_id) else { + interface.record_for( + target, + InterfaceEventKind::SessionRejected { + procedure_id, + hook_id, + started_ns, + finished_ns: interface.now_ns(), + }, + ); + + return; + }; + match S::init(leaf, packet) { + Ok(state) => { + family.entries.push(SessionEntry::new(hook_id, state)); + + interface.record_for( + target, + InterfaceEventKind::SessionCreated { + procedure_id, + hook_id, + started_ns, + finished_ns: interface.now_ns(), + }, + ); + } + Err(SessionInitError::Rejected) => { + interface.record_for( + target, + InterfaceEventKind::SessionRejected { + procedure_id, + hook_id, + started_ns, + finished_ns: interface.now_ns(), + }, + ); + } + Err(SessionInitError::Response { data, end_hook }) => { + let packet = Packet { + hook_id, + end_hook, + path, + procedure_id, + data, + }; + + interface.record_for( + target, + InterfaceEventKind::SessionRejected { + procedure_id, + hook_id, + started_ns, + finished_ns: interface.now_ns(), + }, + ); + + let _ = flush_packet_with_target(endpoint, target, &packet, interface); + } + } +} + +/// Updates every live session in one generated session family with interface logging. +#[cfg(feature = "interface")] +pub fn update_session_family_interface( + endpoint: &mut Endpoint, + leaf_id: u32, + leaf: &mut L, + family: &mut SessionFamily, + interface: &mut InterfaceStore, +) where + S: Session, +{ + for entry in &mut family.entries { + if entry.closed { + continue; + } + + let started_ns = interface.now_ns(); + let status = S::update(leaf, &mut entry.state, &mut entry.inbox, endpoint); + let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id); + + interface.record_for( + target, + InterfaceEventKind::SessionUpdated { + procedure_id: S::PROCEDURE_ID, + hook_id: entry.hook_id, + status, + started_ns, + finished_ns: interface.now_ns(), + }, + ); + + if matches!(status, SessionStatus::Closed) { + entry.closed = true; + } + } + + family.entries.retain(|entry| !entry.closed); +} + +/// Dispatches one packet into a generated one-shot procedure with interface logging. +#[cfg(feature = "interface")] +pub fn dispatch_procedure_interface( leaf_id: u32, leaf: &mut L, endpoint: &mut Endpoint, packet: Packet, outbox: &mut LeafOutbox, - interface: &mut Option<&mut InterfaceStore>, + interface: &mut InterfaceStore, ) where P: Procedure, { - let started_ns = interface.as_ref().and_then(|store| store.now_ns()); + let started_ns = interface.now_ns(); let target = InterfaceTarget::procedure(leaf_id, P::PROCEDURE_ID); - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::Inbound { - packet: packet.clone(), - }, - ); - } + interface.record_for( + target, + InterfaceEventKind::Inbound { + packet: packet.clone(), + }, + ); let hook_id = packet.hook_id; let mut procedure_out = @@ -273,36 +363,35 @@ pub fn dispatch_procedure( let packets = procedure_out.into_packets(); - if let Some(store) = interface.as_mut() { - store.record_for( + interface.record_for( + target, + InterfaceEventKind::ProcedureCalled { + procedure_id: P::PROCEDURE_ID, + hook_id, + started_ns, + finished_ns: interface.now_ns(), + }, + ); + + for packet in &packets { + interface.record_for( target, - InterfaceEventKind::ProcedureCalled { - procedure_id: P::PROCEDURE_ID, - hook_id, - started_ns, - finished_ns: store.now_ns(), + InterfaceEventKind::OutboundQueued { + packet: packet.clone(), }, ); - - for packet in &packets { - store.record_for( - target, - InterfaceEventKind::OutboundQueued { - packet: packet.clone(), - }, - ); - } } outbox.extend_for_target(packets, target); } -/// Flushes a generated leaf-level outbox through endpoint routing. -pub fn flush_leaf_outbox( +/// Flushes a generated leaf-level outbox through endpoint routing with interface logging. +#[cfg(feature = "interface")] +pub fn flush_leaf_outbox_interface( endpoint: &mut Endpoint, leaf_id: u32, outbox: &mut LeafOutbox, - interface: &mut Option<&mut InterfaceStore>, + interface: &mut InterfaceStore, ) -> bool { flush_outbox(endpoint, &mut outbox.packets, interface, |entry| { let target = entry.target.unwrap_or_else(|| { @@ -313,10 +402,11 @@ pub fn flush_leaf_outbox( }) } +#[cfg(feature = "interface")] fn flush_outbox( endpoint: &mut Endpoint, outbox: &mut VecDeque, - interface: &mut Option<&mut InterfaceStore>, + interface: &mut InterfaceStore, mut packet_for: impl FnMut(&T) -> (InterfaceTarget, Packet), ) -> bool { while let Some(item) = outbox.front() { @@ -332,44 +422,39 @@ fn flush_outbox( true } +#[cfg(feature = "interface")] fn flush_packet_with_target( endpoint: &mut Endpoint, target: InterfaceTarget, packet: &Packet, - interface: &mut Option<&mut InterfaceStore>, + interface: &mut InterfaceStore, ) -> bool { - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::RouteAttempt { - packet: packet.clone(), - }, - ); - } + interface.record_for( + target, + InterfaceEventKind::RouteAttempt { + packet: packet.clone(), + }, + ); match endpoint.add_outbound(packet.clone()) { Ok(()) => { - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::RouteSuccess { - packet: packet.clone(), - }, - ); - } + interface.record_for( + target, + InterfaceEventKind::RouteSuccess { + packet: packet.clone(), + }, + ); true } Err(error) => { - if let Some(store) = interface.as_mut() { - store.record_for( - target, - InterfaceEventKind::RouteFailure { - packet: packet.clone(), - error, - }, - ); - } + interface.record_for( + target, + InterfaceEventKind::RouteFailure { + packet: packet.clone(), + error, + }, + ); false } diff --git a/src/protocol/tests/merkle_sync/harness.rs b/src/protocol/tests/merkle_sync/harness.rs index 120ff9a..e3e4a5e 100644 --- a/src/protocol/tests/merkle_sync/harness.rs +++ b/src/protocol/tests/merkle_sync/harness.rs @@ -1,10 +1,13 @@ -use alloc::{boxed::Box, rc::Rc, vec}; +use alloc::{rc::Rc, vec}; use core::cell::RefCell; -use crate::protocol::Endpoint; +use crate::protocol::{Endpoint, Leaf}; use super::{ - constants::{ENDPOINT_CALLER, ENDPOINT_RESPONDENT}, + constants::{ + ENDPOINT_CALLER, ENDPOINT_RESPONDENT, LEAF_MERKLE_CALLER, LEAF_MERKLE_RESPONDENT, + LEAF_MOCK_CONNECTION, + }, leaves::{MerkleCallerLeaf, MerkleRespondentLeaf, MockConnectionLeaf}, state::{CallerReport, RespondentReport}, tree::{MerkleStore, local_fixture, remote_fixture}, @@ -19,6 +22,10 @@ use super::{ pub(super) struct MerkleHarness { pub(super) endpoint_a: Endpoint, pub(super) endpoint_b: Endpoint, + caller_leaf: MerkleCallerLeaf, + caller_connection: MockConnectionLeaf, + respondent_leaf: MerkleRespondentLeaf, + respondent_connection: MockConnectionLeaf, pub(super) caller_report: Rc>, pub(super) respondent_report: Rc>, pub(super) remote_root_hash: u32, @@ -38,37 +45,24 @@ impl MerkleHarness { let (tx_a, rx_a) = crossbeam_channel::unbounded(); let (tx_b, rx_b) = crossbeam_channel::unbounded(); - let mut endpoint_a = Endpoint::new( - ENDPOINT_CALLER, - vec![ - Box::new(MerkleCallerLeaf::new(local, caller_report.clone())), - Box::new(MockConnectionLeaf::new( - tx_b, - rx_a, - ENDPOINT_RESPONDENT, - false, - )), - ], - ); + let mut endpoint_a = Endpoint::new(ENDPOINT_CALLER); endpoint_a.path = vec![ENDPOINT_CALLER]; - let mut endpoint_b = Endpoint::new( - ENDPOINT_RESPONDENT, - vec![ - Box::new(MerkleRespondentLeaf::new(remote, respondent_report.clone())), - Box::new(MockConnectionLeaf::new(tx_a, rx_b, ENDPOINT_CALLER, true)), - ], - ); + let mut endpoint_b = Endpoint::new(ENDPOINT_RESPONDENT); endpoint_b.path = vec![ENDPOINT_CALLER, ENDPOINT_RESPONDENT]; // Register routes before the first caller update so initial packet delivery // does not depend on leaf ordering. - endpoint_a.connections.insert((ENDPOINT_RESPONDENT, false)); - endpoint_b.connections.insert((ENDPOINT_CALLER, true)); + endpoint_a.add_connection(ENDPOINT_RESPONDENT, false); + endpoint_b.add_connection(ENDPOINT_CALLER, true); Self { endpoint_a, endpoint_b, + caller_leaf: MerkleCallerLeaf::new(local, caller_report.clone()), + caller_connection: MockConnectionLeaf::new(tx_b, rx_a, ENDPOINT_RESPONDENT, false), + respondent_leaf: MerkleRespondentLeaf::new(remote, respondent_report.clone()), + respondent_connection: MockConnectionLeaf::new(tx_a, rx_b, ENDPOINT_CALLER, true), caller_report, respondent_report, remote_root_hash, @@ -77,8 +71,10 @@ impl MerkleHarness { /// Drives one deterministic protocol loop. pub(super) fn tick(&mut self) { - self.endpoint_a.update(); - self.endpoint_b.update(); + self.caller_leaf.update(&mut self.endpoint_a); + self.caller_connection.update(&mut self.endpoint_a); + self.respondent_leaf.update(&mut self.endpoint_b); + self.respondent_connection.update(&mut self.endpoint_b); } /// Runs until the caller reports completion. @@ -113,7 +109,9 @@ impl MerkleHarness { /// Verifies the requested four-leaf topology. pub(super) fn assert_four_leaf_topology(&self) { - assert_eq!(self.endpoint_a.leaves.len(), 2); - assert_eq!(self.endpoint_b.leaves.len(), 2); + assert_eq!(self.caller_leaf.get_id(), LEAF_MERKLE_CALLER); + assert_eq!(self.caller_connection.get_id(), LEAF_MOCK_CONNECTION); + assert_eq!(self.respondent_leaf.get_id(), LEAF_MERKLE_RESPONDENT); + assert_eq!(self.respondent_connection.get_id(), LEAF_MOCK_CONNECTION); } } diff --git a/src/protocol/tests/merkle_sync/leaves.rs b/src/protocol/tests/merkle_sync/leaves.rs index 3a8d22f..caf9a73 100644 --- a/src/protocol/tests/merkle_sync/leaves.rs +++ b/src/protocol/tests/merkle_sync/leaves.rs @@ -111,9 +111,7 @@ impl Leaf for MockConnectionLeaf { fn update(&mut self, endpoint: &mut Endpoint) { if !self.started { - endpoint - .connections - .insert((self.remote_id, self.is_authority)); + endpoint.add_connection(self.remote_id, self.is_authority); self.started = true; } diff --git a/src/protocol/tests/merkle_sync/state.rs b/src/protocol/tests/merkle_sync/state.rs index dfcf725..5ea611f 100644 --- a/src/protocol/tests/merkle_sync/state.rs +++ b/src/protocol/tests/merkle_sync/state.rs @@ -34,8 +34,8 @@ pub(super) enum CallerPhase { /// Test-visible caller observations. /// -/// The leaf itself lives behind `Box`, so the harness keeps a shared -/// report handle for assertions without needing downcasts. +/// The harness keeps a shared report handle so assertions can inspect caller +/// behavior without borrowing the concrete leaf for the duration of a protocol run. #[derive(Debug, Default)] pub(super) struct CallerReport { pub(super) done: bool, diff --git a/src/protocol/tests/oneshot/mod.rs b/src/protocol/tests/oneshot/mod.rs index 310c562..e319bcb 100644 --- a/src/protocol/tests/oneshot/mod.rs +++ b/src/protocol/tests/oneshot/mod.rs @@ -1,9 +1,9 @@ mod streams; mod support; -use crate::protocol::{Endpoint, EndpointError, RouteDirection}; +use crate::protocol::{Endpoint, EndpointError, Leaf, RouteDirection}; -use alloc::{boxed::Box, vec}; +use alloc::vec; use support::{ CommsLeaf, ControllerLeaf, ENDPOINT_A, ENDPOINT_B, ENDPOINT_C, ResponderLeaf, @@ -16,66 +16,63 @@ fn test_oneshot() { let (tx_a, rx_a) = crossbeam_channel::unbounded(); let (tx_b, rx_b) = crossbeam_channel::unbounded(); - let mut endpoint_a = Endpoint::new( - ENDPOINT_A, - vec![ - Box::new(ControllerLeaf { has_run: false }), - Box::new(CommsLeaf { - tx: tx_b, - rx: rx_a, - remote_id: ENDPOINT_B, - is_authority: false, - started: false, - }), - ], - ); + let mut endpoint_a = Endpoint::new(ENDPOINT_A); + let mut controller_a = ControllerLeaf { has_run: false }; + let mut comms_a = CommsLeaf { + tx: tx_b, + rx: rx_a, + remote_id: ENDPOINT_B, + is_authority: false, + started: false, + }; endpoint_a.path = vec![ENDPOINT_A]; - let mut endpoint_b = Endpoint::new( - ENDPOINT_B, - vec![ - Box::new(ResponderLeaf), - Box::new(CommsLeaf { - tx: tx_a, - rx: rx_b, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - }), - ], - ); + let mut endpoint_b = Endpoint::new(ENDPOINT_B); + let mut responder_b = ResponderLeaf; + let mut comms_b = CommsLeaf { + tx: tx_a, + rx: rx_b, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }; endpoint_b.path = vec![ENDPOINT_A, ENDPOINT_B]; // Connections are registered routing state. The comms leaves also insert them // during updates, but the first application packet should not depend on leaf order. - endpoint_a.connections.insert((ENDPOINT_B, false)); - endpoint_b.connections.insert((ENDPOINT_A, true)); + endpoint_a.add_connection(ENDPOINT_B, false); + endpoint_b.add_connection(ENDPOINT_A, true); // Cycle 1: A sends request to B - endpoint_a.update(); - endpoint_b.update(); + controller_a.update(&mut endpoint_a); + comms_a.update(&mut endpoint_a); + responder_b.update(&mut endpoint_b); + comms_b.update(&mut endpoint_b); // Cycle 2: B receives request and sends response to A - endpoint_b.update(); - endpoint_a.update(); + responder_b.update(&mut endpoint_b); + comms_b.update(&mut endpoint_b); + controller_a.update(&mut endpoint_a); + comms_a.update(&mut endpoint_a); // Cycle 3: A's CommsLeaf needs one more update to pull the packet from the channel // and put it into the inbound queue. - endpoint_a.update(); + controller_a.update(&mut endpoint_a); + comms_a.update(&mut endpoint_a); // Assertions on state assert!( - endpoint_a.inbound.contains_key(&ENDPOINT_A), + Endpoint::route_contains(ENDPOINT_A, &endpoint_a.inbound), "Endpoint A should have received response" ); assert_eq!( - endpoint_a.inbound.get(&ENDPOINT_A).unwrap().len(), + Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) + .unwrap() + .len(), 1, "Endpoint A should have exactly one packet" ); - let response = &endpoint_a - .inbound - .get(&ENDPOINT_A) + let response = &Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) .unwrap() .front() .unwrap(); @@ -92,7 +89,7 @@ fn test_oneshot() { fn inbound_downward_packet_for_local_endpoint_opens_hook() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let hook_id = endpoint.get_hook_id(); - endpoint.connections.insert((ENDPOINT_A, true)); + endpoint.add_connection(ENDPOINT_A, true); endpoint .add_inbound_from( @@ -106,7 +103,7 @@ fn inbound_downward_packet_for_local_endpoint_opens_hook() { assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B]); assert_hook_present(&endpoint, hook_id); assert_eq!(endpoint.hook_peer(hook_id), Some(ENDPOINT_A)); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] @@ -122,15 +119,15 @@ fn outbound_packet_for_local_endpoint_is_delivered_locally() { assert!(!packet.end_hook); assert_eq!(packet.data, "ABC123".as_bytes()); assert_hook_removed(&endpoint, hook_id); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] fn inbound_downward_packet_routes_to_immediate_child() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let hook_id = endpoint.get_hook_id(); - endpoint.connections.insert((ENDPOINT_A, true)); - endpoint.connections.insert((ENDPOINT_C, false)); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); endpoint .add_inbound_from( @@ -144,7 +141,7 @@ fn inbound_downward_packet_routes_to_immediate_child() { assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); assert_hook_present(&endpoint, hook_id); assert_eq!(endpoint.hook_peer(hook_id), Some(ENDPOINT_C)); - assert!(!endpoint.outbound.contains_key(&ENDPOINT_A)); + assert!(!Endpoint::route_contains(ENDPOINT_A, &endpoint.outbound)); } #[test] @@ -152,7 +149,7 @@ fn outbound_downward_packet_routes_to_immediate_child() { let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); let hook_id = endpoint.get_hook_id(); endpoint.accept_hook(hook_id, ENDPOINT_B); - endpoint.connections.insert((ENDPOINT_B, false)); + endpoint.add_connection(ENDPOINT_B, false); endpoint .add_outbound(echo_packet_with_end( @@ -166,7 +163,7 @@ fn outbound_downward_packet_routes_to_immediate_child() { assert!(packet.end_hook); assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); assert_hook_removed(&endpoint, hook_id); - assert!(!endpoint.outbound.contains_key(&ENDPOINT_C)); + assert!(!Endpoint::route_contains(ENDPOINT_C, &endpoint.outbound)); } #[test] @@ -174,8 +171,8 @@ fn inbound_upward_packet_with_hook_routes_to_parent() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let hook_id = endpoint.get_hook_id(); endpoint.accept_hook(hook_id, ENDPOINT_C); - endpoint.connections.insert((ENDPOINT_A, true)); - endpoint.connections.insert((ENDPOINT_C, false)); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); endpoint .add_inbound_from( @@ -188,15 +185,15 @@ fn inbound_upward_packet_with_hook_routes_to_parent() { assert!(packet.end_hook); assert_eq!(packet.hook_id, hook_id); assert_hook_removed(&endpoint, hook_id); - assert!(!endpoint.outbound.contains_key(&ENDPOINT_C)); + assert!(!Endpoint::route_contains(ENDPOINT_C, &endpoint.outbound)); } #[test] fn inbound_upward_packet_without_hook_is_rejected() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let hook_id = endpoint.get_hook_id(); - endpoint.connections.insert((ENDPOINT_A, true)); - endpoint.connections.insert((ENDPOINT_C, false)); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); let error = endpoint .add_inbound_from( @@ -209,16 +206,16 @@ fn inbound_upward_packet_without_hook_is_rejected() { error, EndpointError::UnknownHook { hook_id: observed_hook_id } if observed_hook_id == hook_id )); - assert!(endpoint.inbound.is_empty()); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] fn forged_upward_packet_with_unknown_hook_is_rejected() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); endpoint.accept_hook(7, ENDPOINT_C); - endpoint.connections.insert((ENDPOINT_A, true)); - endpoint.connections.insert((ENDPOINT_C, false)); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); let error = endpoint .add_inbound_from(ENDPOINT_C, echo_packet_with_end(vec![ENDPOINT_A], 99, true)) @@ -226,7 +223,7 @@ fn forged_upward_packet_with_unknown_hook_is_rejected() { assert!(matches!(error, EndpointError::UnknownHook { hook_id: 99 })); assert_hook_present(&endpoint, 7); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] @@ -234,7 +231,7 @@ fn forged_sideways_packet_is_rejected_as_incorrect_path() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let hook_id = endpoint.get_hook_id(); endpoint.accept_hook(hook_id, ENDPOINT_A); - endpoint.connections.insert((ENDPOINT_A, true)); + endpoint.add_connection(ENDPOINT_A, true); let error = endpoint .add_inbound_from( @@ -245,31 +242,29 @@ fn forged_sideways_packet_is_rejected_as_incorrect_path() { assert!(matches!(error, EndpointError::DestinationOutsideLocalTree)); assert_hook_present(&endpoint, hook_id); - assert!(endpoint.inbound.is_empty()); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] fn malformed_frame_is_dropped_by_comms_leaf() { let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); - let mut endpoint = Endpoint::new( - ENDPOINT_B, - vec![Box::new(CommsLeaf { - tx: tx_unused, - rx: rx_for_endpoint, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - })], - ); + let mut endpoint = Endpoint::new(ENDPOINT_B); + let mut comms = CommsLeaf { + tx: tx_unused, + rx: rx_for_endpoint, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }; endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; tx_to_endpoint.send(vec![0, 1, 2, 3]).unwrap(); - endpoint.update(); + comms.update(&mut endpoint); - assert!(endpoint.inbound.is_empty()); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] @@ -277,16 +272,14 @@ fn malformed_frame_does_not_block_following_valid_packet() { let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); let hook_id = 42; - let mut endpoint = Endpoint::new( - ENDPOINT_B, - vec![Box::new(CommsLeaf { - tx: tx_unused, - rx: rx_for_endpoint, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - })], - ); + let mut endpoint = Endpoint::new(ENDPOINT_B); + let mut comms = CommsLeaf { + tx: tx_unused, + rx: rx_for_endpoint, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }; endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; tx_to_endpoint.send(vec![0, 1, 2, 3]).unwrap(); @@ -297,7 +290,7 @@ fn malformed_frame_does_not_block_following_valid_packet() { .unwrap(), ) .unwrap(); - endpoint.update(); + comms.update(&mut endpoint); let packet = single_inbound_packet(&endpoint, ENDPOINT_B); assert!(!packet.end_hook); @@ -309,19 +302,17 @@ fn malformed_frame_does_not_block_following_valid_packet() { fn forged_frame_without_required_hook_is_dropped_by_comms_leaf() { let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); - let mut endpoint = Endpoint::new( - ENDPOINT_B, - vec![Box::new(CommsLeaf { - tx: tx_unused, - rx: rx_for_endpoint, - remote_id: ENDPOINT_C, - is_authority: false, - started: false, - })], - ); + let mut endpoint = Endpoint::new(ENDPOINT_B); + let mut comms = CommsLeaf { + tx: tx_unused, + rx: rx_for_endpoint, + remote_id: ENDPOINT_C, + is_authority: false, + started: false, + }; endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; endpoint.accept_hook(7, ENDPOINT_C); - endpoint.connections.insert((ENDPOINT_A, true)); + endpoint.add_connection(ENDPOINT_A, true); tx_to_endpoint .send( @@ -330,18 +321,18 @@ fn forged_frame_without_required_hook_is_dropped_by_comms_leaf() { .unwrap(), ) .unwrap(); - endpoint.update(); + comms.update(&mut endpoint); assert_hook_present(&endpoint, 7); - assert!(endpoint.inbound.is_empty()); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] fn upward_outbound_without_hook_is_rejected() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); endpoint.accept_hook(7, ENDPOINT_A); - endpoint.connections.insert((ENDPOINT_A, true)); + endpoint.add_connection(ENDPOINT_A, true); let new_hook = endpoint.get_hook_id(); @@ -354,13 +345,13 @@ fn upward_outbound_without_hook_is_rejected() { EndpointError::UnknownHook { hook_id: observed_hook_id } if observed_hook_id == new_hook )); assert_hook_present(&endpoint, 7); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] fn downward_outbound_without_hook_is_allowed() { let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); - endpoint.connections.insert((ENDPOINT_B, false)); + endpoint.add_connection(ENDPOINT_B, false); let new_hook = endpoint.get_hook_id(); @@ -368,7 +359,12 @@ fn downward_outbound_without_hook_is_allowed() { .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], new_hook)) .unwrap(); - assert_eq!(endpoint.outbound.get(&ENDPOINT_B).unwrap().len(), 1); + assert_eq!( + Endpoint::route_get(ENDPOINT_B, &endpoint.outbound) + .unwrap() + .len(), + 1 + ); assert_hook_present(&endpoint, new_hook); assert_eq!(endpoint.hook_peer(new_hook), Some(ENDPOINT_B)); } @@ -379,14 +375,14 @@ fn deeper_upward_route_uses_parent_as_next_hop() { let new_hook = endpoint.get_hook_id(); endpoint.accept_hook(new_hook, ENDPOINT_B); - endpoint.connections.insert((ENDPOINT_B, true)); + endpoint.add_connection(ENDPOINT_B, true); endpoint .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], new_hook, true)) .unwrap(); - assert!(endpoint.outbound.contains_key(&ENDPOINT_B)); - assert!(!endpoint.outbound.contains_key(&ENDPOINT_A)); + assert!(Endpoint::route_contains(ENDPOINT_B, &endpoint.outbound)); + assert!(!Endpoint::route_contains(ENDPOINT_A, &endpoint.outbound)); assert_hook_removed(&endpoint, new_hook); } @@ -407,7 +403,7 @@ fn downward_route_without_connection_is_rejected() { } )); assert_hook_removed(&endpoint, hook_id); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] @@ -428,7 +424,7 @@ fn upward_route_without_connection_is_rejected_even_with_hook() { } )); assert_hook_present(&endpoint, hook_id); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] @@ -436,7 +432,7 @@ fn end_hook_removes_hook_after_packet_is_queued() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let hook_id = endpoint.get_hook_id(); endpoint.accept_hook(hook_id, ENDPOINT_A); - endpoint.connections.insert((ENDPOINT_A, true)); + endpoint.add_connection(ENDPOINT_A, true); endpoint .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) @@ -467,29 +463,29 @@ fn failed_end_hook_route_keeps_hook_state() { } )); assert_hook_present(&endpoint, hook_id); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } #[test] fn inbound_without_absolute_path_is_rejected() { - let mut endpoint = Endpoint::new(ENDPOINT_A, vec![]); + let mut endpoint = Endpoint::new(ENDPOINT_A); let error = endpoint .add_inbound(echo_packet(vec![ENDPOINT_A], 1)) .unwrap_err(); assert!(matches!(error, EndpointError::EndpointPathUnset)); - assert!(endpoint.inbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); } #[test] fn outbound_without_absolute_path_is_rejected() { - let mut endpoint = Endpoint::new(ENDPOINT_A, vec![]); + let mut endpoint = Endpoint::new(ENDPOINT_A); let error = endpoint .add_outbound(echo_packet(vec![ENDPOINT_A], 1)) .unwrap_err(); assert!(matches!(error, EndpointError::EndpointPathUnset)); - assert!(endpoint.outbound.is_empty()); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); } diff --git a/src/protocol/tests/oneshot/streams.rs b/src/protocol/tests/oneshot/streams.rs index cde3f42..11f4940 100644 --- a/src/protocol/tests/oneshot/streams.rs +++ b/src/protocol/tests/oneshot/streams.rs @@ -3,7 +3,7 @@ use crate::protocol::{Endpoint, Leaf, Packet}; #[cfg(feature = "interface")] use crate::protocol::LeafMeta; -use alloc::{boxed::Box, format, vec, vec::Vec}; +use alloc::{format, vec, vec::Vec}; use super::support::{CommsLeaf, ENDPOINT_A, ENDPOINT_B, assert_hook_present, assert_hook_removed}; @@ -69,6 +69,20 @@ struct StreamState { next_index: usize, } +/// Concrete stream test harness that keeps leaves outside endpoint routing state. +/// +/// This mirrors firmware-style ownership: the endpoint only routes packets while the +/// caller, respondent, and connection leaves are updated explicitly in the same +/// order the old boxed endpoint dispatcher used. +struct StreamHarness { + endpoint_a: Endpoint, + endpoint_b: Endpoint, + caller_a: StreamCallerLeaf, + comms_a: CommsLeaf, + respondent_b: StreamRespondentLeaf, + comms_b: CommsLeaf, +} + impl StreamRespondentLeaf { /// Creates a respondent that will emit `total_packets` stream frames. fn new(total_packets: usize) -> Self { @@ -189,66 +203,57 @@ impl StreamRespondentLeaf { /// Each endpoint has exactly one application leaf and one mock connection leaf. The /// channel leaves are intentionally the same `CommsLeaf` used by the oneshot tests /// so stream behavior exercises the same serialization and routing boundary. -fn stream_endpoints(total_packets: usize) -> (Endpoint, Endpoint) { +fn stream_endpoints(total_packets: usize) -> StreamHarness { let (tx_a, rx_a) = crossbeam_channel::unbounded(); let (tx_b, rx_b) = crossbeam_channel::unbounded(); - let mut endpoint_a = Endpoint::new( - ENDPOINT_A, - vec![ - Box::new(StreamCallerLeaf { has_run: false }), - Box::new(CommsLeaf { - tx: tx_b, - rx: rx_a, - remote_id: ENDPOINT_B, - is_authority: false, - started: false, - }), - ], - ); + let mut endpoint_a = Endpoint::new(ENDPOINT_A); endpoint_a.path = vec![ENDPOINT_A]; - let mut endpoint_b = Endpoint::new( - ENDPOINT_B, - vec![ - Box::new(StreamRespondentLeaf::new(total_packets)), - Box::new(CommsLeaf { - tx: tx_a, - rx: rx_b, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - }), - ], - ); + let mut endpoint_b = Endpoint::new(ENDPOINT_B); endpoint_b.path = vec![ENDPOINT_A, ENDPOINT_B]; // Register routes before the first application packet so leaf order is not a // hidden prerequisite for the initial request leaving endpoint A. - endpoint_a.connections.insert((ENDPOINT_B, false)); - endpoint_b.connections.insert((ENDPOINT_A, true)); + endpoint_a.add_connection(ENDPOINT_B, false); + endpoint_b.add_connection(ENDPOINT_A, true); - (endpoint_a, endpoint_b) + StreamHarness { + endpoint_a, + endpoint_b, + caller_a: StreamCallerLeaf { has_run: false }, + comms_a: CommsLeaf { + tx: tx_b, + rx: rx_a, + remote_id: ENDPOINT_B, + is_authority: false, + started: false, + }, + respondent_b: StreamRespondentLeaf::new(total_packets), + comms_b: CommsLeaf { + tx: tx_a, + rx: rx_b, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }, + } } /// Asserts the requested two-endpoint, four-leaf topology. -fn assert_four_leaf_topology(endpoint_a: &Endpoint, endpoint_b: &Endpoint) { - assert_eq!( - endpoint_a.leaves.len(), - 2, - "caller endpoint should have two leaves" - ); - assert_eq!( - endpoint_b.leaves.len(), - 2, - "respondent endpoint should have two leaves" - ); +fn assert_four_leaf_topology(harness: &StreamHarness) { + assert_eq!(harness.caller_a.get_id(), LEAF_STREAM_CALLER); + assert_eq!(harness.comms_a.get_id(), 101); + assert_eq!(harness.respondent_b.get_id(), LEAF_STREAM_RESPONDENT); + assert_eq!(harness.comms_b.get_id(), 101); } /// Drives the initial request until it is queued locally on endpoint B. -fn deliver_stream_request(endpoint_a: &mut Endpoint, endpoint_b: &mut Endpoint) { - endpoint_a.update(); - endpoint_b.update(); +fn deliver_stream_request(harness: &mut StreamHarness) { + harness.caller_a.update(&mut harness.endpoint_a); + harness.comms_a.update(&mut harness.endpoint_a); + harness.respondent_b.update(&mut harness.endpoint_b); + harness.comms_b.update(&mut harness.endpoint_b); } /// Returns the single hook opened by the stream request on both endpoints. @@ -269,15 +274,13 @@ fn opened_stream_hook_id(endpoint_a: &Endpoint, endpoint_b: &Endpoint) -> u16 { "respondent endpoint should have exactly one stream hook" ); - let (&caller_hook, &caller_peer) = endpoint_a + let &(caller_hook, caller_peer) = endpoint_a .hooks - .iter() - .next() + .first() .expect("caller endpoint should expose the opened hook"); - let (&respondent_hook, &respondent_peer) = endpoint_b + let &(respondent_hook, respondent_peer) = endpoint_b .hooks - .iter() - .next() + .first() .expect("respondent endpoint should expose the opened hook"); assert_eq!( @@ -297,16 +300,16 @@ fn opened_stream_hook_id(endpoint_a: &Endpoint, endpoint_b: &Endpoint) -> u16 { } /// Drives one respondent stream loop and delivers any produced frame to endpoint A. -fn drive_stream_loop(endpoint_a: &mut Endpoint, endpoint_b: &mut Endpoint) { - endpoint_b.update(); - endpoint_a.update(); +fn drive_stream_loop(harness: &mut StreamHarness) { + harness.respondent_b.update(&mut harness.endpoint_b); + harness.comms_b.update(&mut harness.endpoint_b); + harness.caller_a.update(&mut harness.endpoint_a); + harness.comms_a.update(&mut harness.endpoint_a); } /// Returns stream packets that endpoint A has received so far. fn received_stream_packets(endpoint: &Endpoint) -> Vec<&Packet> { - endpoint - .inbound - .get(&ENDPOINT_A) + Endpoint::route_get(ENDPOINT_A, &endpoint.inbound) .map(|queue| queue.iter().collect()) .unwrap_or_default() } @@ -335,77 +338,77 @@ fn assert_received_stream( #[test] fn one_directional_stream_returns_one_packet_per_loop() { let total_packets = 3; - let (mut endpoint_a, mut endpoint_b) = stream_endpoints(total_packets); - assert_four_leaf_topology(&endpoint_a, &endpoint_b); + let mut harness = stream_endpoints(total_packets); + assert_four_leaf_topology(&harness); - deliver_stream_request(&mut endpoint_a, &mut endpoint_b); - let stream_hook_id = opened_stream_hook_id(&endpoint_a, &endpoint_b); + deliver_stream_request(&mut harness); + let stream_hook_id = opened_stream_hook_id(&harness.endpoint_a, &harness.endpoint_b); - assert_received_stream(&endpoint_a, 0, false, stream_hook_id); - assert_hook_present(&endpoint_a, stream_hook_id); - assert_hook_present(&endpoint_b, stream_hook_id); + assert_received_stream(&harness.endpoint_a, 0, false, stream_hook_id); + assert_hook_present(&harness.endpoint_a, stream_hook_id); + assert_hook_present(&harness.endpoint_b, stream_hook_id); for index in 0..total_packets { - drive_stream_loop(&mut endpoint_a, &mut endpoint_b); + drive_stream_loop(&mut harness); let final_seen = index + 1 == total_packets; - assert_received_stream(&endpoint_a, index + 1, final_seen, stream_hook_id); + assert_received_stream(&harness.endpoint_a, index + 1, final_seen, stream_hook_id); if final_seen { - assert_hook_removed(&endpoint_a, stream_hook_id); - assert_hook_removed(&endpoint_b, stream_hook_id); + assert_hook_removed(&harness.endpoint_a, stream_hook_id); + assert_hook_removed(&harness.endpoint_b, stream_hook_id); } else { - assert_hook_present(&endpoint_a, stream_hook_id); - assert_hook_present(&endpoint_b, stream_hook_id); + assert_hook_present(&harness.endpoint_a, stream_hook_id); + assert_hook_present(&harness.endpoint_b, stream_hook_id); } } } #[test] fn stream_does_not_emit_before_request_is_processed_by_respondent() { - let (mut endpoint_a, mut endpoint_b) = stream_endpoints(2); + let mut harness = stream_endpoints(2); - deliver_stream_request(&mut endpoint_a, &mut endpoint_b); - let stream_hook_id = opened_stream_hook_id(&endpoint_a, &endpoint_b); + deliver_stream_request(&mut harness); + let stream_hook_id = opened_stream_hook_id(&harness.endpoint_a, &harness.endpoint_b); - assert_received_stream(&endpoint_a, 0, false, stream_hook_id); - assert!(endpoint_b.outbound.is_empty()); - assert_hook_present(&endpoint_a, stream_hook_id); - assert_hook_present(&endpoint_b, stream_hook_id); + assert_received_stream(&harness.endpoint_a, 0, false, stream_hook_id); + assert!(Endpoint::routes_is_empty(&harness.endpoint_b.outbound)); + assert_hook_present(&harness.endpoint_a, stream_hook_id); + assert_hook_present(&harness.endpoint_b, stream_hook_id); } #[test] fn stream_stops_after_final_packet() { let total_packets = 2; - let (mut endpoint_a, mut endpoint_b) = stream_endpoints(total_packets); + let mut harness = stream_endpoints(total_packets); - deliver_stream_request(&mut endpoint_a, &mut endpoint_b); - let stream_hook_id = opened_stream_hook_id(&endpoint_a, &endpoint_b); - drive_stream_loop(&mut endpoint_a, &mut endpoint_b); - drive_stream_loop(&mut endpoint_a, &mut endpoint_b); - assert_received_stream(&endpoint_a, total_packets, true, stream_hook_id); - assert_hook_removed(&endpoint_b, stream_hook_id); + deliver_stream_request(&mut harness); + let stream_hook_id = opened_stream_hook_id(&harness.endpoint_a, &harness.endpoint_b); + drive_stream_loop(&mut harness); + drive_stream_loop(&mut harness); + assert_received_stream(&harness.endpoint_a, total_packets, true, stream_hook_id); + assert_hook_removed(&harness.endpoint_b, stream_hook_id); - drive_stream_loop(&mut endpoint_a, &mut endpoint_b); - assert_received_stream(&endpoint_a, total_packets, true, stream_hook_id); - assert_hook_removed(&endpoint_b, stream_hook_id); + drive_stream_loop(&mut harness); + assert_received_stream(&harness.endpoint_a, total_packets, true, stream_hook_id); + assert_hook_removed(&harness.endpoint_b, stream_hook_id); } #[test] fn failed_final_stream_route_keeps_hook_and_retries() { - let (mut endpoint_a, mut endpoint_b) = stream_endpoints(1); + let mut harness = stream_endpoints(1); - deliver_stream_request(&mut endpoint_a, &mut endpoint_b); - let stream_hook_id = opened_stream_hook_id(&endpoint_a, &endpoint_b); - endpoint_b.connections.remove(&(ENDPOINT_A, true)); + deliver_stream_request(&mut harness); + let stream_hook_id = opened_stream_hook_id(&harness.endpoint_a, &harness.endpoint_b); + harness.endpoint_b.remove_connection(ENDPOINT_A, true); - drive_stream_loop(&mut endpoint_a, &mut endpoint_b); - assert_received_stream(&endpoint_a, 0, false, stream_hook_id); - assert_hook_present(&endpoint_b, stream_hook_id); + drive_stream_loop(&mut harness); + assert_received_stream(&harness.endpoint_a, 0, false, stream_hook_id); + assert_hook_present(&harness.endpoint_b, stream_hook_id); - endpoint_b.connections.insert((ENDPOINT_A, true)); - drive_stream_loop(&mut endpoint_a, &mut endpoint_b); + harness.endpoint_b.add_connection(ENDPOINT_A, true); + drive_stream_loop(&mut harness); - assert_received_stream(&endpoint_a, 1, true, stream_hook_id); - assert_hook_removed(&endpoint_b, stream_hook_id); + assert_received_stream(&harness.endpoint_a, 1, true, stream_hook_id); + assert_hook_removed(&harness.endpoint_b, stream_hook_id); } diff --git a/src/protocol/tests/oneshot/support.rs b/src/protocol/tests/oneshot/support.rs index 2c1f19b..c1af87a 100644 --- a/src/protocol/tests/oneshot/support.rs +++ b/src/protocol/tests/oneshot/support.rs @@ -40,7 +40,7 @@ pub(super) fn echo_packet_with_end(path: Vec, hook_id: u16, end_hook: bool) /// connection table, and hook table. This helper keeps that setup explicit without /// hiding the routing state that each test is validating. pub(super) fn endpoint_at(id: u32, path: Vec) -> Endpoint { - let mut endpoint = Endpoint::new(id, vec![]); + let mut endpoint = Endpoint::new(id); endpoint.path = path; endpoint } @@ -51,9 +51,7 @@ pub(super) fn endpoint_at(id: u32, path: Vec) -> Endpoint { /// than the immediate neighbor. Tests use this helper to assert both that exactly one /// packet exists and that it was queued for the expected adjacent endpoint. pub(super) fn single_outbound_packet(endpoint: &Endpoint, next_hop: u32) -> &Packet { - let queue = endpoint - .outbound - .get(&next_hop) + let queue = Endpoint::route_get(next_hop, &endpoint.outbound) .unwrap_or_else(|| panic!("expected one outbound queue for {next_hop}")); assert_eq!(queue.len(), 1, "expected exactly one outbound packet"); queue.front().unwrap() @@ -65,9 +63,7 @@ pub(super) fn single_outbound_packet(endpoint: &Endpoint, next_hop: u32) -> &Pac /// assert against the local inbound queue instead of only checking that routing did /// not produce an error. pub(super) fn single_inbound_packet(endpoint: &Endpoint, local_id: u32) -> &Packet { - let queue = endpoint - .inbound - .get(&local_id) + let queue = Endpoint::route_get(local_id, &endpoint.inbound) .unwrap_or_else(|| panic!("expected one inbound queue for {local_id}")); assert_eq!(queue.len(), 1, "expected exactly one inbound packet"); queue.front().unwrap() @@ -154,9 +150,7 @@ impl Leaf for CommsLeaf { fn update(&mut self, endpoint: &mut Endpoint) { if !self.started { - endpoint - .connections - .insert((self.remote_id, self.is_authority)); + endpoint.add_connection(self.remote_id, self.is_authority); self.started = true; } diff --git a/unshell-leaves/leaf-pty/src/tests/interface.rs b/unshell-leaves/leaf-pty/src/tests/interface.rs index e2929f0..5a0af03 100644 --- a/unshell-leaves/leaf-pty/src/tests/interface.rs +++ b/unshell-leaves/leaf-pty/src/tests/interface.rs @@ -107,7 +107,7 @@ fn interface_update_records_failed_direct_route_without_retry() { &[], false, ); - endpoint_b.connections.remove(&(ENDPOINT_A, true)); + endpoint_b.remove_connection(ENDPOINT_A, true); leaf.update_interface(&mut endpoint_b, &mut interface); let session_key = SessionKey { @@ -121,7 +121,7 @@ fn interface_update_records_failed_direct_route_without_retry() { assert_eq!(leaf.pending_packet_count(), 0); assert_eq!(session_view.status, SessionViewStatus::Closed); - endpoint_b.connections.insert((ENDPOINT_A, true)); + endpoint_b.add_connection(ENDPOINT_A, true); leaf.update_interface(&mut endpoint_b, &mut interface); transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); let packets = drain_parent_pty_packets(&mut endpoint_a); diff --git a/unshell-leaves/leaf-pty/src/tests/session.rs b/unshell-leaves/leaf-pty/src/tests/session.rs index 285cf83..a8a30be 100644 --- a/unshell-leaves/leaf-pty/src/tests/session.rs +++ b/unshell-leaves/leaf-pty/src/tests/session.rs @@ -138,14 +138,14 @@ fn failed_final_exit_route_closes_session_without_retry() { &[], false, ); - endpoint_b.connections.remove(&(ENDPOINT_A, true)); + endpoint_b.remove_connection(ENDPOINT_A, true); leaf.update(&mut endpoint_b); assert_eq!(leaf.active_session_count(), 0); assert_eq!(leaf.pending_packet_count(), 0); assert_hook_removed(&endpoint_b, hook_id); - endpoint_b.connections.insert((ENDPOINT_A, true)); + endpoint_b.add_connection(ENDPOINT_A, true); leaf.update(&mut endpoint_b); transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); let packets = drain_parent_pty_packets(&mut endpoint_a); @@ -248,7 +248,7 @@ fn two_pty_sessions_interleave_without_crossing_hooks() { fn pty_leaf_does_not_consume_other_leaf_packets() { let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - endpoint.connections.insert((ENDPOINT_A, true)); + endpoint.add_connection(ENDPOINT_A, true); endpoint .add_inbound_from(ENDPOINT_A, pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], 7)) diff --git a/unshell-leaves/leaf-pty/src/tests/support.rs b/unshell-leaves/leaf-pty/src/tests/support.rs index 4298db7..5c1d179 100644 --- a/unshell-leaves/leaf-pty/src/tests/support.rs +++ b/unshell-leaves/leaf-pty/src/tests/support.rs @@ -12,7 +12,7 @@ pub(super) const PROC_OTHER: u32 = 31; /// Creates a bare endpoint at a known absolute path. pub(super) fn endpoint_at(id: u32, path: Vec) -> Endpoint { - let mut endpoint = Endpoint::new(id, vec![]); + let mut endpoint = Endpoint::new(id); endpoint.path = path; endpoint } @@ -22,8 +22,8 @@ pub(super) fn pty_endpoints() -> (Endpoint, Endpoint) { let mut endpoint_a = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); let mut endpoint_b = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - endpoint_a.connections.insert((ENDPOINT_B, false)); - endpoint_b.connections.insert((ENDPOINT_A, true)); + endpoint_a.add_connection(ENDPOINT_B, false); + endpoint_b.add_connection(ENDPOINT_A, true); (endpoint_a, endpoint_b) } diff --git a/unshell-leaves/leaf-shell/Cargo.toml b/unshell-leaves/leaf-shell/Cargo.toml new file mode 100644 index 0000000..b98e44f --- /dev/null +++ b/unshell-leaves/leaf-shell/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "leaf-shell" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +include.workspace = true + +[dependencies] +unshell = { workspace = true } + +[features] +default = [] +interface = ["unshell/interface"] +interface_ratatui = ["interface", "unshell/interface_ratatui"] + +[lints.rust] +elided_lifetimes_in_paths = "warn" +future_incompatible = { level = "warn", priority = -1 } +nonstandard_style = { level = "warn", priority = -1 } +rust_2018_idioms = { level = "warn", priority = -1 } +rust_2021_prelude_collisions = "warn" +semicolon_in_expressions_from_macros = "warn" +unsafe_op_in_unsafe_fn = "warn" +unused_import_braces = "warn" +unused_lifetimes = "warn" +trivial_casts = "allow" diff --git a/unshell-leaves/leaf-shell/src/lib.rs b/unshell-leaves/leaf-shell/src/lib.rs new file mode 100644 index 0000000..b09f976 --- /dev/null +++ b/unshell-leaves/leaf-shell/src/lib.rs @@ -0,0 +1,3 @@ +mod shell; + +pub use shell::{ShellLeaf, ShellState}; diff --git a/unshell-leaves/leaf-shell/src/shell/mod.rs b/unshell-leaves/leaf-shell/src/shell/mod.rs new file mode 100644 index 0000000..c5d2a7e --- /dev/null +++ b/unshell-leaves/leaf-shell/src/shell/mod.rs @@ -0,0 +1,143 @@ +use std::{ + io::Write, + process::{Child, Command, Stdio}, +}; + +use unshell::{ + crypto::hash_str_32, + protocol::{Endpoint, HookID, Packet, PacketQueue, Session, SessionInitError, SessionStatus}, + unshell_leaf, +}; + +macro_rules! version { + () => { + env!("CARGO_PKG_VERSION") + }; +} + +pub const IDENTIFIER: &str = concat!("dev.unshell.", version!(), ".shell"); +pub const SESSION_ID: &str = concat!("dev.unshell.", version!(), ".shell.session"); + +pub const IDENTIFIER_HASH: u32 = hash_str_32(IDENTIFIER); +pub const SESSION_ID_HASH: u32 = hash_str_32(SESSION_ID); + +unshell_leaf! { + pub leaf ShellLeaf for ShellState { + id: IDENTIFIER_HASH, + meta: unshell::protocol::LeafMeta { + name: "Shell", + identifier: IDENTIFIER, + version: version!(), + authors: vec!["ASTATIN3"], + }, + sessions { + shell: ShellSession, + } + procedures { + // ping: PingProcedure, + } + } +} + +/// Runtime state for the native shell leaf. +/// +/// The process state lives in per-hook [`ShellSessionState`] values because every +/// routed hook owns one child shell. The leaf-level state is intentionally empty for +/// now, but keeping a named type gives callers a stable constructor as the real shell +/// leaf grows environment and policy configuration. +#[derive(Debug, Default)] +pub struct ShellState; + +impl ShellState { + /// Creates a shell leaf state with default local process settings. + pub fn new() -> Self { + Self + } +} + +/// Per-hook native child process state. +/// +/// Hook routing is retained by the generated runtime. This state only owns the child +/// process and stream lifecycle so dropping a session cannot leave a shell orphaned. +struct ShellSession { + _hook_id: HookID, + child: Child, + stdin_closed: bool, +} + +impl ShellSession { + /// Starts the user's interactive shell for one routed session. + /// + /// `/bin/bash` matches the original shell leaf sketch. This should eventually be + /// made configurable at `ShellState`, but hard-coding it here keeps the current + /// migration focused on the session API instead of broadening shell policy. + fn spawn(hook_id: HookID) -> Result { + let child = Command::new("/bin/bash") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .spawn() + .map_err(|_| SessionInitError::rejected())?; + + Ok(Self { + _hook_id: hook_id, + child, + stdin_closed: false, + }) + } + + /// Closes the child's stdin once callers finish writing to the session. + fn close_stdin(&mut self) { + self.stdin_closed = true; + let _ = self.child.stdin.take(); + } +} + +impl Drop for ShellSession { + fn drop(&mut self) { + if matches!(self.child.try_wait(), Ok(Some(_))) { + return; + } + + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +impl Session for ShellSession { + const PROCEDURE_ID: u32 = SESSION_ID_HASH; + + fn init(_leaf: &mut ShellState, packet: Packet) -> Result { + Self::spawn(packet.hook_id) + } + + fn update( + _leaf: &mut ShellState, + session: &mut Self, + incoming: &mut PacketQueue, + _endpoint: &mut Endpoint, + ) -> SessionStatus { + while let Some(packet) = incoming.pop_front() { + if packet.end_hook { + session.close_stdin(); + } + + if packet.data.is_empty() || session.stdin_closed { + continue; + } + + let Some(stdin) = session.child.stdin.as_mut() else { + session.close_stdin(); + continue; + }; + + if stdin.write_all(&packet.data).is_err() { + session.close_stdin(); + } + } + + match session.child.try_wait() { + Ok(Some(_)) | Err(_) => SessionStatus::Closed, + Ok(None) => SessionStatus::Running, + } + } +} From 693ba7c04050c05c42365a9a5e864c201a9ae1df Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:15:31 -0600 Subject: [PATCH 08/12] Unify leaf macro expansion --- src/protocol/leaf_template.rs | 161 ++++++++----- src/protocol/leaf_template/no_procedures.rs | 240 -------------------- 2 files changed, 102 insertions(+), 299 deletions(-) delete mode 100644 src/protocol/leaf_template/no_procedures.rs diff --git a/src/protocol/leaf_template.rs b/src/protocol/leaf_template.rs index 2f0da51..eac74fe 100644 --- a/src/protocol/leaf_template.rs +++ b/src/protocol/leaf_template.rs @@ -1,29 +1,15 @@ -mod no_procedures; - /// Declares a generated leaf wrapper using a small template-like syntax. /// /// The macro deliberately requires callers to name every generated session field. It /// does not infer identifiers, inspect struct fields, or hide behavior inside a proc /// macro. All real dispatch and retry behavior lives in normal Rust helpers. +/// +/// The procedure list is handled by small internal `@...` rules instead of by +/// separate full macro expansions. That keeps the generated shape easy to audit +/// while still allowing empty `procedures {}` leaves to avoid allocating a +/// `LeafOutbox`. #[macro_export] macro_rules! unshell_leaf { - ( - $vis:vis leaf $Leaf:ident for $State:ty { - id: $id:expr, - meta: $meta:expr, - sessions { $( $session_field:ident : $Session:ty ),* $(,)? } - procedures {} - } - ) => { - $crate::__unshell_leaf_no_procedures! { - $vis leaf $Leaf for $State { - id: $id, - meta: $meta, - sessions { $( $session_field : $Session ),* } - } - } - }; - ( $vis:vis leaf $Leaf:ident for $State:ty { id: $id:expr, @@ -34,7 +20,7 @@ macro_rules! unshell_leaf { ) => { $vis struct $Leaf { state: $State, - outbox: $crate::protocol::LeafOutbox, + outbox: $crate::unshell_leaf!(@outbox_type $( $procedure_field : $Procedure ),*), $( $session_field: $crate::protocol::SessionFamily<$Session>, )* @@ -45,7 +31,7 @@ macro_rules! unshell_leaf { pub fn new(state: $State) -> Self { Self { state, - outbox: $crate::protocol::LeafOutbox::new(), + outbox: $crate::unshell_leaf!(@outbox_new $( $procedure_field : $Procedure ),*), $( $session_field: $crate::protocol::SessionFamily::new(), )* @@ -69,14 +55,16 @@ macro_rules! unshell_leaf { /// Returns queued packets owned by this generated leaf. pub fn pending_packet_count(&self) -> usize { - let mut count = self.outbox.len(); - $( - count += self.$session_field.pending_packet_count(); - )* - count + $crate::unshell_leaf!( + @outbox_len + &self.outbox; + $( $procedure_field : $Procedure ),* + ) $(+ self.$session_field.pending_packet_count())* } fn __unshell_packet_is_owned(packet: &$crate::protocol::Packet) -> bool { + let _ = packet; + false $( || packet.procedure_id @@ -93,7 +81,14 @@ macro_rules! unshell_leaf { endpoint: &mut $crate::protocol::Endpoint, ) { let leaf_id = $id; - self.__unshell_flush_all(endpoint); + let _ = leaf_id; + + $crate::unshell_leaf!( + @flush_outbox + endpoint, + &mut self.outbox; + $( $procedure_field : $Procedure ),* + ); let Some(local_id) = endpoint.path.last().copied() else { return; @@ -118,7 +113,12 @@ macro_rules! unshell_leaf { ); )* - self.__unshell_flush_all(endpoint); + $crate::unshell_leaf!( + @flush_outbox + endpoint, + &mut self.outbox; + $( $procedure_field : $Procedure ),* + ); } #[cfg(feature = "interface")] @@ -128,7 +128,16 @@ macro_rules! unshell_leaf { interface: &mut $crate::interface::InterfaceStore, ) { let leaf_id = $id; - self.__unshell_flush_all_interface(endpoint, interface); + let _ = leaf_id; + + $crate::unshell_leaf!( + @flush_outbox_interface + endpoint, + leaf_id, + &mut self.outbox, + interface; + $( $procedure_field : $Procedure ),* + ); let Some(local_id) = endpoint.path.last().copied() else { return; @@ -155,7 +164,14 @@ macro_rules! unshell_leaf { ); )* - self.__unshell_flush_all_interface(endpoint, interface); + $crate::unshell_leaf!( + @flush_outbox_interface + endpoint, + leaf_id, + &mut self.outbox, + interface; + $( $procedure_field : $Procedure ),* + ); } fn __unshell_dispatch_packet( @@ -207,6 +223,7 @@ macro_rules! unshell_leaf { interface: &mut $crate::interface::InterfaceStore, ) { let leaf_id = $id; + let _ = leaf_id; $( if packet.procedure_id @@ -245,35 +262,6 @@ macro_rules! unshell_leaf { let _ = packet; let _ = interface; } - - fn __unshell_flush_all( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - ) { - let leaf_id = $id; - let _ = leaf_id; - - $crate::protocol::flush_leaf_outbox( - endpoint, - &mut self.outbox, - ); - } - - #[cfg(feature = "interface")] - fn __unshell_flush_all_interface( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - interface: &mut $crate::interface::InterfaceStore, - ) { - let leaf_id = $id; - - $crate::protocol::flush_leaf_outbox_interface( - endpoint, - leaf_id, - &mut self.outbox, - interface, - ); - } } impl $crate::protocol::Leaf for $Leaf { @@ -309,6 +297,7 @@ macro_rules! unshell_leaf { interface: &mut $crate::interface::InterfaceStore, ) { let leaf_id = $id; + let _ = (&frame, &area, &interface, leaf_id); $( for entry in &mut self.$session_field.entries { @@ -345,4 +334,58 @@ macro_rules! unshell_leaf { } } }; + + // Select the leaf-level outbox type. Empty procedure lists use `()` so + // session-only leaves carry no retry queue, while non-empty lists share the + // normal procedure response queue. + (@outbox_type) => { + () + }; + + (@outbox_type $first_field:ident : $FirstProcedure:ty $(, $procedure_field:ident : $Procedure:ty )* $(,)?) => { + $crate::protocol::LeafOutbox + }; + + // Construct the procedure outbox selected by `@outbox_type`. + (@outbox_new) => { + () + }; + + (@outbox_new $first_field:ident : $FirstProcedure:ty $(, $procedure_field:ident : $Procedure:ty )* $(,)?) => { + $crate::protocol::LeafOutbox::new() + }; + + // Count queued procedure packets without forcing session-only leaves to own a queue. + (@outbox_len $outbox:expr;) => { + 0usize + }; + + (@outbox_len $outbox:expr; $first_field:ident : $FirstProcedure:ty $(, $procedure_field:ident : $Procedure:ty )* $(,)?) => { + $outbox.len() + }; + + // Flush queued procedure responses when the leaf declares at least one procedure. + (@flush_outbox $endpoint:expr, $outbox:expr;) => {}; + + (@flush_outbox $endpoint:expr, $outbox:expr; $first_field:ident : $FirstProcedure:ty $(, $procedure_field:ident : $Procedure:ty )* $(,)?) => {{ + let _ = stringify!($first_field); + $( + let _ = stringify!($procedure_field); + )* + + $crate::protocol::flush_leaf_outbox($endpoint, $outbox); + }}; + + // Flush queued procedure responses with interface logging when procedures exist. + (@flush_outbox_interface $endpoint:expr, $leaf_id:expr, $outbox:expr, $interface:expr;) => {}; + + (@flush_outbox_interface $endpoint:expr, $leaf_id:expr, $outbox:expr, $interface:expr; $first_field:ident : $FirstProcedure:ty $(, $procedure_field:ident : $Procedure:ty )* $(,)?) => {{ + let _ = stringify!($first_field); + $( + let _ = stringify!($procedure_field); + )* + + $crate::protocol::flush_leaf_outbox_interface($endpoint, $leaf_id, $outbox, $interface); + }}; + } diff --git a/src/protocol/leaf_template/no_procedures.rs b/src/protocol/leaf_template/no_procedures.rs deleted file mode 100644 index ad10a1b..0000000 --- a/src/protocol/leaf_template/no_procedures.rs +++ /dev/null @@ -1,240 +0,0 @@ -/// Expands the `unshell_leaf!` specialization for leaves without one-shot procedures. -/// -/// This helper stays separate from the public macro because the no-procedure shape is -/// intentionally different: it does not allocate a `LeafOutbox`, so tiny leaves such as -/// the shell leaf avoid carrying unused procedure retry machinery in the optimized -/// endpoint binary. -#[doc(hidden)] -#[macro_export] -macro_rules! __unshell_leaf_no_procedures { - ( - $vis:vis leaf $Leaf:ident for $State:ty { - id: $id:expr, - meta: $meta:expr, - sessions { $( $session_field:ident : $Session:ty ),* $(,)? } - } - ) => { - $vis struct $Leaf { - state: $State, - $( - $session_field: $crate::protocol::SessionFamily<$Session>, - )* - } - - impl $Leaf { - /// Creates the generated leaf wrapper around user-owned state. - pub fn new(state: $State) -> Self { - Self { - state, - $( - $session_field: $crate::protocol::SessionFamily::new(), - )* - } - } - - /// Returns immutable access to the user-owned leaf state. - pub fn state(&self) -> &$State { - &self.state - } - - /// Returns mutable access to the user-owned leaf state. - pub fn state_mut(&mut self) -> &mut $State { - &mut self.state - } - - /// Returns the number of active session entries across all families. - pub fn active_session_count(&self) -> usize { - 0usize $(+ self.$session_field.entries.len())* - } - - /// Returns queued packets owned by this generated leaf. - pub fn pending_packet_count(&self) -> usize { - 0usize $(+ self.$session_field.pending_packet_count())* - } - - fn __unshell_packet_is_owned(packet: &$crate::protocol::Packet) -> bool { - false - $( - || packet.procedure_id - == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID - )* - } - - fn __unshell_update_inner( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - ) { - let leaf_id = $id; - let _ = leaf_id; - - let Some(local_id) = endpoint.path.last().copied() else { - return; - }; - - let mut packets = $crate::alloc::vec::Vec::new(); - endpoint.take_inbound_matching( - local_id, - Self::__unshell_packet_is_owned, - |packet| packets.push(packet), - ); - - for packet in packets { - self.__unshell_dispatch_packet(endpoint, packet); - } - - $( - $crate::protocol::update_session_family::<$State, $Session>( - endpoint, - &mut self.state, - &mut self.$session_field, - ); - )* - } - - #[cfg(feature = "interface")] - fn __unshell_update_interface_inner( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - interface: &mut $crate::interface::InterfaceStore, - ) { - let leaf_id = $id; - let _ = leaf_id; - - let Some(local_id) = endpoint.path.last().copied() else { - return; - }; - - let mut packets = $crate::alloc::vec::Vec::new(); - endpoint.take_inbound_matching( - local_id, - Self::__unshell_packet_is_owned, - |packet| packets.push(packet), - ); - - for packet in packets { - self.__unshell_dispatch_packet_interface(endpoint, packet, interface); - } - - $( - $crate::protocol::update_session_family_interface::<$State, $Session>( - endpoint, - leaf_id, - &mut self.state, - &mut self.$session_field, - interface, - ); - )* - } - - fn __unshell_dispatch_packet( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - packet: $crate::protocol::Packet, - ) { - let leaf_id = $id; - let _ = leaf_id; - - $( - if packet.procedure_id - == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID - { - $crate::protocol::dispatch_session::<$State, $Session>( - endpoint, - &mut self.state, - &mut self.$session_field, - packet, - ); - return; - } - )* - - let _ = endpoint; - let _ = packet; - } - - #[cfg(feature = "interface")] - fn __unshell_dispatch_packet_interface( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - packet: $crate::protocol::Packet, - interface: &mut $crate::interface::InterfaceStore, - ) { - let leaf_id = $id; - - $( - if packet.procedure_id - == <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID - { - $crate::protocol::dispatch_session_interface::<$State, $Session>( - endpoint, - leaf_id, - &mut self.state, - &mut self.$session_field, - packet, - interface, - ); - return; - } - )* - - let _ = endpoint; - let _ = packet; - let _ = interface; - } - } - - impl $crate::protocol::Leaf for $Leaf { - fn get_id(&self) -> u32 { - $id - } - - #[inline(never)] - fn update(&mut self, endpoint: &mut $crate::protocol::Endpoint) { - self.__unshell_update_inner(endpoint); - } - - #[cfg(feature = "interface")] - #[inline(never)] - fn update_interface( - &mut self, - endpoint: &mut $crate::protocol::Endpoint, - interface: &mut $crate::interface::InterfaceStore, - ) { - self.__unshell_update_interface_inner(endpoint, interface); - } - - #[cfg(feature = "interface")] - fn get_meta(&self) -> $crate::protocol::LeafMeta { - $meta - } - - #[cfg(feature = "interface_ratatui")] - fn render_ratatui( - &mut self, - frame: &mut $crate::protocol::ratatui::Frame<'_>, - area: $crate::protocol::ratatui::layout::Rect, - interface: &mut $crate::interface::InterfaceStore, - ) { - let leaf_id = $id; - let _ = leaf_id; - - $( - for entry in &mut self.$session_field.entries { - let view = interface.session_view_mut( - leaf_id, - <$Session as $crate::protocol::Session<$State>>::PROCEDURE_ID, - entry.hook_id, - ); - <$Session as $crate::protocol::Session<$State>>::render_ratatui( - &self.state, - &entry.state, - view, - frame, - area, - ); - } - )* - } - } - }; -} From b4344a8d6a89ac4f0496af67cf11962d8d444597 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:35:30 -0600 Subject: [PATCH 09/12] Split generated leaf runtime modules --- .../{runtime.rs => runtime/interface.rs} | 209 +----------------- src/protocol/runtime/mod.rs | 23 ++ src/protocol/runtime/outbox.rs | 84 +++++++ src/protocol/runtime/procedure.rs | 46 ++++ src/protocol/runtime/session.rs | 73 ++++++ 5 files changed, 233 insertions(+), 202 deletions(-) rename src/protocol/{runtime.rs => runtime/interface.rs} (54%) create mode 100644 src/protocol/runtime/mod.rs create mode 100644 src/protocol/runtime/outbox.rs create mode 100644 src/protocol/runtime/procedure.rs create mode 100644 src/protocol/runtime/session.rs diff --git a/src/protocol/runtime.rs b/src/protocol/runtime/interface.rs similarity index 54% rename from src/protocol/runtime.rs rename to src/protocol/runtime/interface.rs index 1448226..5da610c 100644 --- a/src/protocol/runtime.rs +++ b/src/protocol/runtime/interface.rs @@ -1,197 +1,16 @@ use alloc::collections::VecDeque; -#[cfg(feature = "interface")] -use crate::interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}; -use crate::protocol::{ - Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionEntry, SessionFamily, - SessionInitError, SessionStatus, +use crate::{ + interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget}, + protocol::{ + Endpoint, Packet, Procedure, ProcedureOut, Session, SessionEntry, SessionFamily, + SessionInitError, SessionStatus, + }, }; -/// Retry queue shared by generated leaves. -/// -/// Leaf-level retry queue shared by generated leaves. -/// -/// Sessions route directly through `Endpoint` to keep their runtime shape small. This -/// queue remains only for one-shot procedures, whose handlers still use `ProcedureOut` -/// and should not route while the procedure is borrowing leaf state. -pub struct LeafOutbox { - packets: VecDeque, -} - -/// One packet retained by a leaf-level retry queue. -/// -/// Procedure responses from different generated branches share one queue. Storing the -/// owner beside the packet keeps route logging precise without exposing another public -/// queue type. -#[derive(Clone)] -struct LeafOutboxEntry { - packet: Packet, - #[cfg(feature = "interface")] - target: Option, -} - -impl LeafOutbox { - /// Creates an empty leaf-level outbox. - pub fn new() -> Self { - Self { - packets: VecDeque::new(), - } - } - - /// Adds one packet to the retry queue. - pub fn push(&mut self, packet: Packet) { - self.packets.push_back(LeafOutboxEntry { - packet, - #[cfg(feature = "interface")] - target: None, - }); - } - - /// Adds all packets from `packets` in FIFO order. - pub fn extend(&mut self, packets: PacketQueue) { - for packet in packets { - self.push(packet); - } - } - - /// Returns the number of queued packets. - pub fn len(&self) -> usize { - self.packets.len() - } - - /// Returns true when the queue has no pending packets. - pub fn is_empty(&self) -> bool { - self.packets.is_empty() - } - - /// Adds one packet with a runtime-known interface target. - #[cfg(feature = "interface")] - pub(crate) fn push_for_target(&mut self, packet: Packet, target: InterfaceTarget) { - self.packets.push_back(LeafOutboxEntry { - packet, - target: Some(target), - }); - } - - /// Adds all packets with the same runtime-known interface target. - #[cfg(feature = "interface")] - pub(crate) fn extend_for_target(&mut self, packets: PacketQueue, target: InterfaceTarget) { - for packet in packets { - self.push_for_target(packet, target); - } - } -} - -impl Default for LeafOutbox { - fn default() -> Self { - Self::new() - } -} - -/// Dispatches one packet into a generated session family. -/// -/// The macro picks `S` and the family field. This helper owns the boring details: -/// find the hook, initialize missing sessions, and route rejected responses. The -/// interface build uses the sibling logging helper so the smallest endpoint binary -/// does not mention the interface logging types on its hot update path. -pub fn dispatch_session( - endpoint: &mut Endpoint, - leaf: &mut L, - family: &mut SessionFamily, - packet: Packet, -) where - S: Session, -{ - let hook_id = packet.hook_id; - let procedure_id = S::PROCEDURE_ID; - if let Some(entry) = family - .entries - .iter_mut() - .find(|entry| entry.hook_id == hook_id) - { - entry.inbox.push_back(packet); - return; - } - - let Ok(path) = endpoint.hook_path(hook_id) else { - return; - }; - match S::init(leaf, packet) { - Ok(state) => { - family.entries.push(SessionEntry::new(hook_id, state)); - } - Err(SessionInitError::Rejected) => {} - Err(SessionInitError::Response { data, end_hook }) => { - let packet = Packet { - hook_id, - end_hook, - path, - procedure_id, - data, - }; - - let _ = endpoint.add_outbound(packet); - } - } -} - -/// Updates every live session in one generated session family. -pub fn update_session_family( - endpoint: &mut Endpoint, - leaf: &mut L, - family: &mut SessionFamily, -) where - S: Session, -{ - for entry in &mut family.entries { - if entry.closed { - continue; - } - - let status = S::update(leaf, &mut entry.state, &mut entry.inbox, endpoint); - - if matches!(status, SessionStatus::Closed) { - entry.closed = true; - } - } - - family.entries.retain(|entry| !entry.closed); -} - -/// Dispatches one packet into a generated one-shot procedure. -pub fn dispatch_procedure( - leaf: &mut L, - endpoint: &mut Endpoint, - packet: Packet, - outbox: &mut LeafOutbox, -) where - P: Procedure, -{ - let hook_id = packet.hook_id; - let mut procedure_out = - ProcedureOut::new(hook_id, parent_reply_path(endpoint), P::PROCEDURE_ID); - - P::handle(leaf, endpoint, packet, &mut procedure_out); - - let packets = procedure_out.into_packets(); - outbox.extend(packets); -} - -/// Flushes a generated leaf-level outbox through endpoint routing. -pub fn flush_leaf_outbox(endpoint: &mut Endpoint, outbox: &mut LeafOutbox) -> bool { - while let Some(entry) = outbox.packets.front() { - if endpoint.add_outbound(entry.packet.clone()).is_err() { - return false; - } - - outbox.packets.pop_front(); - } - - true -} +use super::{LeafOutbox, procedure::parent_reply_path}; /// Dispatches one packet into a generated session family with interface logging. -#[cfg(feature = "interface")] pub fn dispatch_session_interface( endpoint: &mut Endpoint, leaf_id: u32, @@ -295,7 +114,6 @@ pub fn dispatch_session_interface( } /// Updates every live session in one generated session family with interface logging. -#[cfg(feature = "interface")] pub fn update_session_family_interface( endpoint: &mut Endpoint, leaf_id: u32, @@ -334,7 +152,6 @@ pub fn update_session_family_interface( } /// Dispatches one packet into a generated one-shot procedure with interface logging. -#[cfg(feature = "interface")] pub fn dispatch_procedure_interface( leaf_id: u32, leaf: &mut L, @@ -386,7 +203,6 @@ pub fn dispatch_procedure_interface( } /// Flushes a generated leaf-level outbox through endpoint routing with interface logging. -#[cfg(feature = "interface")] pub fn flush_leaf_outbox_interface( endpoint: &mut Endpoint, leaf_id: u32, @@ -402,7 +218,6 @@ pub fn flush_leaf_outbox_interface( }) } -#[cfg(feature = "interface")] fn flush_outbox( endpoint: &mut Endpoint, outbox: &mut VecDeque, @@ -422,7 +237,6 @@ fn flush_outbox( true } -#[cfg(feature = "interface")] fn flush_packet_with_target( endpoint: &mut Endpoint, target: InterfaceTarget, @@ -460,12 +274,3 @@ fn flush_packet_with_target( } } } - -/// Returns the path used by generated procedure responses. -fn parent_reply_path(endpoint: &Endpoint) -> alloc::vec::Vec { - if endpoint.path.len() > 1 { - endpoint.path[..endpoint.path.len() - 1].to_vec() - } else { - endpoint.path.clone() - } -} diff --git a/src/protocol/runtime/mod.rs b/src/protocol/runtime/mod.rs new file mode 100644 index 0000000..97056ad --- /dev/null +++ b/src/protocol/runtime/mod.rs @@ -0,0 +1,23 @@ +//! Runtime helpers used by generated leaves. +//! +//! The `unshell_leaf!` macro emits static dispatch code and delegates the reusable +//! session, procedure, retry, and interface-observation behavior to this module. +//! Keeping those pieces in normal Rust makes the macro easier to audit and keeps the +//! smallest endpoint builds free of interface-only logging paths. + +mod outbox; +mod procedure; +mod session; + +#[cfg(feature = "interface")] +mod interface; + +pub use outbox::LeafOutbox; +pub use procedure::{dispatch_procedure, flush_leaf_outbox}; +pub use session::{dispatch_session, update_session_family}; + +#[cfg(feature = "interface")] +pub use interface::{ + dispatch_procedure_interface, dispatch_session_interface, flush_leaf_outbox_interface, + update_session_family_interface, +}; diff --git a/src/protocol/runtime/outbox.rs b/src/protocol/runtime/outbox.rs new file mode 100644 index 0000000..f720df4 --- /dev/null +++ b/src/protocol/runtime/outbox.rs @@ -0,0 +1,84 @@ +use alloc::collections::VecDeque; + +#[cfg(feature = "interface")] +use crate::interface::InterfaceTarget; +use crate::protocol::{Packet, PacketQueue}; + +/// Retry queue shared by generated leaves. +/// +/// Sessions route directly through `Endpoint` to keep their runtime shape small. This +/// queue remains only for one-shot procedures, whose handlers still use `ProcedureOut` +/// and should not route while the procedure is borrowing leaf state. +pub struct LeafOutbox { + pub(super) packets: VecDeque, +} + +/// One packet retained by a leaf-level retry queue. +/// +/// Procedure responses from different generated branches share one queue. Storing the +/// owner beside the packet keeps route logging precise without exposing another public +/// queue type. +#[derive(Clone)] +pub(super) struct LeafOutboxEntry { + pub(super) packet: Packet, + #[cfg(feature = "interface")] + pub(super) target: Option, +} + +impl LeafOutbox { + /// Creates an empty leaf-level outbox. + pub fn new() -> Self { + Self { + packets: VecDeque::new(), + } + } + + /// Adds one packet to the retry queue. + pub fn push(&mut self, packet: Packet) { + self.packets.push_back(LeafOutboxEntry { + packet, + #[cfg(feature = "interface")] + target: None, + }); + } + + /// Adds all packets from `packets` in FIFO order. + pub fn extend(&mut self, packets: PacketQueue) { + for packet in packets { + self.push(packet); + } + } + + /// Returns the number of queued packets. + pub fn len(&self) -> usize { + self.packets.len() + } + + /// Returns true when the queue has no pending packets. + pub fn is_empty(&self) -> bool { + self.packets.is_empty() + } + + /// Adds one packet with a runtime-known interface target. + #[cfg(feature = "interface")] + pub(crate) fn push_for_target(&mut self, packet: Packet, target: InterfaceTarget) { + self.packets.push_back(LeafOutboxEntry { + packet, + target: Some(target), + }); + } + + /// Adds all packets with the same runtime-known interface target. + #[cfg(feature = "interface")] + pub(crate) fn extend_for_target(&mut self, packets: PacketQueue, target: InterfaceTarget) { + for packet in packets { + self.push_for_target(packet, target); + } + } +} + +impl Default for LeafOutbox { + fn default() -> Self { + Self::new() + } +} diff --git a/src/protocol/runtime/procedure.rs b/src/protocol/runtime/procedure.rs new file mode 100644 index 0000000..9859283 --- /dev/null +++ b/src/protocol/runtime/procedure.rs @@ -0,0 +1,46 @@ +use alloc::vec::Vec; + +use crate::protocol::{Endpoint, Packet, Procedure, ProcedureOut}; + +use super::LeafOutbox; + +/// Dispatches one packet into a generated one-shot procedure. +pub fn dispatch_procedure( + leaf: &mut L, + endpoint: &mut Endpoint, + packet: Packet, + outbox: &mut LeafOutbox, +) where + P: Procedure, +{ + let hook_id = packet.hook_id; + let mut procedure_out = + ProcedureOut::new(hook_id, parent_reply_path(endpoint), P::PROCEDURE_ID); + + P::handle(leaf, endpoint, packet, &mut procedure_out); + + let packets = procedure_out.into_packets(); + outbox.extend(packets); +} + +/// Flushes a generated leaf-level outbox through endpoint routing. +pub fn flush_leaf_outbox(endpoint: &mut Endpoint, outbox: &mut LeafOutbox) -> bool { + while let Some(entry) = outbox.packets.front() { + if endpoint.add_outbound(entry.packet.clone()).is_err() { + return false; + } + + outbox.packets.pop_front(); + } + + true +} + +/// Returns the path used by generated procedure responses. +pub(super) fn parent_reply_path(endpoint: &Endpoint) -> Vec { + if endpoint.path.len() > 1 { + endpoint.path[..endpoint.path.len() - 1].to_vec() + } else { + endpoint.path.clone() + } +} diff --git a/src/protocol/runtime/session.rs b/src/protocol/runtime/session.rs new file mode 100644 index 0000000..5dda6ae --- /dev/null +++ b/src/protocol/runtime/session.rs @@ -0,0 +1,73 @@ +use crate::protocol::{ + Endpoint, Packet, Session, SessionEntry, SessionFamily, SessionInitError, SessionStatus, +}; + +/// Dispatches one packet into a generated session family. +/// +/// The macro picks `S` and the family field. This helper owns the boring details: +/// find the hook, initialize missing sessions, and route rejected responses. The +/// interface build uses the sibling logging helper so the smallest endpoint binary +/// does not mention the interface logging types on its hot update path. +pub fn dispatch_session( + endpoint: &mut Endpoint, + leaf: &mut L, + family: &mut SessionFamily, + packet: Packet, +) where + S: Session, +{ + let hook_id = packet.hook_id; + let procedure_id = S::PROCEDURE_ID; + if let Some(entry) = family + .entries + .iter_mut() + .find(|entry| entry.hook_id == hook_id) + { + entry.inbox.push_back(packet); + return; + } + + let Ok(path) = endpoint.hook_path(hook_id) else { + return; + }; + match S::init(leaf, packet) { + Ok(state) => { + family.entries.push(SessionEntry::new(hook_id, state)); + } + Err(SessionInitError::Rejected) => {} + Err(SessionInitError::Response { data, end_hook }) => { + let packet = Packet { + hook_id, + end_hook, + path, + procedure_id, + data, + }; + + let _ = endpoint.add_outbound(packet); + } + } +} + +/// Updates every live session in one generated session family. +pub fn update_session_family( + endpoint: &mut Endpoint, + leaf: &mut L, + family: &mut SessionFamily, +) where + S: Session, +{ + for entry in &mut family.entries { + if entry.closed { + continue; + } + + let status = S::update(leaf, &mut entry.state, &mut entry.inbox, endpoint); + + if matches!(status, SessionStatus::Closed) { + entry.closed = true; + } + } + + family.entries.retain(|entry| !entry.closed); +} From 921ea838c48e948831305931405bcc7ef9b1134f Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:39:48 -0600 Subject: [PATCH 10/12] Split protocol internals by responsibility --- src/protocol/endpoint/connections.rs | 39 ++++ src/protocol/endpoint/hook_output.rs | 77 +++++++ src/protocol/endpoint/hooks.rs | 104 +++------ src/protocol/endpoint/mod.rs | 213 ++---------------- src/protocol/endpoint/queues.rs | 127 +++++++++++ src/protocol/procedure/contract.rs | 26 +++ src/protocol/procedure/mod.rs | 7 + .../{procedure.rs => procedure/out.rs} | 29 +-- src/protocol/session.rs | 195 ---------------- src/protocol/session/contract.rs | 72 ++++++ src/protocol/session/error.rs | 39 ++++ src/protocol/session/mod.rs | 11 + src/protocol/session/status.rs | 15 ++ src/protocol/session/storage.rs | 70 ++++++ 14 files changed, 533 insertions(+), 491 deletions(-) create mode 100644 src/protocol/endpoint/connections.rs create mode 100644 src/protocol/endpoint/hook_output.rs create mode 100644 src/protocol/endpoint/queues.rs create mode 100644 src/protocol/procedure/contract.rs create mode 100644 src/protocol/procedure/mod.rs rename src/protocol/{procedure.rs => procedure/out.rs} (58%) delete mode 100644 src/protocol/session.rs create mode 100644 src/protocol/session/contract.rs create mode 100644 src/protocol/session/error.rs create mode 100644 src/protocol/session/mod.rs create mode 100644 src/protocol/session/status.rs create mode 100644 src/protocol/session/storage.rs diff --git a/src/protocol/endpoint/connections.rs b/src/protocol/endpoint/connections.rs new file mode 100644 index 0000000..6b3ad4f --- /dev/null +++ b/src/protocol/endpoint/connections.rs @@ -0,0 +1,39 @@ +use crate::protocol::{Endpoint, EndpointName}; + +impl Endpoint { + /// Registers an adjacent endpoint and returns whether this is a new edge. + /// + /// Endpoint routing tables are intentionally tiny in the minimized firmware + /// profile. A linear vector keeps that profile from linking tree-map machinery + /// while preserving the old set semantics: duplicate connection registrations do + /// not create duplicate route entries. + pub fn add_connection(&mut self, remote_id: EndpointName, is_authority: bool) -> bool { + let connection = (remote_id, is_authority); + + if self.connection_contains(remote_id, is_authority) { + false + } else { + self.connections.push(connection); + true + } + } + + /// Removes an adjacent endpoint registration and reports whether it existed. + pub fn remove_connection(&mut self, remote_id: EndpointName, is_authority: bool) -> bool { + let Some(index) = self + .connections + .iter() + .position(|connection| *connection == (remote_id, is_authority)) + else { + return false; + }; + + self.connections.remove(index); + true + } + + /// Returns whether an adjacent endpoint is registered in the requested direction. + pub fn connection_contains(&self, remote_id: EndpointName, is_authority: bool) -> bool { + self.connections.contains(&(remote_id, is_authority)) + } +} diff --git a/src/protocol/endpoint/hook_output.rs b/src/protocol/endpoint/hook_output.rs new file mode 100644 index 0000000..9afe244 --- /dev/null +++ b/src/protocol/endpoint/hook_output.rs @@ -0,0 +1,77 @@ +use alloc::vec::Vec; + +use crate::protocol::{Endpoint, EndpointError, Packet}; + +use super::HookID; + +impl Endpoint { + /// Returns the destination path for packets sent back over `hook_id`. + /// + /// Hooks record the adjacent peer that paved the return channel. This helper turns + /// that peer into the packet path required by the current router: parent peers map + /// to the parent path, and child peers map to the direct child path. Session logic + /// should not store this path itself. + pub(crate) fn hook_path(&self, hook_id: HookID) -> Result, EndpointError> { + let peer = self + .hook_peer(hook_id) + .ok_or(EndpointError::UnknownHook { hook_id })?; + + if self.path.is_empty() { + return Err(EndpointError::EndpointPathUnset); + } + + if self.path.len() > 1 && self.path[self.path.len() - 2] == peer { + Ok(self.path[..self.path.len() - 1].to_vec()) + } else { + let mut path = self.path.clone(); + path.push(peer); + Ok(path) + } + } + + /// Routes raw response data over an existing hook immediately. + /// + /// This is the compact session-output path: it avoids an intermediate context and + /// retry queue. If a final packet cannot route, the local hook is still removed so + /// an implant does not retain dead hook state forever. + pub fn send_hook_raw( + &mut self, + hook_id: HookID, + procedure_id: u32, + data: Vec, + end_hook: bool, + ) -> Result<(), EndpointError> { + let path = self.hook_path(hook_id)?; + let packet = Packet { + hook_id, + end_hook, + path, + procedure_id, + data, + }; + + let result = self.add_outbound(packet); + + if result.is_err() && end_hook { + self.close_hook(hook_id); + } + + result + } + + /// Routes a one-byte-opcode response frame over an existing hook immediately. + pub fn send_hook_frame( + &mut self, + hook_id: HookID, + procedure_id: u32, + opcode: u8, + payload: &[u8], + end_hook: bool, + ) -> Result<(), EndpointError> { + let mut data = Vec::with_capacity(payload.len() + 1); + data.push(opcode); + data.extend_from_slice(payload); + + self.send_hook_raw(hook_id, procedure_id, data, end_hook) + } +} diff --git a/src/protocol/endpoint/hooks.rs b/src/protocol/endpoint/hooks.rs index 1cf7a03..531389c 100644 --- a/src/protocol/endpoint/hooks.rs +++ b/src/protocol/endpoint/hooks.rs @@ -1,6 +1,4 @@ -use alloc::vec::Vec; - -use crate::protocol::{Endpoint, EndpointError, EndpointName, Packet}; +use crate::protocol::{Endpoint, EndpointError, EndpointName}; /// Compact identifier for one routed return channel. /// @@ -86,76 +84,6 @@ impl Endpoint { self.close_hook(hook_id) } - /// Returns the destination path for packets sent back over `hook_id`. - /// - /// Hooks record the adjacent peer that paved the return channel. This helper turns - /// that peer into the packet path required by the current router: parent peers map - /// to the parent path, and child peers map to the direct child path. Session logic - /// should not store this path itself. - pub(crate) fn hook_path(&self, hook_id: HookID) -> Result, EndpointError> { - let peer = self - .hook_peer(hook_id) - .ok_or(EndpointError::UnknownHook { hook_id })?; - - if self.path.is_empty() { - return Err(EndpointError::EndpointPathUnset); - } - - if self.path.len() > 1 && self.path[self.path.len() - 2] == peer { - Ok(self.path[..self.path.len() - 1].to_vec()) - } else { - let mut path = self.path.clone(); - path.push(peer); - Ok(path) - } - } - - /// Routes raw response data over an existing hook immediately. - /// - /// This is the compact session-output path: it avoids an intermediate context and - /// retry queue. If a final packet cannot route, the local hook is still removed so - /// an implant does not retain dead hook state forever. - pub fn send_hook_raw( - &mut self, - hook_id: HookID, - procedure_id: u32, - data: Vec, - end_hook: bool, - ) -> Result<(), EndpointError> { - let path = self.hook_path(hook_id)?; - let packet = Packet { - hook_id, - end_hook, - path, - procedure_id, - data, - }; - - let result = self.add_outbound(packet); - - if result.is_err() && end_hook { - self.close_hook(hook_id); - } - - result - } - - /// Routes a one-byte-opcode response frame over an existing hook immediately. - pub fn send_hook_frame( - &mut self, - hook_id: HookID, - procedure_id: u32, - opcode: u8, - payload: &[u8], - end_hook: bool, - ) -> Result<(), EndpointError> { - let mut data = Vec::with_capacity(payload.len() + 1); - data.push(opcode); - data.extend_from_slice(payload); - - self.send_hook_raw(hook_id, procedure_id, data, end_hook) - } - /// Validates that `actual_peer` is the peer allowed to use `hook_id`. pub(crate) fn ensure_hook_peer( &self, @@ -186,4 +114,34 @@ impl Endpoint { pub(crate) fn close_hook(&mut self, hook_id: HookID) -> bool { self.hook_remove(hook_id).is_some() } + + /// Inserts or updates a hook and returns the previously associated peer. + pub(crate) fn hook_insert( + &mut self, + hook_id: HookID, + peer: EndpointName, + ) -> Option { + if let Some((_, existing_peer)) = self + .hooks + .iter_mut() + .find(|(existing_hook, _)| *existing_hook == hook_id) + { + let previous = *existing_peer; + *existing_peer = peer; + Some(previous) + } else { + self.hooks.push((hook_id, peer)); + None + } + } + + /// Removes a hook and returns the peer it pointed at. + pub(crate) fn hook_remove(&mut self, hook_id: HookID) -> Option { + let index = self + .hooks + .iter() + .position(|(existing_hook, _)| *existing_hook == hook_id)?; + + Some(self.hooks.remove(index).1) + } } diff --git a/src/protocol/endpoint/mod.rs b/src/protocol/endpoint/mod.rs index f2f11f2..9ebb31d 100644 --- a/src/protocol/endpoint/mod.rs +++ b/src/protocol/endpoint/mod.rs @@ -1,4 +1,7 @@ +mod connections; +mod hook_output; mod hooks; +mod queues; mod routing; pub use hooks::HookID; @@ -7,28 +10,34 @@ use alloc::vec::Vec; use crate::{ crypto::Counter, - protocol::{ConnectionSet, EndpointName, HookMap, Packet, PacketQueue, Path, RouteMap}, + protocol::{ConnectionSet, HookMap, Path, RouteMap}, }; +/// Local routing state for one protocol node. +/// +/// `Endpoint` deliberately owns only route, hook, and connection tables. Leaves are +/// caller-owned concrete values, which keeps small firmware-style binaries from +/// linking dynamic leaf registries or boxed trait objects. pub struct Endpoint { - // This endpoint's identifier + /// This endpoint's identifier. pub id: u32, - // A counter that creates unique hook IDs. + /// Counter used to allocate locally unique hook ids. pub(crate) last_hook: Counter, - // Absolute path for this node. Must be set by some leaf + /// Absolute path for this node. An empty path means routing is not initialized. pub path: Path, - // Map of connections so that we can know what is connected - // and which endpoints are authorities + /// Adjacent endpoints and whether each adjacent endpoint is upstream/authority. pub connections: ConnectionSet, - // Local list of hooks. + /// Active hook id to adjacent peer mappings. pub(crate) hooks: HookMap, - // Map of endpoints to packet queues + /// Packets delivered locally and waiting for leaf consumption. pub(crate) inbound: RouteMap, + + /// Packets queued for adjacent endpoints and waiting for transport leaves. pub(crate) outbound: RouteMap, } @@ -53,192 +62,4 @@ impl Endpoint { outbound: Vec::new(), } } - - /// Registers an adjacent endpoint and returns whether this is a new edge. - /// - /// Endpoint routing tables are intentionally tiny in the minimized firmware - /// profile. A linear vector keeps that profile from linking tree-map machinery - /// while preserving the old set semantics: duplicate connection registrations do - /// not create duplicate route entries. - pub fn add_connection(&mut self, remote_id: EndpointName, is_authority: bool) -> bool { - let connection = (remote_id, is_authority); - - if self.connection_contains(remote_id, is_authority) { - false - } else { - self.connections.push(connection); - true - } - } - - /// Removes an adjacent endpoint registration and reports whether it existed. - pub fn remove_connection(&mut self, remote_id: EndpointName, is_authority: bool) -> bool { - let Some(index) = self - .connections - .iter() - .position(|connection| *connection == (remote_id, is_authority)) - else { - return false; - }; - - self.connections.remove(index); - true - } - - /// Returns whether an adjacent endpoint is registered in the requested direction. - pub fn connection_contains(&self, remote_id: EndpointName, is_authority: bool) -> bool { - self.connections.contains(&(remote_id, is_authority)) - } - - /// Run a function over all inbound packets with some ID then clear it. - pub fn take_inbound_clear(&mut self, path: u32, f: F) - where - F: FnMut(&Packet), - { - Self::take_clear(path, f, &mut self.inbound); - } - - /// Drain inbound packets for `path` that match `predicate` and preserve the rest. - /// - /// Generated leaf dispatch uses this instead of [`Self::take_inbound_clear`] so - /// one leaf can consume only its procedure or session packets without stealing - /// traffic intended for another leaf. Matching packets are passed by value because - /// most handlers need to move payload bytes into application state; unmatched - /// packets are reinserted in their original FIFO order. - pub fn take_inbound_matching(&mut self, path: u32, mut predicate: P, mut f: F) - where - P: FnMut(&Packet) -> bool, - F: FnMut(Packet), - { - let Some(mut queue) = Self::route_remove(path, &mut self.inbound) else { - return; - }; - - let mut unmatched = Vec::new(); - - while let Some(packet) = queue.pop_front() { - if predicate(&packet) { - f(packet); - } else { - unmatched.push(packet); - } - } - - if !unmatched.is_empty() { - Self::route_queue_mut(path, &mut self.inbound).extend(unmatched); - } - } - - /// Run a function over all outbound packets with some ID then clear it. - pub fn take_outbound_clear(&mut self, path: u32, f: F) - where - F: FnMut(&Packet), - { - Self::take_clear(path, f, &mut self.outbound); - } - - fn take_clear(path: u32, mut f: F, queue: &mut RouteMap) - where - F: FnMut(&Packet), - { - if let Some(queue) = Self::route_queue_mut_existing(path, queue) { - for packet in queue.iter() { - f(packet); - } - - queue.clear(); - } - } - - /// Appends a packet to the route queue for `endpoint`. - pub(crate) fn route_push(endpoint: EndpointName, packet: Packet, routes: &mut RouteMap) { - Self::route_queue_mut(endpoint, routes).push_back(packet); - } - - /// Returns the route queue for `endpoint` if one exists. - #[cfg(test)] - pub(crate) fn route_get(endpoint: EndpointName, routes: &RouteMap) -> Option<&PacketQueue> { - routes - .iter() - .find(|(queued_endpoint, _)| *queued_endpoint == endpoint) - .map(|(_, queue)| queue) - } - - /// Removes and returns the queue for `endpoint`. - pub(crate) fn route_remove( - endpoint: EndpointName, - routes: &mut RouteMap, - ) -> Option { - let index = routes - .iter() - .position(|(queued_endpoint, _)| *queued_endpoint == endpoint)?; - - Some(routes.remove(index).1) - } - - /// Returns whether a route queue exists for `endpoint`. - #[cfg(test)] - pub(crate) fn route_contains(endpoint: EndpointName, routes: &RouteMap) -> bool { - Self::route_get(endpoint, routes).is_some() - } - - /// Returns whether no route queues are present. - #[cfg(test)] - pub(crate) fn routes_is_empty(routes: &RouteMap) -> bool { - routes.is_empty() - } - - /// Returns the route queue for `endpoint`, creating it on first use. - fn route_queue_mut(endpoint: EndpointName, routes: &mut RouteMap) -> &mut PacketQueue { - if let Some(index) = routes - .iter() - .position(|(queued_endpoint, _)| *queued_endpoint == endpoint) - { - &mut routes[index].1 - } else { - routes.push((endpoint, PacketQueue::new())); - &mut routes.last_mut().unwrap().1 - } - } - - /// Returns the existing route queue for `endpoint` without allocating a new one. - fn route_queue_mut_existing( - endpoint: EndpointName, - routes: &mut RouteMap, - ) -> Option<&mut PacketQueue> { - routes - .iter_mut() - .find(|(queued_endpoint, _)| *queued_endpoint == endpoint) - .map(|(_, queue)| queue) - } - - /// Inserts or updates a hook and returns the previously associated peer. - pub(crate) fn hook_insert( - &mut self, - hook_id: HookID, - peer: EndpointName, - ) -> Option { - if let Some((_, existing_peer)) = self - .hooks - .iter_mut() - .find(|(existing_hook, _)| *existing_hook == hook_id) - { - let previous = *existing_peer; - *existing_peer = peer; - Some(previous) - } else { - self.hooks.push((hook_id, peer)); - None - } - } - - /// Removes a hook and returns the peer it pointed at. - pub(crate) fn hook_remove(&mut self, hook_id: HookID) -> Option { - let index = self - .hooks - .iter() - .position(|(existing_hook, _)| *existing_hook == hook_id)?; - - Some(self.hooks.remove(index).1) - } } diff --git a/src/protocol/endpoint/queues.rs b/src/protocol/endpoint/queues.rs new file mode 100644 index 0000000..ed09b9e --- /dev/null +++ b/src/protocol/endpoint/queues.rs @@ -0,0 +1,127 @@ +use alloc::vec::Vec; + +use crate::protocol::{Endpoint, EndpointName, Packet, PacketQueue, RouteMap}; + +impl Endpoint { + /// Runs a function over all inbound packets for `path`, then clears that queue. + pub fn take_inbound_clear(&mut self, path: u32, f: F) + where + F: FnMut(&Packet), + { + Self::take_clear(path, f, &mut self.inbound); + } + + /// Drain inbound packets for `path` that match `predicate` and preserve the rest. + /// + /// Generated leaf dispatch uses this instead of [`Self::take_inbound_clear`] so + /// one leaf can consume only its procedure or session packets without stealing + /// traffic intended for another leaf. Matching packets are passed by value because + /// most handlers need to move payload bytes into application state; unmatched + /// packets are reinserted in their original FIFO order. + pub fn take_inbound_matching(&mut self, path: u32, mut predicate: P, mut f: F) + where + P: FnMut(&Packet) -> bool, + F: FnMut(Packet), + { + let Some(mut queue) = Self::route_remove(path, &mut self.inbound) else { + return; + }; + + let mut unmatched = Vec::new(); + + while let Some(packet) = queue.pop_front() { + if predicate(&packet) { + f(packet); + } else { + unmatched.push(packet); + } + } + + if !unmatched.is_empty() { + Self::route_queue_mut(path, &mut self.inbound).extend(unmatched); + } + } + + /// Runs a function over all outbound packets for `path`, then clears that queue. + pub fn take_outbound_clear(&mut self, path: u32, f: F) + where + F: FnMut(&Packet), + { + Self::take_clear(path, f, &mut self.outbound); + } + + fn take_clear(path: u32, mut f: F, queue: &mut RouteMap) + where + F: FnMut(&Packet), + { + if let Some(queue) = Self::route_queue_mut_existing(path, queue) { + for packet in queue.iter() { + f(packet); + } + + queue.clear(); + } + } + + /// Appends a packet to the route queue for `endpoint`. + pub(crate) fn route_push(endpoint: EndpointName, packet: Packet, routes: &mut RouteMap) { + Self::route_queue_mut(endpoint, routes).push_back(packet); + } + + /// Returns the route queue for `endpoint` if one exists. + #[cfg(test)] + pub(crate) fn route_get(endpoint: EndpointName, routes: &RouteMap) -> Option<&PacketQueue> { + routes + .iter() + .find(|(queued_endpoint, _)| *queued_endpoint == endpoint) + .map(|(_, queue)| queue) + } + + /// Removes and returns the queue for `endpoint`. + pub(crate) fn route_remove( + endpoint: EndpointName, + routes: &mut RouteMap, + ) -> Option { + let index = routes + .iter() + .position(|(queued_endpoint, _)| *queued_endpoint == endpoint)?; + + Some(routes.remove(index).1) + } + + /// Returns whether a route queue exists for `endpoint`. + #[cfg(test)] + pub(crate) fn route_contains(endpoint: EndpointName, routes: &RouteMap) -> bool { + Self::route_get(endpoint, routes).is_some() + } + + /// Returns whether no route queues are present. + #[cfg(test)] + pub(crate) fn routes_is_empty(routes: &RouteMap) -> bool { + routes.is_empty() + } + + /// Returns the route queue for `endpoint`, creating it on first use. + fn route_queue_mut(endpoint: EndpointName, routes: &mut RouteMap) -> &mut PacketQueue { + if let Some(index) = routes + .iter() + .position(|(queued_endpoint, _)| *queued_endpoint == endpoint) + { + &mut routes[index].1 + } else { + routes.push((endpoint, PacketQueue::new())); + &mut routes.last_mut().unwrap().1 + } + } + + /// Returns the existing route queue for `endpoint` without allocating a new one. + fn route_queue_mut_existing( + endpoint: EndpointName, + routes: &mut RouteMap, + ) -> Option<&mut PacketQueue> { + routes + .iter_mut() + .find(|(queued_endpoint, _)| *queued_endpoint == endpoint) + .map(|(_, queue)| queue) + } +} diff --git a/src/protocol/procedure/contract.rs b/src/protocol/procedure/contract.rs new file mode 100644 index 0000000..511d704 --- /dev/null +++ b/src/protocol/procedure/contract.rs @@ -0,0 +1,26 @@ +use crate::protocol::{Endpoint, Packet, ProcedureOut}; + +#[cfg(feature = "interface_ratatui")] +use crate::interface::ProcedureView; + +/// Contract implemented by one generated one-packet procedure handler. +/// +/// Procedures are for stateless or short-lived operations such as ping, capabilities, +/// or health checks. Long-running conversations should use [`Session`](crate::protocol::Session) +/// so final packet cleanup and retries remain tied to hook state. +pub trait Procedure { + /// Outer packet procedure id handled by this procedure. + const PROCEDURE_ID: u32; + + /// Handles one packet and optionally queues response packets in `out`. + fn handle(leaf: &mut L, endpoint: &mut Endpoint, packet: Packet, out: &mut ProcedureOut); + + #[cfg(feature = "interface_ratatui")] + fn render_ratatui( + _: &L, + _: &mut ProcedureView, + _: &mut ratatui::Frame<'_>, + _: ratatui::layout::Rect, + ) { + } +} diff --git a/src/protocol/procedure/mod.rs b/src/protocol/procedure/mod.rs new file mode 100644 index 0000000..63ff28d --- /dev/null +++ b/src/protocol/procedure/mod.rs @@ -0,0 +1,7 @@ +//! One-shot procedure contracts and response output helpers. + +mod contract; +mod out; + +pub use contract::Procedure; +pub use out::ProcedureOut; diff --git a/src/protocol/procedure.rs b/src/protocol/procedure/out.rs similarity index 58% rename from src/protocol/procedure.rs rename to src/protocol/procedure/out.rs index 0ee7a00..9c13f53 100644 --- a/src/protocol/procedure.rs +++ b/src/protocol/procedure/out.rs @@ -1,33 +1,8 @@ use alloc::vec::Vec; -use crate::protocol::{Endpoint, HookID, Packet, PacketQueue}; +use crate::protocol::{HookID, Packet, PacketQueue}; -#[cfg(feature = "interface_ratatui")] -use crate::interface::ProcedureView; - -/// Contract implemented by one generated one-packet procedure handler. -/// -/// Procedures are for stateless or short-lived operations such as ping, capabilities, -/// or health checks. Long-running conversations should use [`Session`] so final -/// packet cleanup and retries remain tied to hook state. -pub trait Procedure { - /// Outer packet procedure id handled by this procedure. - const PROCEDURE_ID: u32; - - /// Handles one packet and optionally queues response packets in `out`. - fn handle(leaf: &mut L, endpoint: &mut Endpoint, packet: Packet, out: &mut ProcedureOut); - - #[cfg(feature = "interface_ratatui")] - fn render_ratatui( - _: &L, - _: &mut ProcedureView, - _: &mut ratatui::Frame<'_>, - _: ratatui::layout::Rect, - ) { - } -} - -/// Output accumulator passed to [`Procedure::handle`]. +/// Output accumulator passed to [`Procedure::handle`](super::Procedure::handle). pub struct ProcedureOut { hook_id: HookID, reply_path: Vec, diff --git a/src/protocol/session.rs b/src/protocol/session.rs deleted file mode 100644 index 405b990..0000000 --- a/src/protocol/session.rs +++ /dev/null @@ -1,195 +0,0 @@ -use alloc::vec::Vec; - -use crate::protocol::{Endpoint, HookID, Packet, PacketQueue}; - -#[cfg(feature = "interface_ratatui")] -use crate::interface::SessionView; - -/// Contract implemented by one hook-backed generated session family. -/// -/// A session family maps one outer `procedure_id` to many live hook instances. The -/// generated leaf owns packet grouping, retry-safe output flushing, and final cleanup; -/// the session value owns one hook's application behavior and mutable state. -/// -/// # Example -/// -/// ```rust,ignore -/// impl Session for MySessionState { -/// const PROCEDURE_ID: u32 = 7; -/// -/// fn init( -/// leaf: &mut MyLeafState, -/// packet: Packet, -/// ) -> Result { -/// Ok(MySessionState::from_open(leaf, packet)) -/// } -/// -/// fn update( -/// leaf: &mut MyLeafState, -/// session: &mut Self, -/// incoming: &mut PacketQueue, -/// endpoint: &mut Endpoint, -/// ) -> SessionStatus { -/// while let Some(packet) = incoming.pop_front() { -/// session.apply(leaf, packet, endpoint); -/// } -/// SessionStatus::Running -/// } -/// } -/// ``` -pub trait Session: Sized { - /// Outer packet procedure id used by every packet in this session family. - const PROCEDURE_ID: u32; - - /// Creates one session value from a packet whose hook has no active session. - /// - /// The generated runtime derives all response routing from hook state. Session - /// initialization therefore returns only application state or a protocol-level - /// rejection; it never stores or receives a caller reply path. - fn init(leaf: &mut L, packet: Packet) -> Result; - - /// Advances one active hook session. - /// - /// The generated leaf calls this for every live session on each update tick so - /// sessions can poll external workers even when no new packet arrived. Session - /// output is routed immediately through `endpoint`; callers that need retry - /// semantics should keep their own compact application state and retry on a later - /// tick. - fn update( - leaf: &mut L, - session: &mut Self, - incoming: &mut PacketQueue, - endpoint: &mut Endpoint, - ) -> SessionStatus; - - #[cfg(feature = "interface_ratatui")] - fn render_ratatui( - _: &L, - _: &Self, - _: &mut SessionView, - _: &mut ratatui::Frame<'_>, - _: ratatui::layout::Rect, - ) { - } -} - -/// Error returned when a packet cannot create a new session. -pub enum SessionInitError { - /// The packet was intentionally consumed without creating state or sending output. - Rejected, - - /// The packet was rejected with response data that should be sent on the same hook. - Response { - /// Raw `Packet::data` for the response frame. - data: Vec, - - /// Whether the response should close the hook after successful routing. - end_hook: bool, - }, -} - -impl SessionInitError { - /// Creates a silent session rejection. - pub fn rejected() -> Self { - Self::Rejected - } - - /// Creates a non-final response for a rejected session open. - pub fn response(data: Vec) -> Self { - Self::Response { - data, - end_hook: false, - } - } - - /// Creates a final response for a rejected session open. - pub fn response_final(data: Vec) -> Self { - Self::Response { - data, - end_hook: true, - } - } -} - -/// Session lifecycle status returned from [`Session::update`]. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SessionStatus { - /// The session is active and should receive future update ticks. - Running, - - /// The session is winding down but still needs future update ticks. - Closing, - - /// The session has finished application work. - /// - /// The generated leaf removes the entry after the update tick. Final packets are - /// routed immediately by the session before returning this status. - Closed, -} - -/// Storage entry used by macro-generated session stores. -/// -/// The fields are public so generated code in downstream crates can keep the update -/// loop straightforward and static. Handwritten leaves may also use this type, but it -/// is intentionally small rather than a full session framework. -pub struct SessionEntry { - /// Hook id associated with this live session. - pub hook_id: HookID, - - /// Application-owned session state. - pub state: S, - - /// Packets delivered for this hook but not yet consumed by the session. - pub inbox: PacketQueue, - - /// Whether application logic has finished and should be removed after update. - pub closed: bool, -} - -/// Generated storage for one session family. -/// -/// The macro only names this field and picks the concrete `Session` type. All update, -/// retry, and cleanup behavior lives in normal Rust helpers so the template stays -/// small and readable. -pub struct SessionFamily { - /// Active hook-backed sessions for this family. - pub entries: Vec>, -} - -impl SessionFamily { - /// Creates an empty session family. - pub fn new() -> Self { - Self { - entries: Vec::new(), - } - } - - /// Counts packets retained by this family for retry or future session work. - pub fn pending_packet_count(&self) -> usize { - let mut count = 0usize; - - for entry in &self.entries { - count += entry.inbox.len(); - } - - count - } -} - -impl Default for SessionFamily { - fn default() -> Self { - Self::new() - } -} - -impl SessionEntry { - /// Creates one active session entry for `hook_id`. - pub fn new(hook_id: HookID, state: S) -> Self { - Self { - hook_id, - state, - inbox: PacketQueue::new(), - closed: false, - } - } -} diff --git a/src/protocol/session/contract.rs b/src/protocol/session/contract.rs new file mode 100644 index 0000000..97dcf81 --- /dev/null +++ b/src/protocol/session/contract.rs @@ -0,0 +1,72 @@ +use crate::protocol::{Endpoint, Packet, PacketQueue, SessionInitError, SessionStatus}; + +#[cfg(feature = "interface_ratatui")] +use crate::interface::SessionView; + +/// Contract implemented by one hook-backed generated session family. +/// +/// A session family maps one outer `procedure_id` to many live hook instances. The +/// generated leaf owns packet grouping, retry-safe output flushing, and final cleanup; +/// the session value owns one hook's application behavior and mutable state. +/// +/// # Example +/// +/// ```rust,ignore +/// impl Session for MySessionState { +/// const PROCEDURE_ID: u32 = 7; +/// +/// fn init( +/// leaf: &mut MyLeafState, +/// packet: Packet, +/// ) -> Result { +/// Ok(MySessionState::from_open(leaf, packet)) +/// } +/// +/// fn update( +/// leaf: &mut MyLeafState, +/// session: &mut Self, +/// incoming: &mut PacketQueue, +/// endpoint: &mut Endpoint, +/// ) -> SessionStatus { +/// while let Some(packet) = incoming.pop_front() { +/// session.apply(leaf, packet, endpoint); +/// } +/// SessionStatus::Running +/// } +/// } +/// ``` +pub trait Session: Sized { + /// Outer packet procedure id used by every packet in this session family. + const PROCEDURE_ID: u32; + + /// Creates one session value from a packet whose hook has no active session. + /// + /// The generated runtime derives all response routing from hook state. Session + /// initialization therefore returns only application state or a protocol-level + /// rejection; it never stores or receives a caller reply path. + fn init(leaf: &mut L, packet: Packet) -> Result; + + /// Advances one active hook session. + /// + /// The generated leaf calls this for every live session on each update tick so + /// sessions can poll external workers even when no new packet arrived. Session + /// output is routed immediately through `endpoint`; callers that need retry + /// semantics should keep their own compact application state and retry on a later + /// tick. + fn update( + leaf: &mut L, + session: &mut Self, + incoming: &mut PacketQueue, + endpoint: &mut Endpoint, + ) -> SessionStatus; + + #[cfg(feature = "interface_ratatui")] + fn render_ratatui( + _: &L, + _: &Self, + _: &mut SessionView, + _: &mut ratatui::Frame<'_>, + _: ratatui::layout::Rect, + ) { + } +} diff --git a/src/protocol/session/error.rs b/src/protocol/session/error.rs new file mode 100644 index 0000000..0e6bde4 --- /dev/null +++ b/src/protocol/session/error.rs @@ -0,0 +1,39 @@ +use alloc::vec::Vec; + +/// Error returned when a packet cannot create a new session. +pub enum SessionInitError { + /// The packet was intentionally consumed without creating state or sending output. + Rejected, + + /// The packet was rejected with response data that should be sent on the same hook. + Response { + /// Raw `Packet::data` for the response frame. + data: Vec, + + /// Whether the response should close the hook after successful routing. + end_hook: bool, + }, +} + +impl SessionInitError { + /// Creates a silent session rejection. + pub fn rejected() -> Self { + Self::Rejected + } + + /// Creates a non-final response for a rejected session open. + pub fn response(data: Vec) -> Self { + Self::Response { + data, + end_hook: false, + } + } + + /// Creates a final response for a rejected session open. + pub fn response_final(data: Vec) -> Self { + Self::Response { + data, + end_hook: true, + } + } +} diff --git a/src/protocol/session/mod.rs b/src/protocol/session/mod.rs new file mode 100644 index 0000000..90845c8 --- /dev/null +++ b/src/protocol/session/mod.rs @@ -0,0 +1,11 @@ +//! Hook-backed session contracts and generated session storage. + +mod contract; +mod error; +mod status; +mod storage; + +pub use contract::Session; +pub use error::SessionInitError; +pub use status::SessionStatus; +pub use storage::{SessionEntry, SessionFamily}; diff --git a/src/protocol/session/status.rs b/src/protocol/session/status.rs new file mode 100644 index 0000000..c1b21f7 --- /dev/null +++ b/src/protocol/session/status.rs @@ -0,0 +1,15 @@ +/// Session lifecycle status returned from [`Session::update`](super::Session::update). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionStatus { + /// The session is active and should receive future update ticks. + Running, + + /// The session is winding down but still needs future update ticks. + Closing, + + /// The session has finished application work. + /// + /// The generated leaf removes the entry after the update tick. Final packets are + /// routed immediately by the session before returning this status. + Closed, +} diff --git a/src/protocol/session/storage.rs b/src/protocol/session/storage.rs new file mode 100644 index 0000000..19bdc5a --- /dev/null +++ b/src/protocol/session/storage.rs @@ -0,0 +1,70 @@ +use alloc::vec::Vec; + +use crate::protocol::{HookID, PacketQueue}; + +/// Storage entry used by macro-generated session stores. +/// +/// The fields are public so generated code in downstream crates can keep the update +/// loop straightforward and static. Handwritten leaves may also use this type, but it +/// is intentionally small rather than a full session framework. +pub struct SessionEntry { + /// Hook id associated with this live session. + pub hook_id: HookID, + + /// Application-owned session state. + pub state: S, + + /// Packets delivered for this hook but not yet consumed by the session. + pub inbox: PacketQueue, + + /// Whether application logic has finished and should be removed after update. + pub closed: bool, +} + +/// Generated storage for one session family. +/// +/// The macro only names this field and picks the concrete `Session` type. All update, +/// retry, and cleanup behavior lives in normal Rust helpers so the template stays +/// small and readable. +pub struct SessionFamily { + /// Active hook-backed sessions for this family. + pub entries: Vec>, +} + +impl SessionFamily { + /// Creates an empty session family. + pub fn new() -> Self { + Self { + entries: Vec::new(), + } + } + + /// Counts packets retained by this family for retry or future session work. + pub fn pending_packet_count(&self) -> usize { + let mut count = 0usize; + + for entry in &self.entries { + count += entry.inbox.len(); + } + + count + } +} + +impl Default for SessionFamily { + fn default() -> Self { + Self::new() + } +} + +impl SessionEntry { + /// Creates one active session entry for `hook_id`. + pub fn new(hook_id: HookID, state: S) -> Self { + Self { + hook_id, + state, + inbox: PacketQueue::new(), + closed: false, + } + } +} From 2d5f04a024f054eea65a8bf1c8a2b0eddfe065d6 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:57:56 -0600 Subject: [PATCH 11/12] Reorganize protocol test modules --- src/protocol/mod.rs | 6 +- .../tests/endpoint/downward_routing.rs | 94 ++++ src/protocol/tests/endpoint/hook_lifecycle.rs | 48 ++ src/protocol/tests/endpoint/local_delivery.rs | 70 +++ .../tests/endpoint/malformed_or_forged.rs | 112 ++++ src/protocol/tests/endpoint/mod.rs | 5 + src/protocol/tests/endpoint/upward_routing.rs | 143 +++++ src/protocol/tests/integration/mod.rs | 2 + src/protocol/tests/integration/oneshot.rs | 149 ++++++ .../tests/{oneshot => integration}/streams.rs | 17 +- .../{leaves.rs => leaves/caller.rs} | 221 +------- src/protocol/tests/merkle_sync/leaves/mod.rs | 7 + .../tests/merkle_sync/leaves/respondent.rs | 147 ++++++ .../tests/merkle_sync/leaves/transport.rs | 79 +++ src/protocol/tests/mod.rs | 5 + src/protocol/tests/oneshot/mod.rs | 491 ------------------ src/protocol/tests/oneshot/support.rs | 205 -------- src/protocol/tests/support/assertions.rs | 23 + src/protocol/tests/support/endpoints.rs | 42 ++ src/protocol/tests/support/mod.rs | 4 + src/protocol/tests/support/packets.rs | 23 + src/protocol/tests/support/transport.rs | 63 +++ unshell-leaves/leaf-pty/src/tests/session.rs | 282 ---------- .../leaf-pty/src/tests/session/concurrency.rs | 47 ++ .../leaf-pty/src/tests/session/failure.rs | 41 ++ .../leaf-pty/src/tests/session/filtering.rs | 44 ++ .../src/tests/session/input_output.rs | 33 ++ .../leaf-pty/src/tests/session/lifecycle.rs | 145 ++++++ .../leaf-pty/src/tests/session/mod.rs | 5 + unshell-leaves/leaf-pty/src/tests/support.rs | 137 ----- .../leaf-pty/src/tests/support/assertions.rs | 47 ++ .../leaf-pty/src/tests/support/drains.rs | 23 + .../leaf-pty/src/tests/support/endpoints.rs | 40 ++ .../leaf-pty/src/tests/support/mod.rs | 9 + .../leaf-pty/src/tests/support/packets.rs | 46 ++ 35 files changed, 1519 insertions(+), 1336 deletions(-) create mode 100644 src/protocol/tests/endpoint/downward_routing.rs create mode 100644 src/protocol/tests/endpoint/hook_lifecycle.rs create mode 100644 src/protocol/tests/endpoint/local_delivery.rs create mode 100644 src/protocol/tests/endpoint/malformed_or_forged.rs create mode 100644 src/protocol/tests/endpoint/mod.rs create mode 100644 src/protocol/tests/endpoint/upward_routing.rs create mode 100644 src/protocol/tests/integration/mod.rs create mode 100644 src/protocol/tests/integration/oneshot.rs rename src/protocol/tests/{oneshot => integration}/streams.rs (96%) rename src/protocol/tests/merkle_sync/{leaves.rs => leaves/caller.rs} (53%) create mode 100644 src/protocol/tests/merkle_sync/leaves/mod.rs create mode 100644 src/protocol/tests/merkle_sync/leaves/respondent.rs create mode 100644 src/protocol/tests/merkle_sync/leaves/transport.rs create mode 100644 src/protocol/tests/mod.rs delete mode 100644 src/protocol/tests/oneshot/mod.rs delete mode 100644 src/protocol/tests/oneshot/support.rs create mode 100644 src/protocol/tests/support/assertions.rs create mode 100644 src/protocol/tests/support/endpoints.rs create mode 100644 src/protocol/tests/support/mod.rs create mode 100644 src/protocol/tests/support/packets.rs create mode 100644 src/protocol/tests/support/transport.rs delete mode 100644 unshell-leaves/leaf-pty/src/tests/session.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/session/concurrency.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/session/failure.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/session/filtering.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/session/input_output.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/session/lifecycle.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/session/mod.rs delete mode 100644 unshell-leaves/leaf-pty/src/tests/support.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/support/assertions.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/support/drains.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/support/endpoints.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/support/mod.rs create mode 100644 unshell-leaves/leaf-pty/src/tests/support/packets.rs diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 3c030a8..edc92c7 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -32,8 +32,4 @@ pub type PacketQueue = VecDeque; type RouteMap = Vec<(EndpointName, PacketQueue)>; #[cfg(test)] -mod tests { - mod merkle_sync; - mod oneshot; - mod packet; -} +mod tests; diff --git a/src/protocol/tests/endpoint/downward_routing.rs b/src/protocol/tests/endpoint/downward_routing.rs new file mode 100644 index 0000000..aaabb16 --- /dev/null +++ b/src/protocol/tests/endpoint/downward_routing.rs @@ -0,0 +1,94 @@ +use alloc::vec; + +use crate::protocol::{Endpoint, EndpointError, RouteDirection}; + +use super::super::support::{ + assertions::{assert_hook_present, assert_hook_removed}, + endpoints::{ENDPOINT_A, ENDPOINT_B, ENDPOINT_C, endpoint_at, single_outbound_packet}, + packets::{echo_packet, echo_packet_with_end}, +}; + +#[test] +fn inbound_downward_packet_routes_to_immediate_child() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); + + endpoint + .add_inbound_from( + ENDPOINT_A, + echo_packet(vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C], hook_id), + ) + .unwrap(); + + let packet = single_outbound_packet(&endpoint, ENDPOINT_C); + assert!(!packet.end_hook); + assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); + assert_hook_present(&endpoint, hook_id); + assert_eq!(endpoint.hook_peer(hook_id), Some(ENDPOINT_C)); + assert!(!Endpoint::route_contains(ENDPOINT_A, &endpoint.outbound)); +} + +#[test] +fn outbound_downward_packet_routes_to_immediate_child() { + let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_B); + endpoint.add_connection(ENDPOINT_B, false); + + endpoint + .add_outbound(echo_packet_with_end( + vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C], + hook_id, + true, + )) + .unwrap(); + + let packet = single_outbound_packet(&endpoint, ENDPOINT_B); + assert!(packet.end_hook); + assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); + assert_hook_removed(&endpoint, hook_id); + assert!(!Endpoint::route_contains(ENDPOINT_C, &endpoint.outbound)); +} + +#[test] +fn downward_outbound_without_hook_is_allowed() { + let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); + endpoint.add_connection(ENDPOINT_B, false); + + let new_hook = endpoint.get_hook_id(); + + endpoint + .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], new_hook)) + .unwrap(); + + assert_eq!( + Endpoint::route_get(ENDPOINT_B, &endpoint.outbound) + .unwrap() + .len(), + 1 + ); + assert_hook_present(&endpoint, new_hook); + assert_eq!(endpoint.hook_peer(new_hook), Some(ENDPOINT_B)); +} + +#[test] +fn downward_route_without_connection_is_rejected() { + let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); + let hook_id = endpoint.get_hook_id(); + + let error = endpoint + .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) + .unwrap_err(); + + assert!(matches!( + error, + EndpointError::MissingConnection { + next_hop: ENDPOINT_B, + direction: RouteDirection::Downward, + } + )); + assert_hook_removed(&endpoint, hook_id); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} diff --git a/src/protocol/tests/endpoint/hook_lifecycle.rs b/src/protocol/tests/endpoint/hook_lifecycle.rs new file mode 100644 index 0000000..ff4dad3 --- /dev/null +++ b/src/protocol/tests/endpoint/hook_lifecycle.rs @@ -0,0 +1,48 @@ +use alloc::vec; + +use crate::protocol::{Endpoint, EndpointError, RouteDirection}; + +use super::super::support::{ + assertions::{assert_hook_present, assert_hook_removed}, + endpoints::{ENDPOINT_A, ENDPOINT_B, endpoint_at, single_outbound_packet}, + packets::echo_packet_with_end, +}; + +#[test] +fn end_hook_removes_hook_after_packet_is_queued() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_A); + endpoint.add_connection(ENDPOINT_A, true); + + endpoint + .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) + .unwrap(); + + assert_hook_removed(&endpoint, hook_id); + assert_eq!( + single_outbound_packet(&endpoint, ENDPOINT_A).hook_id, + hook_id + ); +} + +#[test] +fn failed_end_hook_route_keeps_hook_state() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_A); + + let error = endpoint + .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) + .unwrap_err(); + + assert!(matches!( + error, + EndpointError::MissingConnection { + next_hop: ENDPOINT_A, + direction: RouteDirection::Upward, + } + )); + assert_hook_present(&endpoint, hook_id); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} diff --git a/src/protocol/tests/endpoint/local_delivery.rs b/src/protocol/tests/endpoint/local_delivery.rs new file mode 100644 index 0000000..d685065 --- /dev/null +++ b/src/protocol/tests/endpoint/local_delivery.rs @@ -0,0 +1,70 @@ +use alloc::vec; + +use crate::protocol::{Endpoint, EndpointError}; + +use super::super::support::{ + assertions::{assert_hook_present, assert_hook_removed}, + endpoints::{ENDPOINT_A, ENDPOINT_B, endpoint_at, single_inbound_packet}, + packets::echo_packet, +}; + +#[test] +fn inbound_downward_packet_for_local_endpoint_opens_hook() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.add_connection(ENDPOINT_A, true); + + endpoint + .add_inbound_from( + ENDPOINT_A, + echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id), + ) + .unwrap(); + + let packet = single_inbound_packet(&endpoint, ENDPOINT_B); + assert!(!packet.end_hook); + assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B]); + assert_hook_present(&endpoint, hook_id); + assert_eq!(endpoint.hook_peer(hook_id), Some(ENDPOINT_A)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn outbound_packet_for_local_endpoint_is_delivered_locally() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + + endpoint + .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) + .unwrap(); + + let packet = single_inbound_packet(&endpoint, ENDPOINT_B); + assert!(!packet.end_hook); + assert_eq!(packet.data, "ABC123".as_bytes()); + assert_hook_removed(&endpoint, hook_id); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn inbound_without_absolute_path_is_rejected() { + let mut endpoint = Endpoint::new(ENDPOINT_A); + + let error = endpoint + .add_inbound(echo_packet(vec![ENDPOINT_A], 1)) + .unwrap_err(); + + assert!(matches!(error, EndpointError::EndpointPathUnset)); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); +} + +#[test] +fn outbound_without_absolute_path_is_rejected() { + let mut endpoint = Endpoint::new(ENDPOINT_A); + + let error = endpoint + .add_outbound(echo_packet(vec![ENDPOINT_A], 1)) + .unwrap_err(); + + assert!(matches!(error, EndpointError::EndpointPathUnset)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} diff --git a/src/protocol/tests/endpoint/malformed_or_forged.rs b/src/protocol/tests/endpoint/malformed_or_forged.rs new file mode 100644 index 0000000..ea6ca44 --- /dev/null +++ b/src/protocol/tests/endpoint/malformed_or_forged.rs @@ -0,0 +1,112 @@ +use alloc::vec; + +use crate::protocol::{Endpoint, EndpointError, Leaf}; + +use super::super::support::{ + assertions::assert_hook_present, + endpoints::{ENDPOINT_A, ENDPOINT_B, ENDPOINT_C, endpoint_at, single_inbound_packet}, + packets::{echo_packet, echo_packet_with_end}, + transport::CommsLeaf, +}; + +#[test] +fn forged_sideways_packet_is_rejected_as_incorrect_path() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_A); + endpoint.add_connection(ENDPOINT_A, true); + + let error = endpoint + .add_inbound_from( + ENDPOINT_A, + echo_packet(vec![ENDPOINT_A, ENDPOINT_C], hook_id), + ) + .unwrap_err(); + + assert!(matches!(error, EndpointError::DestinationOutsideLocalTree)); + assert_hook_present(&endpoint, hook_id); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn malformed_frame_is_dropped_by_comms_leaf() { + let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); + let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); + let mut endpoint = Endpoint::new(ENDPOINT_B); + let mut comms = CommsLeaf { + tx: tx_unused, + rx: rx_for_endpoint, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }; + endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; + + tx_to_endpoint.send(vec![0, 1, 2, 3]).unwrap(); + comms.update(&mut endpoint); + + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn malformed_frame_does_not_block_following_valid_packet() { + let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); + let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); + let hook_id = 42; + let mut endpoint = Endpoint::new(ENDPOINT_B); + let mut comms = CommsLeaf { + tx: tx_unused, + rx: rx_for_endpoint, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }; + endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; + + tx_to_endpoint.send(vec![0, 1, 2, 3]).unwrap(); + tx_to_endpoint + .send( + echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id) + .serialize() + .unwrap(), + ) + .unwrap(); + comms.update(&mut endpoint); + + let packet = single_inbound_packet(&endpoint, ENDPOINT_B); + assert!(!packet.end_hook); + assert_eq!(packet.hook_id, hook_id); + assert_hook_present(&endpoint, hook_id); +} + +#[test] +fn forged_frame_without_required_hook_is_dropped_by_comms_leaf() { + let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); + let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); + let mut endpoint = Endpoint::new(ENDPOINT_B); + let mut comms = CommsLeaf { + tx: tx_unused, + rx: rx_for_endpoint, + remote_id: ENDPOINT_C, + is_authority: false, + started: false, + }; + endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; + endpoint.accept_hook(7, ENDPOINT_C); + endpoint.add_connection(ENDPOINT_A, true); + + tx_to_endpoint + .send( + echo_packet_with_end(vec![ENDPOINT_A], 12, true) + .serialize() + .unwrap(), + ) + .unwrap(); + comms.update(&mut endpoint); + + assert_hook_present(&endpoint, 7); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} diff --git a/src/protocol/tests/endpoint/mod.rs b/src/protocol/tests/endpoint/mod.rs new file mode 100644 index 0000000..f43c2a8 --- /dev/null +++ b/src/protocol/tests/endpoint/mod.rs @@ -0,0 +1,5 @@ +mod downward_routing; +mod hook_lifecycle; +mod local_delivery; +mod malformed_or_forged; +mod upward_routing; diff --git a/src/protocol/tests/endpoint/upward_routing.rs b/src/protocol/tests/endpoint/upward_routing.rs new file mode 100644 index 0000000..8da8f22 --- /dev/null +++ b/src/protocol/tests/endpoint/upward_routing.rs @@ -0,0 +1,143 @@ +use alloc::vec; + +use crate::protocol::{Endpoint, EndpointError, RouteDirection}; + +use super::super::support::{ + assertions::{assert_hook_present, assert_hook_removed}, + endpoints::{ENDPOINT_A, ENDPOINT_B, ENDPOINT_C, endpoint_at, single_outbound_packet}, + packets::echo_packet_with_end, +}; + +#[test] +fn inbound_upward_packet_with_hook_routes_to_parent() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_C); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); + + endpoint + .add_inbound_from( + ENDPOINT_C, + echo_packet_with_end(vec![ENDPOINT_A], hook_id, true), + ) + .unwrap(); + + let packet = single_outbound_packet(&endpoint, ENDPOINT_A); + assert!(packet.end_hook); + assert_eq!(packet.hook_id, hook_id); + assert_hook_removed(&endpoint, hook_id); + assert!(!Endpoint::route_contains(ENDPOINT_C, &endpoint.outbound)); +} + +#[test] +fn inbound_upward_packet_without_hook_is_rejected() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); + + let error = endpoint + .add_inbound_from( + ENDPOINT_C, + echo_packet_with_end(vec![ENDPOINT_A], hook_id, true), + ) + .unwrap_err(); + + assert!(matches!( + error, + EndpointError::UnknownHook { hook_id: observed_hook_id } if observed_hook_id == hook_id + )); + assert!(Endpoint::routes_is_empty(&endpoint.inbound)); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn forged_upward_packet_with_unknown_hook_is_rejected() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + endpoint.accept_hook(7, ENDPOINT_C); + endpoint.add_connection(ENDPOINT_A, true); + endpoint.add_connection(ENDPOINT_C, false); + + let error = endpoint + .add_inbound_from(ENDPOINT_C, echo_packet_with_end(vec![ENDPOINT_A], 99, true)) + .unwrap_err(); + + assert!(matches!(error, EndpointError::UnknownHook { hook_id: 99 })); + assert_hook_present(&endpoint, 7); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn upward_outbound_without_hook_is_rejected() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + endpoint.accept_hook(7, ENDPOINT_A); + endpoint.add_connection(ENDPOINT_A, true); + + let new_hook = endpoint.get_hook_id(); + + let error = endpoint + .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], new_hook, true)) + .unwrap_err(); + + assert!(matches!( + error, + EndpointError::UnknownHook { hook_id: observed_hook_id } if observed_hook_id == new_hook + )); + assert_hook_present(&endpoint, 7); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn deeper_upward_route_uses_parent_as_next_hop() { + let mut endpoint = endpoint_at(ENDPOINT_C, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); + let new_hook = endpoint.get_hook_id(); + + endpoint.accept_hook(new_hook, ENDPOINT_B); + endpoint.add_connection(ENDPOINT_B, true); + + endpoint + .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], new_hook, true)) + .unwrap(); + + assert!(Endpoint::route_contains(ENDPOINT_B, &endpoint.outbound)); + assert!(!Endpoint::route_contains(ENDPOINT_A, &endpoint.outbound)); + assert_hook_removed(&endpoint, new_hook); +} + +#[test] +fn upward_route_without_connection_is_rejected_even_with_hook() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_A); + + let error = endpoint + .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) + .unwrap_err(); + + assert!(matches!( + error, + EndpointError::MissingConnection { + next_hop: ENDPOINT_A, + direction: RouteDirection::Upward, + } + )); + assert_hook_present(&endpoint, hook_id); + assert!(Endpoint::routes_is_empty(&endpoint.outbound)); +} + +#[test] +fn trusted_upward_packet_without_peer_metadata_checks_hook_existence_only() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let hook_id = endpoint.get_hook_id(); + endpoint.accept_hook(hook_id, ENDPOINT_A); + endpoint.add_connection(ENDPOINT_A, true); + + endpoint + .add_inbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) + .unwrap(); + + let packet = single_outbound_packet(&endpoint, ENDPOINT_A); + assert_eq!(packet.hook_id, hook_id); + assert_hook_removed(&endpoint, hook_id); +} diff --git a/src/protocol/tests/integration/mod.rs b/src/protocol/tests/integration/mod.rs new file mode 100644 index 0000000..c72bf5f --- /dev/null +++ b/src/protocol/tests/integration/mod.rs @@ -0,0 +1,2 @@ +mod oneshot; +mod streams; diff --git a/src/protocol/tests/integration/oneshot.rs b/src/protocol/tests/integration/oneshot.rs new file mode 100644 index 0000000..e356aa9 --- /dev/null +++ b/src/protocol/tests/integration/oneshot.rs @@ -0,0 +1,149 @@ +use alloc::{vec, vec::Vec}; + +use crate::protocol::{Endpoint, Leaf}; + +#[cfg(feature = "interface")] +use crate::protocol::LeafMeta; + +use super::super::support::{ + endpoints::{ENDPOINT_A, ENDPOINT_B}, + packets::{echo_packet, echo_packet_with_end}, + transport::CommsLeaf, +}; + +const LEAF_CONTROLLER: u32 = 100; +const LEAF_RESPONDER: u32 = 102; + +struct ControllerLeaf { + has_run: bool, +} + +struct ResponderLeaf; + +impl Leaf for ControllerLeaf { + fn get_id(&self) -> u32 { + LEAF_CONTROLLER + } + + #[cfg(feature = "interface")] + fn get_meta(&self) -> LeafMeta { + LeafMeta { + name: "Controller Leaf", + identifier: "dev.unshell.test.controller_leaf", + version: "v0", + authors: alloc::vec!["ASTATIN3"], + } + } + + fn update(&mut self, endpoint: &mut Endpoint) { + if !self.has_run { + // The controller starts exactly one request so the end-to-end test can + // assert deterministic routing without accumulating retries. + let hook_id = endpoint.get_hook_id(); + let packet = echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id); + let _ = endpoint.add_outbound(packet); + self.has_run = true; + } + } +} + +impl Leaf for ResponderLeaf { + fn get_id(&self) -> u32 { + LEAF_RESPONDER + } + + #[cfg(feature = "interface")] + fn get_meta(&self) -> LeafMeta { + LeafMeta { + name: "Responder Leaf", + identifier: "dev.unshell.test.responder_leaf", + version: "v0", + authors: alloc::vec!["ASTATIN3"], + } + } + + fn update(&mut self, endpoint: &mut Endpoint) { + let local_id = endpoint.path.last().cloned().unwrap_or(0); + let mut packets = Vec::new(); + + endpoint.take_inbound_clear(local_id, |packet| { + let mut response = echo_packet_with_end(vec![ENDPOINT_A], packet.hook_id, true); + response.hook_id = packet.hook_id; + response.data = packet.data.clone(); + packets.push(response); + }); + + for packet in packets { + let _ = endpoint.add_outbound(packet); + } + } +} + +#[test] +fn request_response_round_trip_over_mock_transport() { + let (tx_a, rx_a) = crossbeam_channel::unbounded(); + let (tx_b, rx_b) = crossbeam_channel::unbounded(); + + let mut endpoint_a = Endpoint::new(ENDPOINT_A); + let mut controller_a = ControllerLeaf { has_run: false }; + let mut comms_a = CommsLeaf { + tx: tx_b, + rx: rx_a, + remote_id: ENDPOINT_B, + is_authority: false, + started: false, + }; + endpoint_a.path = vec![ENDPOINT_A]; + + let mut endpoint_b = Endpoint::new(ENDPOINT_B); + let mut responder_b = ResponderLeaf; + let mut comms_b = CommsLeaf { + tx: tx_a, + rx: rx_b, + remote_id: ENDPOINT_A, + is_authority: true, + started: false, + }; + endpoint_b.path = vec![ENDPOINT_A, ENDPOINT_B]; + + // Connections are registered routing state. The comms leaves also insert them + // during updates, but the first application packet should not depend on leaf order. + endpoint_a.add_connection(ENDPOINT_B, false); + endpoint_b.add_connection(ENDPOINT_A, true); + + // Cycle 1: A sends request to B. + controller_a.update(&mut endpoint_a); + comms_a.update(&mut endpoint_a); + responder_b.update(&mut endpoint_b); + comms_b.update(&mut endpoint_b); + + // Cycle 2: B receives request and sends response to A. + responder_b.update(&mut endpoint_b); + comms_b.update(&mut endpoint_b); + controller_a.update(&mut endpoint_a); + comms_a.update(&mut endpoint_a); + + // Cycle 3: A's transport leaf needs one more update to pull the response bytes + // from the channel and put the packet into the inbound queue. + controller_a.update(&mut endpoint_a); + comms_a.update(&mut endpoint_a); + + assert!( + Endpoint::route_contains(ENDPOINT_A, &endpoint_a.inbound), + "Endpoint A should have received response" + ); + assert_eq!( + Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) + .unwrap() + .len(), + 1, + "Endpoint A should have exactly one packet" + ); + let response = &Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) + .unwrap() + .front() + .unwrap(); + assert!(response.end_hook); + assert_eq!(response.data, "ABC123".as_bytes()); + assert_eq!(endpoint_b.hook_count(), 0); +} diff --git a/src/protocol/tests/oneshot/streams.rs b/src/protocol/tests/integration/streams.rs similarity index 96% rename from src/protocol/tests/oneshot/streams.rs rename to src/protocol/tests/integration/streams.rs index 11f4940..84d12c1 100644 --- a/src/protocol/tests/oneshot/streams.rs +++ b/src/protocol/tests/integration/streams.rs @@ -5,10 +5,15 @@ use crate::protocol::LeafMeta; use alloc::{format, vec, vec::Vec}; -use super::support::{CommsLeaf, ENDPOINT_A, ENDPOINT_B, assert_hook_present, assert_hook_removed}; +use super::super::support::{ + assertions::{assert_hook_present, assert_hook_removed}, + endpoints::{ENDPOINT_A, ENDPOINT_B}, + transport::CommsLeaf, +}; const LEAF_STREAM_CALLER: u32 = 200; const LEAF_STREAM_RESPONDENT: u32 = 201; +const LEAF_COMMS: u32 = 101; /// Builds the initial downwards packet that opens the stream on the respondent. /// @@ -104,7 +109,7 @@ impl Leaf for StreamCallerLeaf { name: "Stream Caller Leaf", identifier: "dev.unshell.test.stream_caller_leaf", version: "v0", - authors: vec!["ASTATIN3"], + authors: alloc::vec!["ASTATIN3"], } } @@ -127,10 +132,10 @@ impl Leaf for StreamRespondentLeaf { #[cfg(feature = "interface")] fn get_meta(&self) -> LeafMeta { LeafMeta { - name: "Stream Respondant Leaf", + name: "Stream Respondent Leaf", identifier: "dev.unshell.test.stream_respondent_leaf", version: "v0", - authors: vec!["ASTATIN3"], + authors: alloc::vec!["ASTATIN3"], } } @@ -243,9 +248,9 @@ fn stream_endpoints(total_packets: usize) -> StreamHarness { /// Asserts the requested two-endpoint, four-leaf topology. fn assert_four_leaf_topology(harness: &StreamHarness) { assert_eq!(harness.caller_a.get_id(), LEAF_STREAM_CALLER); - assert_eq!(harness.comms_a.get_id(), 101); + assert_eq!(harness.comms_a.get_id(), LEAF_COMMS); assert_eq!(harness.respondent_b.get_id(), LEAF_STREAM_RESPONDENT); - assert_eq!(harness.comms_b.get_id(), 101); + assert_eq!(harness.comms_b.get_id(), LEAF_COMMS); } /// Drives the initial request until it is queued locally on endpoint B. diff --git a/src/protocol/tests/merkle_sync/leaves.rs b/src/protocol/tests/merkle_sync/leaves/caller.rs similarity index 53% rename from src/protocol/tests/merkle_sync/leaves.rs rename to src/protocol/tests/merkle_sync/leaves/caller.rs index caf9a73..751d59b 100644 --- a/src/protocol/tests/merkle_sync/leaves.rs +++ b/src/protocol/tests/merkle_sync/leaves/caller.rs @@ -1,43 +1,25 @@ -use alloc::{collections::VecDeque, rc::Rc, vec, vec::Vec}; +use alloc::{collections::VecDeque, rc::Rc, vec::Vec}; use core::cell::RefCell; -use crossbeam_channel::{Receiver, Sender}; - use crate::protocol::{Endpoint, Leaf, Packet}; #[cfg(feature = "interface")] use crate::protocol::LeafMeta; -use super::{ +use super::super::{ codec::{decode_block_chunk, decode_child_summary, decode_u32}, constants::{ - ENDPOINT_CALLER, ENDPOINT_RESPONDENT, LEAF_MERKLE_CALLER, LEAF_MERKLE_RESPONDENT, - LEAF_MOCK_CONNECTION, PROC_BLOCK_CHUNK, PROC_CHILD_HASH_ENTRY, PROC_GET_BLOCK_STREAM, - PROC_GET_CHILD_HASHES, PROC_GET_ROOT_HASH, PROC_ROOT_HASH, ROOT_NODE, + ENDPOINT_CALLER, LEAF_MERKLE_CALLER, PROC_BLOCK_CHUNK, PROC_CHILD_HASH_ENTRY, + PROC_GET_BLOCK_STREAM, PROC_GET_CHILD_HASHES, PROC_GET_ROOT_HASH, PROC_ROOT_HASH, + ROOT_NODE, }, - rpc::{ - block_chunk_frame, block_stream_request, child_hash_frame, child_hashes_request, - root_hash_frame, root_hash_request, - }, - state::{CallerPhase, CallerReport, RespondentReport, ResponseStream}, - tree::{BlockChunk, ChildKind, MerkleStore}, + rpc::{block_stream_request, child_hashes_request, root_hash_request}, + state::{CallerPhase, CallerReport}, + tree::{ChildKind, MerkleStore}, }; -/// Leaf that simulates a serialized transport connection with crossbeam channels. -/// -/// This is intentionally tiny and reusable. Both endpoints in the Merkle test have -/// exactly one of these leaves, giving the requested four-leaf topology: caller, -/// respondent, and two mock connections. -pub(super) struct MockConnectionLeaf { - pub(super) tx: Sender>, - pub(super) rx: Receiver>, - pub(super) remote_id: u32, - pub(super) is_authority: bool, - pub(super) started: bool, -} - /// Caller leaf that drives the Merkle synchronization algorithm. -pub(super) struct MerkleCallerLeaf { +pub(crate) struct MerkleCallerLeaf { local: MerkleStore, phase: CallerPhase, pending_nodes: VecDeque, @@ -45,34 +27,9 @@ pub(super) struct MerkleCallerLeaf { report: Rc>, } -/// Respondent leaf that serves Merkle hash and block streams. -pub(super) struct MerkleRespondentLeaf { - remote: MerkleStore, - active_stream: Option, - report: Rc>, -} - -impl MockConnectionLeaf { - /// Creates one side of a mock connection. - pub(super) fn new( - tx: Sender>, - rx: Receiver>, - remote_id: u32, - is_authority: bool, - ) -> Self { - Self { - tx, - rx, - remote_id, - is_authority, - started: false, - } - } -} - impl MerkleCallerLeaf { /// Creates a caller with a local store and externally visible report. - pub(super) fn new(local: MerkleStore, report: Rc>) -> Self { + pub(crate) fn new(local: MerkleStore, report: Rc>) -> Self { Self { local, phase: CallerPhase::NeedRoot, @@ -83,55 +40,6 @@ impl MerkleCallerLeaf { } } -impl MerkleRespondentLeaf { - /// Creates a respondent backed by the authoritative remote store. - pub(super) fn new(remote: MerkleStore, report: Rc>) -> Self { - Self { - remote, - active_stream: None, - report, - } - } -} - -impl Leaf for MockConnectionLeaf { - fn get_id(&self) -> u32 { - LEAF_MOCK_CONNECTION - } - - #[cfg(feature = "interface")] - fn get_meta(&self) -> LeafMeta { - LeafMeta { - name: "Merke Connection Leaf", - identifier: "dev.unshell.test.merkle.connection", - version: "v0", - authors: vec!["ASTATIN3"], - } - } - - fn update(&mut self, endpoint: &mut Endpoint) { - if !self.started { - endpoint.add_connection(self.remote_id, self.is_authority); - self.started = true; - } - - while !self.rx.is_empty() { - let data = self.rx.recv().unwrap(); - - // Mock transports move untrusted bytes. Malformed frames are dropped so - // the sync state machine is tested only after packet parsing succeeds. - if let Ok(packet) = Packet::deserialize(&data) { - let _ = endpoint.add_inbound_from(self.remote_id, packet); - } - } - - endpoint.take_outbound_clear(self.remote_id, |packet| { - let data = packet.serialize().unwrap(); - let _ = self.tx.send(data); - }); - } -} - impl Leaf for MerkleCallerLeaf { fn get_id(&self) -> u32 { LEAF_MERKLE_CALLER @@ -143,7 +51,7 @@ impl Leaf for MerkleCallerLeaf { name: "Merke Caller Leaf", identifier: "dev.unshell.test.merkle.caller", version: "v0", - authors: vec!["ASTATIN3"], + authors: alloc::vec!["ASTATIN3"], } } @@ -153,27 +61,6 @@ impl Leaf for MerkleCallerLeaf { } } -impl Leaf for MerkleRespondentLeaf { - fn get_id(&self) -> u32 { - LEAF_MERKLE_RESPONDENT - } - - #[cfg(feature = "interface")] - fn get_meta(&self) -> LeafMeta { - LeafMeta { - name: "Merke Respondent Leaf", - identifier: "dev.unshell.test.merkle.respondent", - version: "v0", - authors: vec!["ASTATIN3"], - } - } - - fn update(&mut self, endpoint: &mut Endpoint) { - self.open_stream_from_request(endpoint); - self.send_one_response_frame(endpoint); - } -} - impl MerkleCallerLeaf { /// Consumes all response packets currently delivered to endpoint A. fn receive_responses(&mut self, endpoint: &mut Endpoint) { @@ -346,89 +233,3 @@ impl MerkleCallerLeaf { report.final_root_hash = Some(self.local.root_hash()); } } - -impl MerkleRespondentLeaf { - /// Opens one response stream from the first pending local request. - fn open_stream_from_request(&mut self, endpoint: &mut Endpoint) { - if self.active_stream.is_some() { - return; - } - - let mut request = None; - endpoint.take_inbound_clear(ENDPOINT_RESPONDENT, |packet| { - if request.is_none() { - request = Some((packet.hook_id, packet.procedure_id, packet.data.clone())); - } - }); - - let Some((hook_id, procedure_id, data)) = request else { - return; - }; - - let frames = self.frames_for_request(procedure_id, &data); - - self.report.borrow_mut().requests_seen.push(procedure_id); - if !frames.is_empty() { - self.report.borrow_mut().streams_started += 1; - self.active_stream = Some(ResponseStream::new(hook_id, frames)); - } - } - - /// Builds response frames for one request procedure. - fn frames_for_request(&self, procedure_id: u32, data: &[u8]) -> Vec { - match procedure_id { - PROC_GET_ROOT_HASH => vec![root_hash_frame(self.remote.root_hash())], - PROC_GET_CHILD_HASHES => { - let node_id = decode_u32(data).expect("child hash request node id"); - self.remote - .child_summaries(node_id) - .into_iter() - .map(child_hash_frame) - .collect() - } - PROC_GET_BLOCK_STREAM => { - let block_id = decode_u32(data).expect("block stream request block id"); - let chunks = self.remote.block_chunks(block_id); - let total = chunks.len() as u32; - chunks - .into_iter() - .enumerate() - .map(|(index, data)| { - block_chunk_frame(BlockChunk { - block_id, - index: index as u32, - total, - data, - }) - }) - .collect() - } - _ => Vec::new(), - } - } - - /// Sends at most one response frame per update loop. - fn send_one_response_frame(&mut self, endpoint: &mut Endpoint) { - let Some(stream) = self.active_stream.as_mut() else { - return; - }; - - if stream.is_empty() { - self.active_stream = None; - return; - } - - let packet = stream.next_packet().expect("active stream frame"); - if endpoint.add_outbound(packet).is_err() { - return; - } - - self.report.borrow_mut().frames_sent += 1; - stream.advance(); - - if stream.is_complete() { - self.report.borrow_mut().streams_completed += 1; - self.active_stream = None; - } - } -} diff --git a/src/protocol/tests/merkle_sync/leaves/mod.rs b/src/protocol/tests/merkle_sync/leaves/mod.rs new file mode 100644 index 0000000..2e1ced9 --- /dev/null +++ b/src/protocol/tests/merkle_sync/leaves/mod.rs @@ -0,0 +1,7 @@ +mod caller; +mod respondent; +mod transport; + +pub(crate) use caller::MerkleCallerLeaf; +pub(crate) use respondent::MerkleRespondentLeaf; +pub(crate) use transport::MockConnectionLeaf; diff --git a/src/protocol/tests/merkle_sync/leaves/respondent.rs b/src/protocol/tests/merkle_sync/leaves/respondent.rs new file mode 100644 index 0000000..1998054 --- /dev/null +++ b/src/protocol/tests/merkle_sync/leaves/respondent.rs @@ -0,0 +1,147 @@ +use alloc::{rc::Rc, vec::Vec}; +use core::cell::RefCell; + +use crate::protocol::{Endpoint, Leaf}; + +#[cfg(feature = "interface")] +use crate::protocol::LeafMeta; + +use super::super::{ + codec::decode_u32, + constants::{ + ENDPOINT_RESPONDENT, LEAF_MERKLE_RESPONDENT, PROC_GET_BLOCK_STREAM, PROC_GET_CHILD_HASHES, + PROC_GET_ROOT_HASH, + }, + rpc::{block_chunk_frame, child_hash_frame, root_hash_frame}, + state::{RespondentReport, ResponseStream}, + tree::{BlockChunk, MerkleStore}, +}; + +/// Respondent leaf that serves Merkle hash and block streams. +pub(crate) struct MerkleRespondentLeaf { + remote: MerkleStore, + active_stream: Option, + report: Rc>, +} + +impl MerkleRespondentLeaf { + /// Creates a respondent backed by the authoritative remote store. + pub(crate) fn new(remote: MerkleStore, report: Rc>) -> Self { + Self { + remote, + active_stream: None, + report, + } + } +} + +impl Leaf for MerkleRespondentLeaf { + fn get_id(&self) -> u32 { + LEAF_MERKLE_RESPONDENT + } + + #[cfg(feature = "interface")] + fn get_meta(&self) -> LeafMeta { + LeafMeta { + name: "Merke Respondent Leaf", + identifier: "dev.unshell.test.merkle.respondent", + version: "v0", + authors: alloc::vec!["ASTATIN3"], + } + } + + fn update(&mut self, endpoint: &mut Endpoint) { + self.open_stream_from_request(endpoint); + self.send_one_response_frame(endpoint); + } +} + +impl MerkleRespondentLeaf { + /// Opens one response stream from the first pending local request. + fn open_stream_from_request(&mut self, endpoint: &mut Endpoint) { + if self.active_stream.is_some() { + return; + } + + let mut request = None; + endpoint.take_inbound_clear(ENDPOINT_RESPONDENT, |packet| { + if request.is_none() { + request = Some((packet.hook_id, packet.procedure_id, packet.data.clone())); + } + }); + + let Some((hook_id, procedure_id, data)) = request else { + return; + }; + + let frames = self.frames_for_request(procedure_id, &data); + + self.report.borrow_mut().requests_seen.push(procedure_id); + if !frames.is_empty() { + self.report.borrow_mut().streams_started += 1; + self.active_stream = Some(ResponseStream::new(hook_id, frames)); + } + } + + /// Builds response frames for one request procedure. + fn frames_for_request( + &self, + procedure_id: u32, + data: &[u8], + ) -> Vec { + match procedure_id { + PROC_GET_ROOT_HASH => alloc::vec![root_hash_frame(self.remote.root_hash())], + PROC_GET_CHILD_HASHES => { + let node_id = decode_u32(data).expect("child hash request node id"); + self.remote + .child_summaries(node_id) + .into_iter() + .map(child_hash_frame) + .collect() + } + PROC_GET_BLOCK_STREAM => { + let block_id = decode_u32(data).expect("block stream request block id"); + let chunks = self.remote.block_chunks(block_id); + let total = chunks.len() as u32; + chunks + .into_iter() + .enumerate() + .map(|(index, data)| { + block_chunk_frame(BlockChunk { + block_id, + index: index as u32, + total, + data, + }) + }) + .collect() + } + _ => Vec::new(), + } + } + + /// Sends at most one response frame per update loop. + fn send_one_response_frame(&mut self, endpoint: &mut Endpoint) { + let Some(stream) = self.active_stream.as_mut() else { + return; + }; + + if stream.is_empty() { + self.active_stream = None; + return; + } + + let packet = stream.next_packet().expect("active stream frame"); + if endpoint.add_outbound(packet).is_err() { + return; + } + + self.report.borrow_mut().frames_sent += 1; + stream.advance(); + + if stream.is_complete() { + self.report.borrow_mut().streams_completed += 1; + self.active_stream = None; + } + } +} diff --git a/src/protocol/tests/merkle_sync/leaves/transport.rs b/src/protocol/tests/merkle_sync/leaves/transport.rs new file mode 100644 index 0000000..05fbae4 --- /dev/null +++ b/src/protocol/tests/merkle_sync/leaves/transport.rs @@ -0,0 +1,79 @@ +use alloc::vec::Vec; + +use crossbeam_channel::{Receiver, Sender}; + +use crate::protocol::{Endpoint, Leaf, Packet}; + +#[cfg(feature = "interface")] +use crate::protocol::LeafMeta; + +use super::super::constants::LEAF_MOCK_CONNECTION; + +/// Leaf that simulates a serialized transport connection with crossbeam channels. +/// +/// This is intentionally tiny and reusable. Both endpoints in the Merkle test have +/// exactly one of these leaves, giving the requested four-leaf topology: caller, +/// respondent, and two mock connections. +pub(crate) struct MockConnectionLeaf { + pub(crate) tx: Sender>, + pub(crate) rx: Receiver>, + pub(crate) remote_id: u32, + pub(crate) is_authority: bool, + pub(crate) started: bool, +} + +impl MockConnectionLeaf { + /// Creates one side of a mock connection. + pub(crate) fn new( + tx: Sender>, + rx: Receiver>, + remote_id: u32, + is_authority: bool, + ) -> Self { + Self { + tx, + rx, + remote_id, + is_authority, + started: false, + } + } +} + +impl Leaf for MockConnectionLeaf { + fn get_id(&self) -> u32 { + LEAF_MOCK_CONNECTION + } + + #[cfg(feature = "interface")] + fn get_meta(&self) -> LeafMeta { + LeafMeta { + name: "Merke Connection Leaf", + identifier: "dev.unshell.test.merkle.connection", + version: "v0", + authors: alloc::vec!["ASTATIN3"], + } + } + + fn update(&mut self, endpoint: &mut Endpoint) { + if !self.started { + endpoint.add_connection(self.remote_id, self.is_authority); + self.started = true; + } + + while !self.rx.is_empty() { + let data = self.rx.recv().unwrap(); + + // Mock transports move untrusted bytes. Malformed frames are dropped so + // the sync state machine is tested only after packet parsing succeeds. + if let Ok(packet) = Packet::deserialize(&data) { + let _ = endpoint.add_inbound_from(self.remote_id, packet); + } + } + + endpoint.take_outbound_clear(self.remote_id, |packet| { + let data = packet.serialize().unwrap(); + let _ = self.tx.send(data); + }); + } +} diff --git a/src/protocol/tests/mod.rs b/src/protocol/tests/mod.rs new file mode 100644 index 0000000..f60b960 --- /dev/null +++ b/src/protocol/tests/mod.rs @@ -0,0 +1,5 @@ +mod endpoint; +mod integration; +mod merkle_sync; +mod packet; +mod support; diff --git a/src/protocol/tests/oneshot/mod.rs b/src/protocol/tests/oneshot/mod.rs deleted file mode 100644 index e319bcb..0000000 --- a/src/protocol/tests/oneshot/mod.rs +++ /dev/null @@ -1,491 +0,0 @@ -mod streams; -mod support; - -use crate::protocol::{Endpoint, EndpointError, Leaf, RouteDirection}; - -use alloc::vec; - -use support::{ - CommsLeaf, ControllerLeaf, ENDPOINT_A, ENDPOINT_B, ENDPOINT_C, ResponderLeaf, - assert_hook_present, assert_hook_removed, echo_packet, echo_packet_with_end, endpoint_at, - single_inbound_packet, single_outbound_packet, -}; - -#[test] -fn test_oneshot() { - let (tx_a, rx_a) = crossbeam_channel::unbounded(); - let (tx_b, rx_b) = crossbeam_channel::unbounded(); - - let mut endpoint_a = Endpoint::new(ENDPOINT_A); - let mut controller_a = ControllerLeaf { has_run: false }; - let mut comms_a = CommsLeaf { - tx: tx_b, - rx: rx_a, - remote_id: ENDPOINT_B, - is_authority: false, - started: false, - }; - endpoint_a.path = vec![ENDPOINT_A]; - - let mut endpoint_b = Endpoint::new(ENDPOINT_B); - let mut responder_b = ResponderLeaf; - let mut comms_b = CommsLeaf { - tx: tx_a, - rx: rx_b, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - }; - endpoint_b.path = vec![ENDPOINT_A, ENDPOINT_B]; - - // Connections are registered routing state. The comms leaves also insert them - // during updates, but the first application packet should not depend on leaf order. - endpoint_a.add_connection(ENDPOINT_B, false); - endpoint_b.add_connection(ENDPOINT_A, true); - - // Cycle 1: A sends request to B - controller_a.update(&mut endpoint_a); - comms_a.update(&mut endpoint_a); - responder_b.update(&mut endpoint_b); - comms_b.update(&mut endpoint_b); - - // Cycle 2: B receives request and sends response to A - responder_b.update(&mut endpoint_b); - comms_b.update(&mut endpoint_b); - controller_a.update(&mut endpoint_a); - comms_a.update(&mut endpoint_a); - - // Cycle 3: A's CommsLeaf needs one more update to pull the packet from the channel - // and put it into the inbound queue. - controller_a.update(&mut endpoint_a); - comms_a.update(&mut endpoint_a); - - // Assertions on state - assert!( - Endpoint::route_contains(ENDPOINT_A, &endpoint_a.inbound), - "Endpoint A should have received response" - ); - assert_eq!( - Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) - .unwrap() - .len(), - 1, - "Endpoint A should have exactly one packet" - ); - let response = &Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) - .unwrap() - .front() - .unwrap(); - assert!(response.end_hook); - assert_eq!(response.data, "ABC123".as_bytes()); - assert!( - endpoint_b.hook_count() == 0, - "responder hook should be cleaned after the upward response" - ); - // assert_eq!(response.hook_id, HOOK_ECHO); -} - -#[test] -fn inbound_downward_packet_for_local_endpoint_opens_hook() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.add_connection(ENDPOINT_A, true); - - endpoint - .add_inbound_from( - ENDPOINT_A, - echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id), - ) - .unwrap(); - - let packet = single_inbound_packet(&endpoint, ENDPOINT_B); - assert!(!packet.end_hook); - assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B]); - assert_hook_present(&endpoint, hook_id); - assert_eq!(endpoint.hook_peer(hook_id), Some(ENDPOINT_A)); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn outbound_packet_for_local_endpoint_is_delivered_locally() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - - endpoint - .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) - .unwrap(); - - let packet = single_inbound_packet(&endpoint, ENDPOINT_B); - assert!(!packet.end_hook); - assert_eq!(packet.data, "ABC123".as_bytes()); - assert_hook_removed(&endpoint, hook_id); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn inbound_downward_packet_routes_to_immediate_child() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.add_connection(ENDPOINT_A, true); - endpoint.add_connection(ENDPOINT_C, false); - - endpoint - .add_inbound_from( - ENDPOINT_A, - echo_packet(vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C], hook_id), - ) - .unwrap(); - - let packet = single_outbound_packet(&endpoint, ENDPOINT_C); - assert!(!packet.end_hook); - assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); - assert_hook_present(&endpoint, hook_id); - assert_eq!(endpoint.hook_peer(hook_id), Some(ENDPOINT_C)); - assert!(!Endpoint::route_contains(ENDPOINT_A, &endpoint.outbound)); -} - -#[test] -fn outbound_downward_packet_routes_to_immediate_child() { - let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); - let hook_id = endpoint.get_hook_id(); - endpoint.accept_hook(hook_id, ENDPOINT_B); - endpoint.add_connection(ENDPOINT_B, false); - - endpoint - .add_outbound(echo_packet_with_end( - vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C], - hook_id, - true, - )) - .unwrap(); - - let packet = single_outbound_packet(&endpoint, ENDPOINT_B); - assert!(packet.end_hook); - assert_eq!(packet.path, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); - assert_hook_removed(&endpoint, hook_id); - assert!(!Endpoint::route_contains(ENDPOINT_C, &endpoint.outbound)); -} - -#[test] -fn inbound_upward_packet_with_hook_routes_to_parent() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.accept_hook(hook_id, ENDPOINT_C); - endpoint.add_connection(ENDPOINT_A, true); - endpoint.add_connection(ENDPOINT_C, false); - - endpoint - .add_inbound_from( - ENDPOINT_C, - echo_packet_with_end(vec![ENDPOINT_A], hook_id, true), - ) - .unwrap(); - - let packet = single_outbound_packet(&endpoint, ENDPOINT_A); - assert!(packet.end_hook); - assert_eq!(packet.hook_id, hook_id); - assert_hook_removed(&endpoint, hook_id); - assert!(!Endpoint::route_contains(ENDPOINT_C, &endpoint.outbound)); -} - -#[test] -fn inbound_upward_packet_without_hook_is_rejected() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.add_connection(ENDPOINT_A, true); - endpoint.add_connection(ENDPOINT_C, false); - - let error = endpoint - .add_inbound_from( - ENDPOINT_C, - echo_packet_with_end(vec![ENDPOINT_A], hook_id, true), - ) - .unwrap_err(); - - assert!(matches!( - error, - EndpointError::UnknownHook { hook_id: observed_hook_id } if observed_hook_id == hook_id - )); - assert!(Endpoint::routes_is_empty(&endpoint.inbound)); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn forged_upward_packet_with_unknown_hook_is_rejected() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - endpoint.accept_hook(7, ENDPOINT_C); - endpoint.add_connection(ENDPOINT_A, true); - endpoint.add_connection(ENDPOINT_C, false); - - let error = endpoint - .add_inbound_from(ENDPOINT_C, echo_packet_with_end(vec![ENDPOINT_A], 99, true)) - .unwrap_err(); - - assert!(matches!(error, EndpointError::UnknownHook { hook_id: 99 })); - assert_hook_present(&endpoint, 7); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn forged_sideways_packet_is_rejected_as_incorrect_path() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.accept_hook(hook_id, ENDPOINT_A); - endpoint.add_connection(ENDPOINT_A, true); - - let error = endpoint - .add_inbound_from( - ENDPOINT_A, - echo_packet(vec![ENDPOINT_A, ENDPOINT_C], hook_id), - ) - .unwrap_err(); - - assert!(matches!(error, EndpointError::DestinationOutsideLocalTree)); - assert_hook_present(&endpoint, hook_id); - assert!(Endpoint::routes_is_empty(&endpoint.inbound)); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn malformed_frame_is_dropped_by_comms_leaf() { - let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); - let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); - let mut endpoint = Endpoint::new(ENDPOINT_B); - let mut comms = CommsLeaf { - tx: tx_unused, - rx: rx_for_endpoint, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - }; - endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; - - tx_to_endpoint.send(vec![0, 1, 2, 3]).unwrap(); - comms.update(&mut endpoint); - - assert!(Endpoint::routes_is_empty(&endpoint.inbound)); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn malformed_frame_does_not_block_following_valid_packet() { - let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); - let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); - let hook_id = 42; - let mut endpoint = Endpoint::new(ENDPOINT_B); - let mut comms = CommsLeaf { - tx: tx_unused, - rx: rx_for_endpoint, - remote_id: ENDPOINT_A, - is_authority: true, - started: false, - }; - endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; - - tx_to_endpoint.send(vec![0, 1, 2, 3]).unwrap(); - tx_to_endpoint - .send( - echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id) - .serialize() - .unwrap(), - ) - .unwrap(); - comms.update(&mut endpoint); - - let packet = single_inbound_packet(&endpoint, ENDPOINT_B); - assert!(!packet.end_hook); - assert_eq!(packet.hook_id, hook_id); - assert_hook_present(&endpoint, hook_id); -} - -#[test] -fn forged_frame_without_required_hook_is_dropped_by_comms_leaf() { - let (tx_to_endpoint, rx_for_endpoint) = crossbeam_channel::unbounded(); - let (tx_unused, _rx_unused) = crossbeam_channel::unbounded(); - let mut endpoint = Endpoint::new(ENDPOINT_B); - let mut comms = CommsLeaf { - tx: tx_unused, - rx: rx_for_endpoint, - remote_id: ENDPOINT_C, - is_authority: false, - started: false, - }; - endpoint.path = vec![ENDPOINT_A, ENDPOINT_B]; - endpoint.accept_hook(7, ENDPOINT_C); - endpoint.add_connection(ENDPOINT_A, true); - - tx_to_endpoint - .send( - echo_packet_with_end(vec![ENDPOINT_A], 12, true) - .serialize() - .unwrap(), - ) - .unwrap(); - comms.update(&mut endpoint); - - assert_hook_present(&endpoint, 7); - assert!(Endpoint::routes_is_empty(&endpoint.inbound)); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn upward_outbound_without_hook_is_rejected() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - endpoint.accept_hook(7, ENDPOINT_A); - endpoint.add_connection(ENDPOINT_A, true); - - let new_hook = endpoint.get_hook_id(); - - let error = endpoint - .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], new_hook, true)) - .unwrap_err(); - - assert!(matches!( - error, - EndpointError::UnknownHook { hook_id: observed_hook_id } if observed_hook_id == new_hook - )); - assert_hook_present(&endpoint, 7); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn downward_outbound_without_hook_is_allowed() { - let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); - endpoint.add_connection(ENDPOINT_B, false); - - let new_hook = endpoint.get_hook_id(); - - endpoint - .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], new_hook)) - .unwrap(); - - assert_eq!( - Endpoint::route_get(ENDPOINT_B, &endpoint.outbound) - .unwrap() - .len(), - 1 - ); - assert_hook_present(&endpoint, new_hook); - assert_eq!(endpoint.hook_peer(new_hook), Some(ENDPOINT_B)); -} - -#[test] -fn deeper_upward_route_uses_parent_as_next_hop() { - let mut endpoint = endpoint_at(ENDPOINT_C, vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C]); - let new_hook = endpoint.get_hook_id(); - - endpoint.accept_hook(new_hook, ENDPOINT_B); - endpoint.add_connection(ENDPOINT_B, true); - - endpoint - .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], new_hook, true)) - .unwrap(); - - assert!(Endpoint::route_contains(ENDPOINT_B, &endpoint.outbound)); - assert!(!Endpoint::route_contains(ENDPOINT_A, &endpoint.outbound)); - assert_hook_removed(&endpoint, new_hook); -} - -#[test] -fn downward_route_without_connection_is_rejected() { - let mut endpoint = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); - let hook_id = endpoint.get_hook_id(); - - let error = endpoint - .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) - .unwrap_err(); - - assert!(matches!( - error, - EndpointError::MissingConnection { - next_hop: ENDPOINT_B, - direction: RouteDirection::Downward, - } - )); - assert_hook_removed(&endpoint, hook_id); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn upward_route_without_connection_is_rejected_even_with_hook() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.accept_hook(hook_id, ENDPOINT_A); - - let error = endpoint - .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) - .unwrap_err(); - - assert!(matches!( - error, - EndpointError::MissingConnection { - next_hop: ENDPOINT_A, - direction: RouteDirection::Upward, - } - )); - assert_hook_present(&endpoint, hook_id); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn end_hook_removes_hook_after_packet_is_queued() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.accept_hook(hook_id, ENDPOINT_A); - endpoint.add_connection(ENDPOINT_A, true); - - endpoint - .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) - .unwrap(); - - assert_hook_removed(&endpoint, hook_id); - assert_eq!( - single_outbound_packet(&endpoint, ENDPOINT_A).hook_id, - hook_id - ); -} - -#[test] -fn failed_end_hook_route_keeps_hook_state() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let hook_id = endpoint.get_hook_id(); - endpoint.accept_hook(hook_id, ENDPOINT_A); - - let error = endpoint - .add_outbound(echo_packet_with_end(vec![ENDPOINT_A], hook_id, true)) - .unwrap_err(); - - assert!(matches!( - error, - EndpointError::MissingConnection { - next_hop: ENDPOINT_A, - direction: RouteDirection::Upward, - } - )); - assert_hook_present(&endpoint, hook_id); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} - -#[test] -fn inbound_without_absolute_path_is_rejected() { - let mut endpoint = Endpoint::new(ENDPOINT_A); - - let error = endpoint - .add_inbound(echo_packet(vec![ENDPOINT_A], 1)) - .unwrap_err(); - - assert!(matches!(error, EndpointError::EndpointPathUnset)); - assert!(Endpoint::routes_is_empty(&endpoint.inbound)); -} - -#[test] -fn outbound_without_absolute_path_is_rejected() { - let mut endpoint = Endpoint::new(ENDPOINT_A); - - let error = endpoint - .add_outbound(echo_packet(vec![ENDPOINT_A], 1)) - .unwrap_err(); - - assert!(matches!(error, EndpointError::EndpointPathUnset)); - assert!(Endpoint::routes_is_empty(&endpoint.outbound)); -} diff --git a/src/protocol/tests/oneshot/support.rs b/src/protocol/tests/oneshot/support.rs deleted file mode 100644 index c1af87a..0000000 --- a/src/protocol/tests/oneshot/support.rs +++ /dev/null @@ -1,205 +0,0 @@ -use crate::protocol::{Endpoint, Leaf, Packet}; - -#[cfg(feature = "interface")] -use crate::protocol::LeafMeta; - -use alloc::{vec, vec::Vec}; -use crossbeam_channel::{Receiver, Sender}; - -pub(super) const ENDPOINT_A: u32 = 0; -pub(super) const ENDPOINT_B: u32 = 1; -pub(super) const ENDPOINT_C: u32 = 2; - -const LEAF_CONTROLLER: u32 = 100; -const LEAF_COMMS: u32 = 101; -const LEAF_RESPONDER: u32 = 102; - -/// Builds a test packet whose route is the only field varied by routing tests. -/// -/// Keeping the payload stable makes each assertion about endpoint behavior rather -/// than packet construction, which is important because forged and malformed cases -/// should fail before any leaf-level procedure handling would matter. -pub(super) fn echo_packet(path: Vec, hook_id: u16) -> Packet { - echo_packet_with_end(path, hook_id, false) -} - -/// Builds a test packet with an explicit hook-lifetime marker. -pub(super) fn echo_packet_with_end(path: Vec, hook_id: u16, end_hook: bool) -> Packet { - Packet { - hook_id, - end_hook, - path, - procedure_id: 1, - data: "ABC123".as_bytes().to_vec(), - } -} - -/// Creates a bare endpoint at a known absolute path. -/// -/// Most routing tests do not need leaves; they only need the endpoint's local path, -/// connection table, and hook table. This helper keeps that setup explicit without -/// hiding the routing state that each test is validating. -pub(super) fn endpoint_at(id: u32, path: Vec) -> Endpoint { - let mut endpoint = Endpoint::new(id); - endpoint.path = path; - endpoint -} - -/// Returns the only outbound packet queued for `next_hop`. -/// -/// Routing bugs often show up as packets being sent to the final destination rather -/// than the immediate neighbor. Tests use this helper to assert both that exactly one -/// packet exists and that it was queued for the expected adjacent endpoint. -pub(super) fn single_outbound_packet(endpoint: &Endpoint, next_hop: u32) -> &Packet { - let queue = Endpoint::route_get(next_hop, &endpoint.outbound) - .unwrap_or_else(|| panic!("expected one outbound queue for {next_hop}")); - assert_eq!(queue.len(), 1, "expected exactly one outbound packet"); - queue.front().unwrap() -} - -/// Returns the only inbound packet delivered to `local_id`. -/// -/// Local delivery is intentionally separate from transit forwarding, so the tests -/// assert against the local inbound queue instead of only checking that routing did -/// not produce an error. -pub(super) fn single_inbound_packet(endpoint: &Endpoint, local_id: u32) -> &Packet { - let queue = Endpoint::route_get(local_id, &endpoint.inbound) - .unwrap_or_else(|| panic!("expected one inbound queue for {local_id}")); - assert_eq!(queue.len(), 1, "expected exactly one inbound packet"); - queue.front().unwrap() -} - -/// Asserts that local hook state still contains `hook_id`. -/// -/// Tests use this instead of open-coded map checks so every lifecycle assertion -/// explains the intended routing invariant when it fails. -pub(super) fn assert_hook_present(endpoint: &Endpoint, hook_id: u16) { - assert!( - endpoint.has_hook(hook_id), - "expected hook {hook_id} to remain registered" - ); -} - -/// Asserts that local hook state no longer contains `hook_id`. -/// -/// Upward `end_hook` packets are the only cases that should remove hook state; -/// downward and local packets with the same flag must leave hooks alone. -pub(super) fn assert_hook_removed(endpoint: &Endpoint, hook_id: u16) { - assert!( - !endpoint.has_hook(hook_id), - "expected hook {hook_id} to be cleaned up" - ); -} - -pub(super) struct ControllerLeaf { - pub(super) has_run: bool, -} - -pub(super) struct CommsLeaf { - pub(super) tx: Sender>, - pub(super) rx: Receiver>, - - pub(super) remote_id: u32, - pub(super) is_authority: bool, - pub(super) started: bool, -} - -pub(super) struct ResponderLeaf; - -impl Leaf for ControllerLeaf { - fn get_id(&self) -> u32 { - LEAF_CONTROLLER - } - - #[cfg(feature = "interface")] - fn get_meta(&self) -> LeafMeta { - LeafMeta { - name: "Controller Leaf", - identifier: "dev.unshell.test.controller_leaf", - version: "v0", - authors: vec!["ASTATIN3"], - } - } - - fn update(&mut self, endpoint: &mut Endpoint) { - if !self.has_run { - // The controller starts exactly one request so the end-to-end test can - // assert deterministic routing without accumulating retries. - let hook_id = endpoint.get_hook_id(); - let packet = echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id); - let _ = endpoint.add_outbound(packet); - self.has_run = true; - } - } -} - -impl Leaf for CommsLeaf { - fn get_id(&self) -> u32 { - LEAF_COMMS - } - - #[cfg(feature = "interface")] - fn get_meta(&self) -> LeafMeta { - LeafMeta { - name: "Comms Leaf", - identifier: "dev.unshell.test.comms_leaf", - version: "v0", - authors: vec!["ASTATIN3"], - } - } - - fn update(&mut self, endpoint: &mut Endpoint) { - if !self.started { - endpoint.add_connection(self.remote_id, self.is_authority); - self.started = true; - } - - while !self.rx.is_empty() { - let data = self.rx.recv().unwrap(); - - // Transport bytes are untrusted. Dropping malformed frames here keeps - // the oneshot harness faithful to a router boundary: invalid wire data - // must not panic or poison later valid packets on the same connection. - if let Ok(packet) = Packet::deserialize(&data) { - let _ = endpoint.add_inbound_from(self.remote_id, packet); - } - } - - endpoint.take_outbound_clear(self.remote_id, |packet| { - let data = packet.serialize().unwrap(); - let _ = self.tx.send(data); - }); - } -} - -impl Leaf for ResponderLeaf { - fn get_id(&self) -> u32 { - LEAF_RESPONDER - } - - #[cfg(feature = "interface")] - fn get_meta(&self) -> LeafMeta { - LeafMeta { - name: "Responder Leaf", - identifier: "dev.unshell.test.responder_leaf", - version: "v0", - authors: vec!["ASTATIN3"], - } - } - - fn update(&mut self, endpoint: &mut Endpoint) { - let local_id = endpoint.path.last().cloned().unwrap_or(0); - let mut packets = Vec::new(); - - endpoint.take_inbound_clear(local_id, |packet| { - let mut response = echo_packet_with_end(vec![ENDPOINT_A], packet.hook_id, true); - response.hook_id = packet.hook_id; - response.data = packet.data.clone(); - packets.push(response); - }); - - for packet in packets { - let _ = endpoint.add_outbound(packet); - } - } -} diff --git a/src/protocol/tests/support/assertions.rs b/src/protocol/tests/support/assertions.rs new file mode 100644 index 0000000..f17aca8 --- /dev/null +++ b/src/protocol/tests/support/assertions.rs @@ -0,0 +1,23 @@ +use crate::protocol::Endpoint; + +/// Asserts that local hook state still contains `hook_id`. +/// +/// Tests use this instead of open-coded map checks so every lifecycle assertion +/// explains the intended routing invariant when it fails. +pub(crate) fn assert_hook_present(endpoint: &Endpoint, hook_id: u16) { + assert!( + endpoint.has_hook(hook_id), + "expected hook {hook_id} to remain registered" + ); +} + +/// Asserts that local hook state no longer contains `hook_id`. +/// +/// Upward `end_hook` packets are the only cases that should remove hook state; +/// downward and local packets with the same flag must leave hooks alone. +pub(crate) fn assert_hook_removed(endpoint: &Endpoint, hook_id: u16) { + assert!( + !endpoint.has_hook(hook_id), + "expected hook {hook_id} to be cleaned up" + ); +} diff --git a/src/protocol/tests/support/endpoints.rs b/src/protocol/tests/support/endpoints.rs new file mode 100644 index 0000000..ae45ede --- /dev/null +++ b/src/protocol/tests/support/endpoints.rs @@ -0,0 +1,42 @@ +use alloc::vec::Vec; + +use crate::protocol::{Endpoint, Packet}; + +pub(crate) const ENDPOINT_A: u32 = 0; +pub(crate) const ENDPOINT_B: u32 = 1; +pub(crate) const ENDPOINT_C: u32 = 2; + +/// Creates a bare endpoint at a known absolute path. +/// +/// Most routing tests do not need leaves; they only need the endpoint's local path, +/// connection table, and hook table. This helper keeps that setup explicit without +/// hiding the routing state that each test is validating. +pub(crate) fn endpoint_at(id: u32, path: Vec) -> Endpoint { + let mut endpoint = Endpoint::new(id); + endpoint.path = path; + endpoint +} + +/// Returns the only outbound packet queued for `next_hop`. +/// +/// Routing bugs often show up as packets being sent to the final destination rather +/// than the immediate neighbor. Tests use this helper to assert both that exactly one +/// packet exists and that it was queued for the expected adjacent endpoint. +pub(crate) fn single_outbound_packet(endpoint: &Endpoint, next_hop: u32) -> &Packet { + let queue = Endpoint::route_get(next_hop, &endpoint.outbound) + .unwrap_or_else(|| panic!("expected one outbound queue for {next_hop}")); + assert_eq!(queue.len(), 1, "expected exactly one outbound packet"); + queue.front().unwrap() +} + +/// Returns the only inbound packet delivered to `local_id`. +/// +/// Local delivery is intentionally separate from transit forwarding, so the tests +/// assert against the local inbound queue instead of only checking that routing did +/// not produce an error. +pub(crate) fn single_inbound_packet(endpoint: &Endpoint, local_id: u32) -> &Packet { + let queue = Endpoint::route_get(local_id, &endpoint.inbound) + .unwrap_or_else(|| panic!("expected one inbound queue for {local_id}")); + assert_eq!(queue.len(), 1, "expected exactly one inbound packet"); + queue.front().unwrap() +} diff --git a/src/protocol/tests/support/mod.rs b/src/protocol/tests/support/mod.rs new file mode 100644 index 0000000..c1eafbf --- /dev/null +++ b/src/protocol/tests/support/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod assertions; +pub(crate) mod endpoints; +pub(crate) mod packets; +pub(crate) mod transport; diff --git a/src/protocol/tests/support/packets.rs b/src/protocol/tests/support/packets.rs new file mode 100644 index 0000000..900ae3d --- /dev/null +++ b/src/protocol/tests/support/packets.rs @@ -0,0 +1,23 @@ +use alloc::vec::Vec; + +use crate::protocol::Packet; + +/// Builds a test packet whose route is the only field varied by routing tests. +/// +/// Keeping the payload stable makes each assertion about endpoint behavior rather +/// than packet construction, which is important because forged and malformed cases +/// should fail before any leaf-level procedure handling would matter. +pub(crate) fn echo_packet(path: Vec, hook_id: u16) -> Packet { + echo_packet_with_end(path, hook_id, false) +} + +/// Builds a test packet with an explicit hook-lifetime marker. +pub(crate) fn echo_packet_with_end(path: Vec, hook_id: u16, end_hook: bool) -> Packet { + Packet { + hook_id, + end_hook, + path, + procedure_id: 1, + data: "ABC123".as_bytes().to_vec(), + } +} diff --git a/src/protocol/tests/support/transport.rs b/src/protocol/tests/support/transport.rs new file mode 100644 index 0000000..cb4f008 --- /dev/null +++ b/src/protocol/tests/support/transport.rs @@ -0,0 +1,63 @@ +use alloc::vec::Vec; + +use crossbeam_channel::{Receiver, Sender}; + +use crate::protocol::{Endpoint, Leaf, Packet}; + +#[cfg(feature = "interface")] +use crate::protocol::LeafMeta; + +const LEAF_COMMS: u32 = 101; + +/// Mock transport leaf that serializes outbound packets through a channel pair. +/// +/// This is intentionally shared by protocol integration tests: it is the boundary +/// where structured packets become untrusted bytes and malformed frames get dropped +/// before reaching endpoint routing. +pub(crate) struct CommsLeaf { + pub(crate) tx: Sender>, + pub(crate) rx: Receiver>, + + pub(crate) remote_id: u32, + pub(crate) is_authority: bool, + pub(crate) started: bool, +} + +impl Leaf for CommsLeaf { + fn get_id(&self) -> u32 { + LEAF_COMMS + } + + #[cfg(feature = "interface")] + fn get_meta(&self) -> LeafMeta { + LeafMeta { + name: "Comms Leaf", + identifier: "dev.unshell.test.comms_leaf", + version: "v0", + authors: alloc::vec!["ASTATIN3"], + } + } + + fn update(&mut self, endpoint: &mut Endpoint) { + if !self.started { + endpoint.add_connection(self.remote_id, self.is_authority); + self.started = true; + } + + while !self.rx.is_empty() { + let data = self.rx.recv().unwrap(); + + // Transport bytes are untrusted. Dropping malformed frames here keeps + // integration harnesses faithful to a router boundary: invalid wire data + // must not panic or poison later valid packets on the same connection. + if let Ok(packet) = Packet::deserialize(&data) { + let _ = endpoint.add_inbound_from(self.remote_id, packet); + } + } + + endpoint.take_outbound_clear(self.remote_id, |packet| { + let data = packet.serialize().unwrap(); + let _ = self.tx.send(data); + }); + } +} diff --git a/unshell-leaves/leaf-pty/src/tests/session.rs b/unshell-leaves/leaf-pty/src/tests/session.rs deleted file mode 100644 index a8a30be..0000000 --- a/unshell-leaves/leaf-pty/src/tests/session.rs +++ /dev/null @@ -1,282 +0,0 @@ -use alloc::{vec, vec::Vec}; - -use unshell::protocol::{Leaf, Packet}; - -use crate::{ - FakePtyLeaf, FakePtyState, OP_ABORT, OP_ERROR, OP_EXIT, OP_INPUT, OP_OUTPUT, OP_STDIN_EOF, - OP_TERMINATE, pty_open_packet, -}; - -use super::support::{ - ENDPOINT_A, ENDPOINT_B, PROC_OTHER, assert_frame, assert_hook_present, assert_hook_removed, - assert_opened, drain_parent_pty_packets, endpoint_at, has_frame, open_pty_session, - pty_endpoints, send_downward_frame, transfer_packets, -}; - -#[test] -fn open_pty_paves_hook_and_creates_session() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - - let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(leaf.active_session_count(), 1); - assert_eq!(leaf.state().active_count, 1); - assert_eq!(leaf.state().total_opened, 1); - assert_hook_present(&endpoint_a, hook_id); - assert_hook_present(&endpoint_b, hook_id); - assert_eq!(packets.len(), 1); - assert_opened(&packets[0], hook_id); -} - -#[test] -fn input_and_output_share_one_hook() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_INPUT, - b"hello", - false, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_OUTPUT, false, b"hello"); - assert_hook_present(&endpoint_a, hook_id); - assert_hook_present(&endpoint_b, hook_id); -} - -#[test] -fn stdin_eof_keeps_hook_until_exit() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_STDIN_EOF, - &[], - false, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - - assert_eq!(leaf.state().last_stdin_eof_hook, Some(hook_id)); - assert!(drain_parent_pty_packets(&mut endpoint_a).is_empty()); - assert_hook_present(&endpoint_a, hook_id); - assert_hook_present(&endpoint_b, hook_id); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_TERMINATE, - &[], - false, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); - assert_eq!(leaf.active_session_count(), 0); - assert_hook_removed(&endpoint_a, hook_id); - assert_hook_removed(&endpoint_b, hook_id); -} - -#[test] -fn exit_end_hook_cleans_route_and_session() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_TERMINATE, - &[], - false, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); - assert_eq!(leaf.active_session_count(), 0); - assert_hook_removed(&endpoint_a, hook_id); - assert_hook_removed(&endpoint_b, hook_id); -} - -#[test] -fn failed_final_exit_route_closes_session_without_retry() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_TERMINATE, - &[], - false, - ); - endpoint_b.remove_connection(ENDPOINT_A, true); - leaf.update(&mut endpoint_b); - - assert_eq!(leaf.active_session_count(), 0); - assert_eq!(leaf.pending_packet_count(), 0); - assert_hook_removed(&endpoint_b, hook_id); - - endpoint_b.add_connection(ENDPOINT_A, true); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert!(packets.is_empty()); - assert_eq!(leaf.active_session_count(), 0); - assert_hook_present(&endpoint_a, hook_id); - assert_hook_removed(&endpoint_b, hook_id); -} - -#[test] -fn abort_downward_end_hook_closes_without_ack() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_ABORT, - &[], - true, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - - assert_eq!(leaf.active_session_count(), 0); - assert!(drain_parent_pty_packets(&mut endpoint_a).is_empty()); - assert_hook_removed(&endpoint_a, hook_id); - assert_hook_removed(&endpoint_b, hook_id); -} - -#[test] -fn unknown_session_input_returns_error_end_hook() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - let hook_id = endpoint_a.get_hook_id(); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - hook_id, - OP_INPUT, - b"orphan", - false, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(packets.len(), 1); - assert_frame(&packets[0], hook_id, OP_ERROR, true, b"unknown-session"); - assert_eq!(leaf.active_session_count(), 0); - assert_hook_removed(&endpoint_a, hook_id); - assert_hook_removed(&endpoint_b, hook_id); -} - -#[test] -fn two_pty_sessions_interleave_without_crossing_hooks() { - let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - - let first_hook = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - let second_hook = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); - drain_parent_pty_packets(&mut endpoint_a); - - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - second_hook, - OP_INPUT, - b"second", - false, - ); - send_downward_frame( - &mut endpoint_a, - &mut endpoint_b, - first_hook, - OP_INPUT, - b"first", - false, - ); - leaf.update(&mut endpoint_b); - transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); - let packets = drain_parent_pty_packets(&mut endpoint_a); - - assert_eq!(leaf.active_session_count(), 2); - assert_eq!(packets.len(), 2); - assert!(has_frame(&packets, first_hook, OP_OUTPUT, b"first")); - assert!(has_frame(&packets, second_hook, OP_OUTPUT, b"second")); - assert_hook_present(&endpoint_a, first_hook); - assert_hook_present(&endpoint_a, second_hook); - assert_hook_present(&endpoint_b, first_hook); - assert_hook_present(&endpoint_b, second_hook); -} - -#[test] -fn pty_leaf_does_not_consume_other_leaf_packets() { - let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - let mut leaf = FakePtyLeaf::new(FakePtyState::new()); - endpoint.add_connection(ENDPOINT_A, true); - - endpoint - .add_inbound_from(ENDPOINT_A, pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], 7)) - .unwrap(); - endpoint - .add_inbound_from( - ENDPOINT_A, - Packet { - hook_id: 8, - end_hook: false, - path: vec![ENDPOINT_A, ENDPOINT_B], - procedure_id: PROC_OTHER, - data: b"leave-me".to_vec(), - }, - ) - .unwrap(); - - leaf.update(&mut endpoint); - - let mut other_packets = Vec::new(); - endpoint.take_inbound_matching( - ENDPOINT_B, - |packet| packet.procedure_id == PROC_OTHER, - |packet| other_packets.push(packet), - ); - - assert_eq!(leaf.active_session_count(), 1); - assert_eq!(other_packets.len(), 1); - assert_eq!(other_packets[0].procedure_id, PROC_OTHER); - assert_eq!(other_packets[0].data, b"leave-me".to_vec()); -} diff --git a/unshell-leaves/leaf-pty/src/tests/session/concurrency.rs b/unshell-leaves/leaf-pty/src/tests/session/concurrency.rs new file mode 100644 index 0000000..a10b692 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/session/concurrency.rs @@ -0,0 +1,47 @@ +use unshell::protocol::Leaf; + +use crate::{FakePtyLeaf, FakePtyState, OP_INPUT, OP_OUTPUT}; + +use super::super::support::{ + ENDPOINT_A, ENDPOINT_B, assert_hook_present, drain_parent_pty_packets, has_frame, + open_pty_session, pty_endpoints, send_downward_frame, transfer_packets, +}; + +#[test] +fn two_pty_sessions_interleave_without_crossing_hooks() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + + let first_hook = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + let second_hook = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + second_hook, + OP_INPUT, + b"second", + false, + ); + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + first_hook, + OP_INPUT, + b"first", + false, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert_eq!(leaf.active_session_count(), 2); + assert_eq!(packets.len(), 2); + assert!(has_frame(&packets, first_hook, OP_OUTPUT, b"first")); + assert!(has_frame(&packets, second_hook, OP_OUTPUT, b"second")); + assert_hook_present(&endpoint_a, first_hook); + assert_hook_present(&endpoint_a, second_hook); + assert_hook_present(&endpoint_b, first_hook); + assert_hook_present(&endpoint_b, second_hook); +} diff --git a/unshell-leaves/leaf-pty/src/tests/session/failure.rs b/unshell-leaves/leaf-pty/src/tests/session/failure.rs new file mode 100644 index 0000000..c01078b --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/session/failure.rs @@ -0,0 +1,41 @@ +use unshell::protocol::Leaf; + +use crate::{FakePtyLeaf, FakePtyState, OP_TERMINATE}; + +use super::super::support::{ + ENDPOINT_A, ENDPOINT_B, assert_hook_present, assert_hook_removed, drain_parent_pty_packets, + open_pty_session, pty_endpoints, send_downward_frame, transfer_packets, +}; + +#[test] +fn failed_final_exit_route_closes_session_without_retry() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_TERMINATE, + &[], + false, + ); + endpoint_b.remove_connection(ENDPOINT_A, true); + leaf.update(&mut endpoint_b); + + assert_eq!(leaf.active_session_count(), 0); + assert_eq!(leaf.pending_packet_count(), 0); + assert_hook_removed(&endpoint_b, hook_id); + + endpoint_b.add_connection(ENDPOINT_A, true); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert!(packets.is_empty()); + assert_eq!(leaf.active_session_count(), 0); + assert_hook_present(&endpoint_a, hook_id); + assert_hook_removed(&endpoint_b, hook_id); +} diff --git a/unshell-leaves/leaf-pty/src/tests/session/filtering.rs b/unshell-leaves/leaf-pty/src/tests/session/filtering.rs new file mode 100644 index 0000000..becb18b --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/session/filtering.rs @@ -0,0 +1,44 @@ +use alloc::{vec, vec::Vec}; + +use unshell::protocol::{Leaf, Packet}; + +use crate::{FakePtyLeaf, FakePtyState, pty_open_packet}; + +use super::super::support::{ENDPOINT_A, ENDPOINT_B, PROC_OTHER, endpoint_at}; + +#[test] +fn pty_leaf_does_not_consume_other_leaf_packets() { + let mut endpoint = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + endpoint.add_connection(ENDPOINT_A, true); + + endpoint + .add_inbound_from(ENDPOINT_A, pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], 7)) + .unwrap(); + endpoint + .add_inbound_from( + ENDPOINT_A, + Packet { + hook_id: 8, + end_hook: false, + path: vec![ENDPOINT_A, ENDPOINT_B], + procedure_id: PROC_OTHER, + data: b"leave-me".to_vec(), + }, + ) + .unwrap(); + + leaf.update(&mut endpoint); + + let mut other_packets = Vec::new(); + endpoint.take_inbound_matching( + ENDPOINT_B, + |packet| packet.procedure_id == PROC_OTHER, + |packet| other_packets.push(packet), + ); + + assert_eq!(leaf.active_session_count(), 1); + assert_eq!(other_packets.len(), 1); + assert_eq!(other_packets[0].procedure_id, PROC_OTHER); + assert_eq!(other_packets[0].data, b"leave-me".to_vec()); +} diff --git a/unshell-leaves/leaf-pty/src/tests/session/input_output.rs b/unshell-leaves/leaf-pty/src/tests/session/input_output.rs new file mode 100644 index 0000000..ae585d4 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/session/input_output.rs @@ -0,0 +1,33 @@ +use unshell::protocol::Leaf; + +use crate::{FakePtyLeaf, FakePtyState, OP_INPUT, OP_OUTPUT}; + +use super::super::support::{ + ENDPOINT_A, ENDPOINT_B, assert_frame, assert_hook_present, drain_parent_pty_packets, + open_pty_session, pty_endpoints, send_downward_frame, transfer_packets, +}; + +#[test] +fn input_and_output_share_one_hook() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_INPUT, + b"hello", + false, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert_eq!(packets.len(), 1); + assert_frame(&packets[0], hook_id, OP_OUTPUT, false, b"hello"); + assert_hook_present(&endpoint_a, hook_id); + assert_hook_present(&endpoint_b, hook_id); +} diff --git a/unshell-leaves/leaf-pty/src/tests/session/lifecycle.rs b/unshell-leaves/leaf-pty/src/tests/session/lifecycle.rs new file mode 100644 index 0000000..deedeb3 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/session/lifecycle.rs @@ -0,0 +1,145 @@ +use unshell::protocol::Leaf; + +use crate::{ + FakePtyLeaf, FakePtyState, OP_ABORT, OP_ERROR, OP_EXIT, OP_INPUT, OP_STDIN_EOF, OP_TERMINATE, +}; + +use super::super::support::{ + ENDPOINT_A, ENDPOINT_B, assert_frame, assert_hook_present, assert_hook_removed, assert_opened, + drain_parent_pty_packets, open_pty_session, pty_endpoints, send_downward_frame, + transfer_packets, +}; + +#[test] +fn open_pty_paves_hook_and_creates_session() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + + let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert_eq!(leaf.active_session_count(), 1); + assert_eq!(leaf.state().active_count, 1); + assert_eq!(leaf.state().total_opened, 1); + assert_hook_present(&endpoint_a, hook_id); + assert_hook_present(&endpoint_b, hook_id); + assert_eq!(packets.len(), 1); + assert_opened(&packets[0], hook_id); +} + +#[test] +fn stdin_eof_keeps_hook_until_exit() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_STDIN_EOF, + &[], + false, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + + assert_eq!(leaf.state().last_stdin_eof_hook, Some(hook_id)); + assert!(drain_parent_pty_packets(&mut endpoint_a).is_empty()); + assert_hook_present(&endpoint_a, hook_id); + assert_hook_present(&endpoint_b, hook_id); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_TERMINATE, + &[], + false, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert_eq!(packets.len(), 1); + assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); + assert_eq!(leaf.active_session_count(), 0); + assert_hook_removed(&endpoint_a, hook_id); + assert_hook_removed(&endpoint_b, hook_id); +} + +#[test] +fn exit_end_hook_cleans_route_and_session() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_TERMINATE, + &[], + false, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert_eq!(packets.len(), 1); + assert_frame(&packets[0], hook_id, OP_EXIT, true, &[0]); + assert_eq!(leaf.active_session_count(), 0); + assert_hook_removed(&endpoint_a, hook_id); + assert_hook_removed(&endpoint_b, hook_id); +} + +#[test] +fn abort_downward_end_hook_closes_without_ack() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let hook_id = open_pty_session(&mut endpoint_a, &mut endpoint_b, &mut leaf); + drain_parent_pty_packets(&mut endpoint_a); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_ABORT, + &[], + true, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + + assert_eq!(leaf.active_session_count(), 0); + assert!(drain_parent_pty_packets(&mut endpoint_a).is_empty()); + assert_hook_removed(&endpoint_a, hook_id); + assert_hook_removed(&endpoint_b, hook_id); +} + +#[test] +fn unknown_session_input_returns_error_end_hook() { + let (mut endpoint_a, mut endpoint_b) = pty_endpoints(); + let mut leaf = FakePtyLeaf::new(FakePtyState::new()); + let hook_id = endpoint_a.get_hook_id(); + + send_downward_frame( + &mut endpoint_a, + &mut endpoint_b, + hook_id, + OP_INPUT, + b"orphan", + false, + ); + leaf.update(&mut endpoint_b); + transfer_packets(&mut endpoint_b, &mut endpoint_a, ENDPOINT_A, ENDPOINT_B); + let packets = drain_parent_pty_packets(&mut endpoint_a); + + assert_eq!(packets.len(), 1); + assert_frame(&packets[0], hook_id, OP_ERROR, true, b"unknown-session"); + assert_eq!(leaf.active_session_count(), 0); + assert_hook_removed(&endpoint_a, hook_id); + assert_hook_removed(&endpoint_b, hook_id); +} diff --git a/unshell-leaves/leaf-pty/src/tests/session/mod.rs b/unshell-leaves/leaf-pty/src/tests/session/mod.rs new file mode 100644 index 0000000..c35a33b --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/session/mod.rs @@ -0,0 +1,5 @@ +mod concurrency; +mod failure; +mod filtering; +mod input_output; +mod lifecycle; diff --git a/unshell-leaves/leaf-pty/src/tests/support.rs b/unshell-leaves/leaf-pty/src/tests/support.rs deleted file mode 100644 index 5c1d179..0000000 --- a/unshell-leaves/leaf-pty/src/tests/support.rs +++ /dev/null @@ -1,137 +0,0 @@ -use alloc::{vec, vec::Vec}; - -use unshell::protocol::{Endpoint, Leaf, Packet}; - -use crate::{ - FakePtyLeaf, OP_OPENED, PROC_PTY, frame_opcode, frame_payload, pty_open_packet, pty_packet, -}; - -pub(super) const ENDPOINT_A: u32 = 0; -pub(super) const ENDPOINT_B: u32 = 1; -pub(super) const PROC_OTHER: u32 = 31; - -/// Creates a bare endpoint at a known absolute path. -pub(super) fn endpoint_at(id: u32, path: Vec) -> Endpoint { - let mut endpoint = Endpoint::new(id); - endpoint.path = path; - endpoint -} - -/// Creates the parent/child endpoint pair used by PTY session tests. -pub(super) fn pty_endpoints() -> (Endpoint, Endpoint) { - let mut endpoint_a = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); - let mut endpoint_b = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); - - endpoint_a.add_connection(ENDPOINT_B, false); - endpoint_b.add_connection(ENDPOINT_A, true); - - (endpoint_a, endpoint_b) -} - -/// Transfers every queued packet for `next_hop` into `receiver` as `remote_id` traffic. -pub(super) fn transfer_packets( - sender: &mut Endpoint, - receiver: &mut Endpoint, - next_hop: u32, - remote_id: u32, -) { - let mut packets = Vec::new(); - sender.take_outbound_clear(next_hop, |packet| packets.push(packet.clone())); - - for packet in packets { - receiver.add_inbound_from(remote_id, packet).unwrap(); - } -} - -/// Sends one downward PTY frame from endpoint A to endpoint B. -pub(super) fn send_downward_frame( - endpoint_a: &mut Endpoint, - endpoint_b: &mut Endpoint, - hook_id: u16, - opcode: u8, - payload: &[u8], - end_hook: bool, -) { - endpoint_a - .add_outbound(pty_packet( - vec![ENDPOINT_A, ENDPOINT_B], - hook_id, - end_hook, - opcode, - payload, - )) - .unwrap(); - transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); -} - -/// Opens a fake PTY session and delivers the `Opened` response to endpoint A. -pub(super) fn open_pty_session( - endpoint_a: &mut Endpoint, - endpoint_b: &mut Endpoint, - leaf: &mut FakePtyLeaf, -) -> u16 { - let hook_id = endpoint_a.get_hook_id(); - endpoint_a - .add_outbound(pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) - .unwrap(); - - transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); - leaf.update(endpoint_b); - transfer_packets(endpoint_b, endpoint_a, ENDPOINT_A, ENDPOINT_B); - - hook_id -} - -/// Drains packets for `procedure_id` delivered to endpoint A. -pub(super) fn drain_parent_packets(endpoint: &mut Endpoint, procedure_id: u32) -> Vec { - let mut packets = Vec::new(); - endpoint.take_inbound_matching( - ENDPOINT_A, - |packet| packet.procedure_id == procedure_id, - |packet| packets.push(packet), - ); - packets -} - -/// Drains PTY packets delivered to endpoint A. -pub(super) fn drain_parent_pty_packets(endpoint: &mut Endpoint) -> Vec { - drain_parent_packets(endpoint, PROC_PTY) -} - -/// Asserts that local hook state still contains `hook_id`. -pub(super) fn assert_hook_present(endpoint: &Endpoint, hook_id: u16) { - assert!(endpoint.has_hook(hook_id)); -} - -/// Asserts that local hook state no longer contains `hook_id`. -pub(super) fn assert_hook_removed(endpoint: &Endpoint, hook_id: u16) { - assert!(!endpoint.has_hook(hook_id)); -} - -/// Asserts that `packet` carries the expected PTY frame. -pub(super) fn assert_frame( - packet: &Packet, - hook_id: u16, - opcode: u8, - end_hook: bool, - payload: &[u8], -) { - assert_eq!(packet.hook_id, hook_id); - assert_eq!(packet.end_hook, end_hook); - assert_eq!(frame_opcode(packet), Some(opcode)); - assert_eq!(frame_payload(packet), payload); -} - -/// Returns true when `packets` contains the requested frame. -pub(super) fn has_frame(packets: &[Packet], hook_id: u16, opcode: u8, payload: &[u8]) -> bool { - packets.iter().any(|packet| { - packet.hook_id == hook_id - && frame_opcode(packet) == Some(opcode) - && frame_payload(packet) == payload - }) -} - -/// Asserts that a packet is the fake PTY open acknowledgement. -pub(super) fn assert_opened(packet: &Packet, hook_id: u16) { - assert_frame(packet, hook_id, OP_OPENED, false, &[]); -} diff --git a/unshell-leaves/leaf-pty/src/tests/support/assertions.rs b/unshell-leaves/leaf-pty/src/tests/support/assertions.rs new file mode 100644 index 0000000..87481b7 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/support/assertions.rs @@ -0,0 +1,47 @@ +use unshell::protocol::{Endpoint, Packet}; + +use crate::{OP_OPENED, frame_opcode, frame_payload}; + +/// Asserts that local hook state still contains `hook_id`. +pub(crate) fn assert_hook_present(endpoint: &Endpoint, hook_id: u16) { + assert!( + endpoint.has_hook(hook_id), + "expected hook {hook_id} to remain registered" + ); +} + +/// Asserts that local hook state no longer contains `hook_id`. +pub(crate) fn assert_hook_removed(endpoint: &Endpoint, hook_id: u16) { + assert!( + !endpoint.has_hook(hook_id), + "expected hook {hook_id} to be cleaned up" + ); +} + +/// Asserts that `packet` carries the expected PTY frame. +pub(crate) fn assert_frame( + packet: &Packet, + hook_id: u16, + opcode: u8, + end_hook: bool, + payload: &[u8], +) { + assert_eq!(packet.hook_id, hook_id); + assert_eq!(packet.end_hook, end_hook); + assert_eq!(frame_opcode(packet), Some(opcode)); + assert_eq!(frame_payload(packet), payload); +} + +/// Returns true when `packets` contains the requested frame. +pub(crate) fn has_frame(packets: &[Packet], hook_id: u16, opcode: u8, payload: &[u8]) -> bool { + packets.iter().any(|packet| { + packet.hook_id == hook_id + && frame_opcode(packet) == Some(opcode) + && frame_payload(packet) == payload + }) +} + +/// Asserts that a packet is the fake PTY open acknowledgement. +pub(crate) fn assert_opened(packet: &Packet, hook_id: u16) { + assert_frame(packet, hook_id, OP_OPENED, false, &[]); +} diff --git a/unshell-leaves/leaf-pty/src/tests/support/drains.rs b/unshell-leaves/leaf-pty/src/tests/support/drains.rs new file mode 100644 index 0000000..344747c --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/support/drains.rs @@ -0,0 +1,23 @@ +use alloc::vec::Vec; + +use unshell::protocol::{Endpoint, Packet}; + +use crate::PROC_PTY; + +use super::ENDPOINT_A; + +/// Drains packets for `procedure_id` delivered to endpoint A. +pub(crate) fn drain_parent_packets(endpoint: &mut Endpoint, procedure_id: u32) -> Vec { + let mut packets = Vec::new(); + endpoint.take_inbound_matching( + ENDPOINT_A, + |packet| packet.procedure_id == procedure_id, + |packet| packets.push(packet), + ); + packets +} + +/// Drains PTY packets delivered to endpoint A. +pub(crate) fn drain_parent_pty_packets(endpoint: &mut Endpoint) -> Vec { + drain_parent_packets(endpoint, PROC_PTY) +} diff --git a/unshell-leaves/leaf-pty/src/tests/support/endpoints.rs b/unshell-leaves/leaf-pty/src/tests/support/endpoints.rs new file mode 100644 index 0000000..a8d6413 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/support/endpoints.rs @@ -0,0 +1,40 @@ +use alloc::{vec, vec::Vec}; + +use unshell::protocol::{Endpoint, Packet}; + +pub(crate) const ENDPOINT_A: u32 = 0; +pub(crate) const ENDPOINT_B: u32 = 1; +pub(crate) const PROC_OTHER: u32 = 31; + +/// Creates a bare endpoint at a known absolute path. +pub(crate) fn endpoint_at(id: u32, path: Vec) -> Endpoint { + let mut endpoint = Endpoint::new(id); + endpoint.path = path; + endpoint +} + +/// Creates the parent/child endpoint pair used by PTY session tests. +pub(crate) fn pty_endpoints() -> (Endpoint, Endpoint) { + let mut endpoint_a = endpoint_at(ENDPOINT_A, vec![ENDPOINT_A]); + let mut endpoint_b = endpoint_at(ENDPOINT_B, vec![ENDPOINT_A, ENDPOINT_B]); + + endpoint_a.add_connection(ENDPOINT_B, false); + endpoint_b.add_connection(ENDPOINT_A, true); + + (endpoint_a, endpoint_b) +} + +/// Transfers every queued packet for `next_hop` into `receiver` as `remote_id` traffic. +pub(crate) fn transfer_packets( + sender: &mut Endpoint, + receiver: &mut Endpoint, + next_hop: u32, + remote_id: u32, +) { + let mut packets = Vec::::new(); + sender.take_outbound_clear(next_hop, |packet| packets.push(packet.clone())); + + for packet in packets { + receiver.add_inbound_from(remote_id, packet).unwrap(); + } +} diff --git a/unshell-leaves/leaf-pty/src/tests/support/mod.rs b/unshell-leaves/leaf-pty/src/tests/support/mod.rs new file mode 100644 index 0000000..96d60ae --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/support/mod.rs @@ -0,0 +1,9 @@ +mod assertions; +mod drains; +mod endpoints; +mod packets; + +pub(crate) use assertions::*; +pub(crate) use drains::*; +pub(crate) use endpoints::*; +pub(crate) use packets::*; diff --git a/unshell-leaves/leaf-pty/src/tests/support/packets.rs b/unshell-leaves/leaf-pty/src/tests/support/packets.rs new file mode 100644 index 0000000..44f50e9 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests/support/packets.rs @@ -0,0 +1,46 @@ +use alloc::vec; + +use unshell::protocol::{Endpoint, Leaf}; + +use crate::{FakePtyLeaf, pty_open_packet, pty_packet}; + +use super::{ENDPOINT_A, ENDPOINT_B, transfer_packets}; + +/// Sends one downward PTY frame from endpoint A to endpoint B. +pub(crate) fn send_downward_frame( + endpoint_a: &mut Endpoint, + endpoint_b: &mut Endpoint, + hook_id: u16, + opcode: u8, + payload: &[u8], + end_hook: bool, +) { + endpoint_a + .add_outbound(pty_packet( + vec![ENDPOINT_A, ENDPOINT_B], + hook_id, + end_hook, + opcode, + payload, + )) + .unwrap(); + transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); +} + +/// Opens a fake PTY session and delivers the `Opened` response to endpoint A. +pub(crate) fn open_pty_session( + endpoint_a: &mut Endpoint, + endpoint_b: &mut Endpoint, + leaf: &mut FakePtyLeaf, +) -> u16 { + let hook_id = endpoint_a.get_hook_id(); + endpoint_a + .add_outbound(pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id)) + .unwrap(); + + transfer_packets(endpoint_a, endpoint_b, ENDPOINT_B, ENDPOINT_A); + leaf.update(endpoint_b); + transfer_packets(endpoint_b, endpoint_a, ENDPOINT_A, ENDPOINT_B); + + hook_id +} From 9ab130a620dc9600bda4058e2142a0c08e6d9dd1 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Mon, 1 Jun 2026 15:53:38 -0600 Subject: [PATCH 12/12] Add tcp_simple transport crate --- Cargo.lock | 7 + Cargo.toml | 2 +- src/protocol/endpoint/queues.rs | 9 + src/protocol/packet.rs | 17 +- unshell-leaves/tcp_simple/Cargo.toml | 28 ++ unshell-leaves/tcp_simple/src/client/mod.rs | 44 +++ unshell-leaves/tcp_simple/src/lib.rs | 33 ++ unshell-leaves/tcp_simple/src/server/mod.rs | 83 +++++ unshell-leaves/tcp_simple/src/transport.rs | 348 ++++++++++++++++++++ 9 files changed, 568 insertions(+), 3 deletions(-) create mode 100644 unshell-leaves/tcp_simple/Cargo.toml create mode 100644 unshell-leaves/tcp_simple/src/client/mod.rs create mode 100644 unshell-leaves/tcp_simple/src/lib.rs create mode 100644 unshell-leaves/tcp_simple/src/server/mod.rs create mode 100644 unshell-leaves/tcp_simple/src/transport.rs diff --git a/Cargo.lock b/Cargo.lock index 6e9a5a8..3b201aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1685,6 +1685,13 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "tcp_simple" +version = "0.1.0" +dependencies = [ + "unshell", +] + [[package]] name = "terminfo" version = "0.9.0" diff --git a/Cargo.toml b/Cargo.toml index f0abdc6..6c49826 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ members = [ "ush-obfuscate", "base62", - "unshell-leaves/leaf-pty", "unshell-leaves/leaf-shell", "examples/endpoint_test", + "unshell-leaves/leaf-pty", "unshell-leaves/leaf-shell", "examples/endpoint_test", "unshell-leaves/tcp_simple", ] resolver = "2" diff --git a/src/protocol/endpoint/queues.rs b/src/protocol/endpoint/queues.rs index ed09b9e..0f17d5b 100644 --- a/src/protocol/endpoint/queues.rs +++ b/src/protocol/endpoint/queues.rs @@ -50,6 +50,15 @@ impl Endpoint { Self::take_clear(path, f, &mut self.outbound); } + /// Removes and returns all outbound packets queued for `path`. + /// + /// Transport leaves use this when they need packet ownership instead of a borrowed + /// callback. Keeping this non-generic avoids creating a new closure-shaped copy of + /// the queue-draining loop for each concrete transport implementation. + pub fn take_outbound_queue(&mut self, path: u32) -> Option { + Self::route_remove(path, &mut self.outbound) + } + fn take_clear(path: u32, mut f: F, queue: &mut RouteMap) where F: FnMut(&Packet), diff --git a/src/protocol/packet.rs b/src/protocol/packet.rs index 3d07ecc..bf30bf1 100644 --- a/src/protocol/packet.rs +++ b/src/protocol/packet.rs @@ -31,6 +31,17 @@ impl Packet { /// validation path. That makes deserialization a single full-packet parse, /// which matches how the endpoint mock transports actually consume packets. pub fn serialize(&self) -> Result, SerializeError> { + let mut buf = Vec::new(); + self.serialize_into(&mut buf)?; + Ok(buf) + } + + /// Appends this packet's serialized frame to an existing byte buffer. + /// + /// Transports use this to avoid allocating a temporary frame only to copy it into + /// their socket write buffer. The method performs all size checks before writing so + /// serialization errors do not leave a partial frame in `buf`. + pub fn serialize_into(&self, buf: &mut Vec) -> Result<(), SerializeError> { let path_len = u32::try_from(self.path.len()).map_err(|_| SerializeError::PathTooLarge)?; // body = fixed procedure_id field + data bytes @@ -49,7 +60,8 @@ impl Packet { .and_then(|n| n.checked_add(4)) .and_then(|n| n.checked_add(body_payload_len)) .ok_or(SerializeError::BodyTooLarge)?; - let mut buf = Vec::with_capacity(total); + + buf.reserve(total); // ── header ──────────────────────────────────────────────────────────── let flags = self.end_hook as u8; @@ -66,7 +78,7 @@ impl Packet { buf.extend_from_slice(&self.procedure_id.to_le_bytes()); buf.extend_from_slice(&self.data); - Ok(buf) + Ok(()) } /// Deserializes a full packet from untrusted transport bytes. @@ -75,6 +87,7 @@ impl Packet { /// partial parse path was removed because current routing tests and mock /// transports always deserialize before calling endpoint routing, so keeping a /// borrowed header API only preserved unused unsafe casting complexity. + #[inline(never)] pub fn deserialize(buf: &[u8]) -> Result { // fixed prefix: hook_id (2) + flags (1) + padding (1) + path_len (4) if buf.len() < 8 { diff --git a/unshell-leaves/tcp_simple/Cargo.toml b/unshell-leaves/tcp_simple/Cargo.toml new file mode 100644 index 0000000..4c9deb7 --- /dev/null +++ b/unshell-leaves/tcp_simple/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "tcp_simple" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +include.workspace = true + +[dependencies] +unshell = { workspace = true } + +[features] +default = [] +interface = ["unshell/interface"] +interface_ratatui = ["interface", "unshell/interface_ratatui"] + +[lints.rust] +elided_lifetimes_in_paths = "warn" +future_incompatible = { level = "warn", priority = -1 } +nonstandard_style = { level = "warn", priority = -1 } +rust_2018_idioms = { level = "warn", priority = -1 } +rust_2021_prelude_collisions = "warn" +semicolon_in_expressions_from_macros = "warn" +unsafe_op_in_unsafe_fn = "warn" +unused_import_braces = "warn" +unused_lifetimes = "warn" +trivial_casts = "allow" diff --git a/unshell-leaves/tcp_simple/src/client/mod.rs b/unshell-leaves/tcp_simple/src/client/mod.rs new file mode 100644 index 0000000..0892e76 --- /dev/null +++ b/unshell-leaves/tcp_simple/src/client/mod.rs @@ -0,0 +1,44 @@ +use std::{io, net::TcpStream, net::ToSocketAddrs}; + +use unshell::protocol::{Endpoint, Leaf}; + +use crate::transport::TcpBridge; + +/// TCP client-side transport leaf for one upstream endpoint. +/// +/// This is the mirror of [`crate::TCPServerLeaf`]: bytes from the connected server +/// are routed through [`Endpoint::add_inbound_from`], and packets queued for the +/// parent endpoint are serialized back onto the TCP stream. +#[derive(Debug)] +pub struct TCPClientLeaf { + bridge: TcpBridge, +} + +impl TCPClientLeaf { + /// Connects to an upstream TCP server and registers it as the authority peer. + /// + /// `parent_endpoint_id` must be the adjacent parent segment in this endpoint's + /// path. The connection is made during construction so failed startup is explicit + /// instead of being hidden as a permanently idle leaf. + pub fn new(connect_addr: A, parent_endpoint_id: u32) -> io::Result + where + A: ToSocketAddrs, + { + let stream = TcpStream::connect(connect_addr)?; + let mut bridge = TcpBridge::new(parent_endpoint_id, true); + bridge.set_stream(stream)?; + + Ok(Self { bridge }) + } +} + +impl Leaf for TCPClientLeaf { + fn get_id(&self) -> u32 { + crate::IDENTIFIER_CLIENT_HASH + } + + fn update(&mut self, endpoint: &mut Endpoint) { + self.bridge.register(endpoint); + self.bridge.update(endpoint); + } +} diff --git a/unshell-leaves/tcp_simple/src/lib.rs b/unshell-leaves/tcp_simple/src/lib.rs new file mode 100644 index 0000000..9a160b2 --- /dev/null +++ b/unshell-leaves/tcp_simple/src/lib.rs @@ -0,0 +1,33 @@ +//! Minimal TCP transport leaves for adjacent UnShell endpoints. +//! +//! This crate deliberately stays small: it does not own an [`unshell::protocol::Endpoint`] +//! or run a scheduler. Callers keep their endpoint and application leaves, then tick a +//! TCP leaf to move serialized packets between the endpoint's outbound queues and a +//! nonblocking socket. + +use unshell::crypto::hash_str_32; + +mod client; +mod server; +mod transport; + +pub use client::TCPClientLeaf; +pub use server::TCPServerLeaf; + +macro_rules! version { + () => { + env!("CARGO_PKG_VERSION") + }; +} + +/// Stable interface identifier for the listening TCP bridge leaf. +pub const IDENTIFIER_SERVER: &str = concat!("dev.unshell.", version!(), ".tcp_simple.server"); + +/// Numeric identifier for [`TCPServerLeaf`]. +pub const IDENTIFIER_SERVER_HASH: u32 = hash_str_32(IDENTIFIER_SERVER); + +/// Stable interface identifier for the connecting TCP bridge leaf. +pub const IDENTIFIER_CLIENT: &str = concat!("dev.unshell.", version!(), ".tcp_simple.client"); + +/// Numeric identifier for [`TCPClientLeaf`]. +pub const IDENTIFIER_CLIENT_HASH: u32 = hash_str_32(IDENTIFIER_CLIENT); diff --git a/unshell-leaves/tcp_simple/src/server/mod.rs b/unshell-leaves/tcp_simple/src/server/mod.rs new file mode 100644 index 0000000..fd398dc --- /dev/null +++ b/unshell-leaves/tcp_simple/src/server/mod.rs @@ -0,0 +1,83 @@ +use std::{ + io, + net::{Ipv4Addr, TcpListener, ToSocketAddrs}, +}; + +use unshell::protocol::{Endpoint, Leaf}; + +use crate::transport::TcpBridge; + +/// TCP server-side transport leaf for one downstream endpoint. +/// +/// The protocol endpoint is intentionally leaf-owned by the caller, so this type +/// only bridges bytes: accepted TCP frames are deserialized into inbound packets, +/// and outbound packets queued for `child_endpoint_id` are serialized back onto the +/// same stream. Use this on the authority/parent side of a two-endpoint link. +#[derive(Debug)] +pub struct TCPServerLeaf { + listener: TcpListener, + bridge: TcpBridge, +} + +impl TCPServerLeaf { + /// Binds a nonblocking TCP listener for a child endpoint connection. + /// + /// `child_endpoint_id` must match the adjacent endpoint segment used in packet + /// paths. The server registers that endpoint as downstream so inbound bytes from + /// the child are treated as upward traffic by [`Endpoint::add_inbound_from`]. + pub fn new(listen_addr: A, child_endpoint_id: u32) -> io::Result + where + A: ToSocketAddrs, + { + let listener = TcpListener::bind(listen_addr)?; + listener.set_nonblocking(true)?; + + Ok(Self { + listener, + bridge: TcpBridge::new(child_endpoint_id, false), + }) + } + + /// Binds a nonblocking IPv4 listener for minimized fixed-address endpoints. + /// + /// This avoids making tiny binaries instantiate the fully generic public + /// constructor when they already know the concrete IPv4 address and port. + pub fn bind_ipv4(addr: Ipv4Addr, port: u16, child_endpoint_id: u32) -> io::Result { + let listener = TcpListener::bind((addr, port))?; + listener.set_nonblocking(true)?; + + Ok(Self { + listener, + bridge: TcpBridge::new(child_endpoint_id, false), + }) + } +} + +impl Leaf for TCPServerLeaf { + fn get_id(&self) -> u32 { + crate::IDENTIFIER_SERVER_HASH + } + + fn update(&mut self, endpoint: &mut Endpoint) { + self.bridge.register(endpoint); + self.accept_connection(); + self.bridge.update(endpoint); + } +} + +impl TCPServerLeaf { + /// Accepts at most one active stream without blocking the endpoint loop. + /// + /// A second accepted stream would make packet ownership ambiguous for the same + /// `child_endpoint_id`, so the minimal bridge keeps the first live connection and + /// waits for it to disconnect before accepting another. + fn accept_connection(&mut self) { + if self.bridge.is_connected() { + return; + } + + if let Ok((stream, _)) = self.listener.accept() { + let _ = self.bridge.set_stream(stream); + } + } +} diff --git a/unshell-leaves/tcp_simple/src/transport.rs b/unshell-leaves/tcp_simple/src/transport.rs new file mode 100644 index 0000000..a20283b --- /dev/null +++ b/unshell-leaves/tcp_simple/src/transport.rs @@ -0,0 +1,348 @@ +use std::{ + io::{self, Read, Write}, + net::TcpStream, +}; + +use unshell::protocol::{Endpoint, Packet}; + +#[cfg(target_os = "linux")] +const WOULD_BLOCK: i32 = 11; + +/// Returns whether `error` is the expected nonblocking-socket retry signal. +/// +/// Linux minimized endpoints use the raw `EAGAIN`/`EWOULDBLOCK` value to avoid +/// linking the broader `ErrorKind` classification path. Other targets keep the +/// portable standard-library classification because their raw values differ. +#[inline(always)] +fn is_would_block(error: &io::Error) -> bool { + #[cfg(target_os = "linux")] + { + error.raw_os_error() == Some(WOULD_BLOCK) + } + + #[cfg(not(target_os = "linux"))] + { + error.kind() == io::ErrorKind::WouldBlock + } +} + +/// Shared packet-to-TCP bridge used by the server and client leaves. +/// +/// TCP is a byte stream, while the protocol serializer emits one self-delimiting +/// packet frame at a time. This helper keeps just enough buffering to rebuild full +/// frames from arbitrary reads, route them through the endpoint, and preserve +/// partially written outbound bytes across nonblocking update ticks. +#[derive(Debug)] +pub(crate) struct TcpBridge { + remote_id: u32, + is_authority: bool, + stream: Option, + read_buffer: Vec, + write_buffer: Vec, + registered: bool, +} + +impl TcpBridge { + /// Creates bridge state for one adjacent endpoint. + /// + /// `is_authority` is passed directly to [`Endpoint::add_connection`]. Use `true` + /// when the remote endpoint is the parent/authority and `false` when it is a + /// child, matching the endpoint routing contract. + pub(crate) fn new(remote_id: u32, is_authority: bool) -> Self { + Self { + remote_id, + is_authority, + stream: None, + read_buffer: Vec::new(), + write_buffer: Vec::new(), + registered: false, + } + } + + /// Registers the transport edge once so endpoint routing accepts this peer. + pub(crate) fn register(&mut self, endpoint: &mut Endpoint) { + if !self.registered { + endpoint.add_connection(self.remote_id, self.is_authority); + self.registered = true; + } + } + + /// Returns whether there is an active TCP stream for this bridge. + pub(crate) fn is_connected(&self) -> bool { + self.stream.is_some() + } + + /// Installs a newly connected stream and makes it nonblocking for update loops. + /// + /// Stale buffers are cleared before replacing the socket because a partial packet + /// from an old TCP stream cannot be resumed safely on a new stream. TCP only gives + /// byte ordering inside one connection, not across reconnects. + pub(crate) fn set_stream(&mut self, stream: TcpStream) -> io::Result<()> { + stream.set_nonblocking(true)?; + self.read_buffer.clear(); + self.write_buffer.clear(); + self.stream = Some(stream); + Ok(()) + } + + /// Moves all currently available TCP frames into the endpoint and flushes queued output. + #[inline(never)] + pub(crate) fn update(&mut self, endpoint: &mut Endpoint) { + self.read_available(); + self.route_complete_frames(endpoint); + + if self.stream.is_none() { + return; + } + + self.collect_outbound(endpoint); + self.flush_pending(); + } + + /// Reads until the nonblocking stream would block or disconnects. + fn read_available(&mut self) { + let Some(stream) = self.stream.as_mut() else { + return; + }; + + let mut chunk = [0u8; 1024]; + + loop { + match stream.read(&mut chunk) { + Ok(0) => { + self.disconnect(); + break; + } + Ok(read) => self.read_buffer.extend_from_slice(&chunk[..read]), + Err(error) if is_would_block(&error) => break, + Err(_) => { + self.disconnect(); + break; + } + } + } + } + + /// Routes each complete serialized packet frame currently buffered from TCP. + fn route_complete_frames(&mut self, endpoint: &mut Endpoint) { + while let Some(frame_len) = next_frame_len(&self.read_buffer) { + // Transport input is untrusted. Bad frames and route failures are dropped + // so a peer cannot wedge the bridge with one malformed packet. + if let Ok(packet) = Packet::deserialize(&self.read_buffer[..frame_len]) { + let _ = endpoint.add_inbound_from(self.remote_id, packet); + } + + // `Packet::deserialize` owns the decoded path/data, so the byte frame can + // be discarded after routing without allocating a second temporary buffer. + self.read_buffer.copy_within(frame_len.., 0); + self.read_buffer + .truncate(self.read_buffer.len() - frame_len); + } + } + + /// Serializes endpoint packets queued for this remote into the pending write buffer. + fn collect_outbound(&mut self, endpoint: &mut Endpoint) { + let Some(queue) = endpoint.take_outbound_queue(self.remote_id) else { + return; + }; + + for packet in queue { + let _ = packet.serialize_into(&mut self.write_buffer); + } + } + + /// Writes pending bytes without blocking the endpoint loop. + fn flush_pending(&mut self) { + while !self.write_buffer.is_empty() { + let Some(stream) = self.stream.as_mut() else { + return; + }; + + match stream.write(&self.write_buffer) { + Ok(0) => { + self.disconnect(); + return; + } + Ok(written) => { + self.write_buffer.copy_within(written.., 0); + self.write_buffer + .truncate(self.write_buffer.len() - written); + } + Err(error) if is_would_block(&error) => return, + Err(_) => { + self.disconnect(); + return; + } + } + } + } + + /// Drops socket-local state; routing registration remains the intended topology. + fn disconnect(&mut self) { + self.stream = None; + self.read_buffer.clear(); + self.write_buffer.clear(); + } +} + +/// Returns the byte length of the next complete serialized packet in `buf`. +/// +/// The packet format has no outer TCP length prefix, so the bridge derives the frame +/// boundary from `path_len` and `body_len`. `None` means either more bytes are needed +/// or the advertised lengths overflowed; in both cases the safest small transport +/// behavior is to wait rather than guess at packet boundaries. +fn next_frame_len(buf: &[u8]) -> Option { + if buf.len() < 8 { + return None; + } + + let path_len = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]) as usize; + let path_bytes = path_len.checked_mul(4)?; + let body_len_offset = 8usize.checked_add(path_bytes)?; + + if buf.len() < body_len_offset.checked_add(4)? { + return None; + } + + let body_len = u32::from_le_bytes([ + buf[body_len_offset], + buf[body_len_offset + 1], + buf[body_len_offset + 2], + buf[body_len_offset + 3], + ]) as usize; + + let frame_len = body_len_offset.checked_add(4)?.checked_add(body_len)?; + + (buf.len() >= frame_len).then_some(frame_len) +} + +#[cfg(test)] +mod tests { + use std::{ + io::{Read, Write}, + net::{TcpListener, TcpStream}, + time::Duration, + }; + + use unshell::protocol::{Endpoint, Packet}; + + use super::{TcpBridge, next_frame_len}; + + const PARENT: u32 = 0x1000_0001; + const CHILD: u32 = 0x1000_0002; + const PROCEDURE: u32 = 0x2000_0001; + + /// Builds the parent side of the two-node topology used by bridge tests. + /// + /// The real endpoint constructor intentionally starts with an empty path so callers + /// can attach it anywhere in the tree. Transport tests set the path explicitly to + /// exercise the same routing contract production callers must satisfy. + fn parent_endpoint() -> Endpoint { + let mut endpoint = Endpoint::new(PARENT); + endpoint.path = vec![PARENT]; + endpoint + } + + /// Creates a local TCP pair without depending on a fixed port. + fn connected_pair() -> (TcpStream, TcpStream) { + let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap(); + let addr = listener.local_addr().unwrap(); + let client = TcpStream::connect(addr).unwrap(); + let (server, _) = listener.accept().unwrap(); + + client + .set_read_timeout(Some(Duration::from_secs(1))) + .unwrap(); + client + .set_write_timeout(Some(Duration::from_secs(1))) + .unwrap(); + + (server, client) + } + + /// Reads exactly one serialized packet frame from a blocking test stream. + fn read_frame(stream: &mut TcpStream) -> Vec { + let mut frame = Vec::new(); + let mut chunk = [0u8; 64]; + + loop { + let read = stream.read(&mut chunk).unwrap(); + assert_ne!(read, 0, "test TCP stream closed before a packet arrived"); + frame.extend_from_slice(&chunk[..read]); + + if let Some(frame_len) = next_frame_len(&frame) { + assert_eq!(frame_len, frame.len()); + return frame; + } + } + } + + /// Creates a downward packet that paves a return hook from parent to child. + fn downward_packet(hook_id: u16) -> Packet { + Packet { + hook_id, + end_hook: false, + path: vec![PARENT, CHILD], + procedure_id: PROCEDURE, + data: vec![1, 2, 3], + } + } + + #[test] + fn update_keeps_outbound_queued_until_connected() { + let mut endpoint = parent_endpoint(); + let mut bridge = TcpBridge::new(CHILD, false); + bridge.register(&mut endpoint); + + endpoint.add_outbound(downward_packet(7)).unwrap(); + bridge.update(&mut endpoint); + + let mut queued = 0usize; + endpoint.take_outbound_clear(CHILD, |_| queued += 1); + + assert_eq!(queued, 1); + } + + #[test] + fn bridge_writes_outbound_and_routes_inbound_reply() { + let mut endpoint = parent_endpoint(); + let mut bridge = TcpBridge::new(CHILD, false); + let (server, mut client) = connected_pair(); + bridge.register(&mut endpoint); + bridge.set_stream(server).unwrap(); + + endpoint.add_outbound(downward_packet(9)).unwrap(); + bridge.update(&mut endpoint); + + let sent = Packet::deserialize(&read_frame(&mut client)).unwrap(); + assert_eq!(sent.hook_id, 9); + assert_eq!(sent.path, vec![PARENT, CHILD]); + assert_eq!(sent.data, vec![1, 2, 3]); + + let reply = Packet { + hook_id: 9, + end_hook: true, + path: vec![PARENT], + procedure_id: PROCEDURE, + data: vec![4, 5, 6], + }; + client.write_all(&reply.serialize().unwrap()).unwrap(); + bridge.update(&mut endpoint); + + let mut received = Vec::new(); + endpoint.take_inbound_clear(PARENT, |packet| received.push(packet.clone())); + + assert_eq!(received.len(), 1); + assert_eq!(received[0].hook_id, 9); + assert_eq!(received[0].path, vec![PARENT]); + assert_eq!(received[0].data, vec![4, 5, 6]); + } + + #[test] + fn frame_length_waits_for_complete_packet() { + let frame = downward_packet(3).serialize().unwrap(); + + assert_eq!(next_frame_len(&frame[..frame.len() - 1]), None); + assert_eq!(next_frame_len(&frame), Some(frame.len())); + } +}