From fc82f4f921862a1a47d235aff4f1775805fdc2b6 Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Thu, 28 May 2026 18:17:01 -0600 Subject: [PATCH] Make macro system and PTY test leaf --- Cargo.lock | 24 + Cargo.toml | 22 +- src/lib.rs | 57 +-- unshell-leaves/leaf-pty/Cargo.toml | 20 + unshell-leaves/leaf-pty/src/codec.rs | 106 +++++ unshell-leaves/leaf-pty/src/constants.rs | 35 ++ unshell-leaves/leaf-pty/src/lib.rs | 26 ++ unshell-leaves/leaf-pty/src/session.rs | 116 +++++ unshell-leaves/leaf-pty/src/state.rs | 37 ++ unshell-leaves/leaf-pty/src/tests.rs | 393 ++++++++++++++++ unshell-leaves/src/lib.rs | 1 - unshell-macros-core/Cargo.toml | 22 + unshell-macros-core/src/leaf/args.rs | 78 ++++ unshell-macros-core/src/leaf/generator.rs | 434 ++++++++++++++++++ unshell-macros-core/src/leaf/mod.rs | 76 +++ unshell-macros-core/src/leaf/names.rs | 58 +++ unshell-macros-core/src/lib.rs | 9 + {unshell-leaves => unshell-macros}/Cargo.toml | 17 +- unshell-macros/src/lib.rs | 15 + unshell-protocol/src/endpoint/mod.rs | 31 ++ unshell-protocol/src/leaf.rs | 359 +++++++++++++++ unshell-protocol/src/lib.rs | 14 +- unshell-protocol/src/packet.rs | 2 +- 23 files changed, 1866 insertions(+), 86 deletions(-) create mode 100644 unshell-leaves/leaf-pty/Cargo.toml create mode 100644 unshell-leaves/leaf-pty/src/codec.rs create mode 100644 unshell-leaves/leaf-pty/src/constants.rs create mode 100644 unshell-leaves/leaf-pty/src/lib.rs create mode 100644 unshell-leaves/leaf-pty/src/session.rs create mode 100644 unshell-leaves/leaf-pty/src/state.rs create mode 100644 unshell-leaves/leaf-pty/src/tests.rs delete mode 100644 unshell-leaves/src/lib.rs create mode 100644 unshell-macros-core/Cargo.toml create mode 100644 unshell-macros-core/src/leaf/args.rs create mode 100644 unshell-macros-core/src/leaf/generator.rs create mode 100644 unshell-macros-core/src/leaf/mod.rs create mode 100644 unshell-macros-core/src/leaf/names.rs create mode 100644 unshell-macros-core/src/lib.rs rename {unshell-leaves => unshell-macros}/Cargo.toml (57%) create mode 100644 unshell-macros/src/lib.rs create mode 100644 unshell-protocol/src/leaf.rs diff --git a/Cargo.lock b/Cargo.lock index eef28b9..207f8b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -399,6 +399,13 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "leaf-pty" +version = "0.1.0" +dependencies = [ + "unshell", +] + [[package]] name = "leb128fmt" version = "0.1.0" @@ -847,9 +854,26 @@ dependencies = [ "rkyv", "static_init", "thiserror", + "unshell-macros", "unshell-protocol", ] +[[package]] +name = "unshell-macros" +version = "0.1.0" +dependencies = [ + "unshell-macros-core", +] + +[[package]] +name = "unshell-macros-core" +version = "0.1.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "unshell-protocol" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 1d73213..7c0ebd8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,11 +4,12 @@ cargo-features = ["trim-paths", "panic-immediate-abort"] members = [ "ush-obfuscate", "base62", - # "unshell-macros", + "unshell-macros-core", + "unshell-macros", + "unshell-protocol", - # "unshell-runtime", - # "unshell-leaves", - # "treetest", + + "unshell-leaves/leaf-pty", ] resolver = "2" @@ -30,15 +31,19 @@ quote = "1.0.45" proc-macro2 = "1.0.106" portable-pty = "0.9.0" crossbeam-channel = "0.5.15" + unshell = { path = "." } unshell-protocol = { path = "./unshell-protocol" } -# unshell-runtime = { path = "./unshell-runtime" } -# unshell-leaves = { path = "./unshell-leaves" } -# unshell-macros = { path = "./unshell-macros" } +unshell-macros-core = { path = "./unshell-macros-core" } +unshell-macros = { path = "./unshell-macros" } # ush-obfuscate = { path = "./ush-obfuscate" } # base62 = { path = "./base62" } +# Leaves +leaf-pty = { path = "./unshell-leaves/leaf-pty" } + + [package] name = "unshell" version.workspace = true @@ -63,7 +68,8 @@ thiserror = { workspace = true, optional = true } chrono = { workspace = true, optional = true } # ush-obfuscate = { workspace = true } static_init = { workspace = true } -# unshell-macros = { workspace = true } + +unshell-macros = { workspace = true } unshell-protocol = { workspace = true } # unshell-runtime = { workspace = true } # unshell-leaves = { workspace = true } diff --git a/src/lib.rs b/src/lib.rs index 3f7a827..365f64c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,62 +12,11 @@ #![no_std] pub extern crate alloc; -// Re-export derive macros against a stable `::unshell` path, including when the -// macros are used inside this crate's own examples and tests. -#[allow(unused_extern_crates)] -extern crate self as unshell; pub mod logger; -/// Re-export the protocol crate behind the historical `unshell::protocol` path so -/// proc-macro output and downstream code do not need a second migration. -pub use unshell_protocol as protocol; +pub mod protocol { + pub use unshell_protocol::*; -// Re-export the leaf library crate behind the historical `unshell::leaves` path -// once the leaf crate is part of the active workspace again. -// pub use unshell_leaves as leaves; - -// Re-export the runtime crate behind the `unshell::runtime` path once the runtime -// crate is part of the active workspace again. -// pub use unshell_runtime as runtime; - -// pub use unshell_macros::{Procedure, leaf, procedures}; - -/// Creates a root-assumed endpoint from one local identifier plus any number of leaf hosts. -/// -/// What it is: a convenience macro that builds a `ProtocolEndpoint` whose protocol path starts at -/// root, with no parent or children, and whose leaf inventory is inferred from the supplied host -/// values. -/// -/// Why it exists: the common bootstrap case should not require callers to manually construct an -/// empty path, `Vec`, and a `Vec` when they already have leaf host values. -/// -/// # Example -/// -/// ```rust,ignore -/// use unshell::{create_endpoint, leaf}; -/// use unshell::protocol::tree::Endpoint; -/// -/// #[derive(Default)] -/// struct DemoLeaf; -/// -/// #[leaf(id = "org.example.v1.demo", procedures = ["ping"], endpoint_struct = DemoLeaf)] -/// struct Demo; -/// -/// let endpoint = create_endpoint!("demo", DemoLeaf::default()); -/// assert!(endpoint.path().is_empty()); -/// assert_eq!(endpoint.local_id(), Some("demo")); -/// ``` -#[macro_export] -macro_rules! create_endpoint { - ($id:expr $(, $leaf:expr )* $(,)?) => {{ - let mut __unshell_leaf_specs = ::unshell::alloc::vec::Vec::new(); - $( - let __unshell_leaf = $leaf; - __unshell_leaf_specs.push(::unshell::protocol::tree::leaf_spec_of(&__unshell_leaf)); - )* - ::unshell::protocol::tree::ProtocolEndpoint::root($id, __unshell_leaf_specs) - }}; + pub use unshell_macros::unshell_leaf; } - -// pub use ush_obfuscate as obfuscate; diff --git a/unshell-leaves/leaf-pty/Cargo.toml b/unshell-leaves/leaf-pty/Cargo.toml new file mode 100644 index 0000000..e16c044 --- /dev/null +++ b/unshell-leaves/leaf-pty/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "leaf-pty" +version.workspace = true +edition.workspace = true +description = "Hook-backed PTY leaf implementation for UnShell" + +[dependencies] +unshell = { workspace = true } + +[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-pty/src/codec.rs b/unshell-leaves/leaf-pty/src/codec.rs new file mode 100644 index 0000000..f421385 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/codec.rs @@ -0,0 +1,106 @@ +use alloc::vec::Vec; + +use unshell::protocol::{HookID, Packet}; + +use crate::{OP_ERROR, OP_OPEN, PROC_PTY}; + +/// Encodes a tiny PTY frame into `Packet::data`. +pub fn encode_frame(opcode: u8, payload: &[u8]) -> Vec { + let mut data = Vec::with_capacity(1 + payload.len()); + data.push(opcode); + data.extend_from_slice(payload); + 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) +} + +/// Returns the opcode byte from a PTY packet, if present. +pub fn frame_opcode(packet: &Packet) -> Option { + packet.data.first().copied() +} + +/// Returns the frame payload after the opcode byte. +pub fn frame_payload(packet: &Packet) -> &[u8] { + if packet.data.len() > 1 { + &packet.data[1..] + } else { + &[] + } +} + +/// Builds an outer PTY packet for callers and tests. +pub fn pty_packet( + path: Vec, + hook_id: HookID, + end_hook: bool, + opcode: u8, + payload: &[u8], +) -> Packet { + Packet { + hook_id, + end_hook, + path, + procedure_id: PROC_PTY, + data: encode_frame(opcode, payload), + } +} + +/// 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 { + 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() + } +} diff --git a/unshell-leaves/leaf-pty/src/constants.rs b/unshell-leaves/leaf-pty/src/constants.rs new file mode 100644 index 0000000..a73be52 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/constants.rs @@ -0,0 +1,35 @@ +/// Leaf id used by the generated fake PTY wrapper. +pub const LEAF_FAKE_PTY: u32 = 300; + +/// Outer procedure id used by all fake PTY session packets. +pub const PROC_PTY: u32 = 30; + +/// Downward opcode that opens one PTY session. +pub const OP_OPEN: u8 = 0; + +/// Upward opcode acknowledging an opened PTY session. +pub const OP_OPENED: u8 = 1; + +/// Downward opcode carrying PTY stdin bytes. +pub const OP_INPUT: u8 = 2; + +/// Downward opcode representing terminal resize. +pub const OP_RESIZE: u8 = 3; + +/// Downward opcode closing PTY stdin without closing the session hook. +pub const OP_STDIN_EOF: u8 = 4; + +/// Downward opcode asking the remote process to terminate gracefully. +pub const OP_TERMINATE: u8 = 5; + +/// Downward opcode aborting the session without an acknowledgement. +pub const OP_ABORT: u8 = 6; + +/// Upward opcode carrying PTY stdout/stderr bytes. +pub const OP_OUTPUT: u8 = 7; + +/// Upward final opcode carrying the process exit status. +pub const OP_EXIT: u8 = 8; + +/// Upward final opcode carrying a fatal PTY protocol error. +pub const OP_ERROR: u8 = 9; diff --git a/unshell-leaves/leaf-pty/src/lib.rs b/unshell-leaves/leaf-pty/src/lib.rs new file mode 100644 index 0000000..ab00097 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/lib.rs @@ -0,0 +1,26 @@ +//! PTY leaf support for UnShell. +//! +//! This crate currently contains a deterministic fake PTY session used to prove the +//! macro-generated leaf shape. The fake leaf exercises the same hook-backed protocol +//! invariants as a real PTY worker without pulling OS-specific PTY code into +//! `unshell-protocol`. + +#![no_std] + +extern crate alloc; + +mod codec; +mod constants; +mod session; +mod state; + +pub use codec::{ + decode_open_reply_path, encode_frame, encode_open, frame_opcode, frame_payload, + pty_open_packet, pty_packet, +}; +pub use constants::*; +pub use session::{PtySession, PtySessionState}; +pub use state::{FakePtyLeaf, FakePtyState}; + +#[cfg(test)] +mod tests; diff --git a/unshell-leaves/leaf-pty/src/session.rs b/unshell-leaves/leaf-pty/src/session.rs new file mode 100644 index 0000000..eccd085 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/session.rs @@ -0,0 +1,116 @@ +use alloc::vec::Vec; + +use unshell::protocol::{ + HookID, Packet, PacketQueue, Session, SessionCtx, SessionInit, SessionInitResult, SessionStatus, +}; + +use crate::{ + codec::{ + decode_open_reply_path, error_packet, frame_opcode, frame_payload, + reply_path_from_destination, + }, + constants::{ + OP_ABORT, OP_ERROR, OP_EXIT, OP_INPUT, OP_OPEN, OP_OPENED, OP_OUTPUT, OP_STDIN_EOF, + OP_TERMINATE, PROC_PTY, + }, + 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. +pub struct PtySessionState { + hook_id: HookID, + reply_path: Vec, + opened_pending: bool, + stdin_closed: bool, +} + +impl Session for PtySession { + 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 { + if frame_opcode(&packet) != Some(OP_OPEN) { + return SessionInitResult::RejectedWith(error_packet( + ctx.hook_id(), + reply_path_from_destination(ctx.packet_path()), + 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, + opened_pending: true, + stdin_closed: false, + }) + } + + fn update( + leaf: &mut FakePtyState, + session: &mut Self::State, + incoming: &mut PacketQueue, + ctx: &mut SessionCtx<'_>, + ) -> SessionStatus { + if session.opened_pending { + ctx.send(OP_OPENED, &[]); + 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_STDIN_EOF) => { + session.stdin_closed = true; + leaf.last_stdin_eof_hook = Some(session.hook_id); + } + Some(OP_TERMINATE) => { + ctx.send_final(OP_EXIT, &[0]); + close_session(leaf); + return SessionStatus::Closed; + } + Some(OP_ABORT) => { + close_session(leaf); + return SessionStatus::Closed; + } + Some(OP_OPEN) => { + ctx.send_final(OP_ERROR, b"duplicate-open"); + close_session(leaf); + return SessionStatus::Closed; + } + _ => { + ctx.send_final(OP_ERROR, b"unknown-opcode"); + close_session(leaf); + return SessionStatus::Closed; + } + } + } + + SessionStatus::Running + } +} + +/// Decrements the active-session counter exactly once for a terminal session path. +fn close_session(leaf: &mut FakePtyState) { + leaf.active_count = leaf.active_count.saturating_sub(1); +} diff --git a/unshell-leaves/leaf-pty/src/state.rs b/unshell-leaves/leaf-pty/src/state.rs new file mode 100644 index 0000000..eaf771b --- /dev/null +++ b/unshell-leaves/leaf-pty/src/state.rs @@ -0,0 +1,37 @@ +use unshell::protocol::{HookID, unshell_leaf}; + +use crate::{constants::LEAF_FAKE_PTY, session::PtySession}; + +/// User-owned state for the generated fake PTY leaf. +/// +/// The macro-generated `FakePtyLeaf` wrapper stores sessions and retry queues around +/// this struct. Keeping counters here makes tests and future procedures observe leaf +/// behavior without reaching into generated session storage. +#[unshell_leaf(leaf = FakePtyLeaf, id = LEAF_FAKE_PTY, sessions(PtySession))] +pub struct FakePtyState { + /// Number of sessions that application logic considers active. + pub active_count: usize, + + /// Total number of successfully opened sessions. + pub total_opened: u64, + + /// Last hook that received stdin EOF. + pub last_stdin_eof_hook: Option, +} + +impl FakePtyState { + /// Creates a fake PTY state with no active sessions. + pub fn new() -> Self { + Self { + active_count: 0, + total_opened: 0, + last_stdin_eof_hook: None, + } + } +} + +impl Default for FakePtyState { + fn default() -> Self { + Self::new() + } +} diff --git a/unshell-leaves/leaf-pty/src/tests.rs b/unshell-leaves/leaf-pty/src/tests.rs new file mode 100644 index 0000000..513fcf0 --- /dev/null +++ b/unshell-leaves/leaf-pty/src/tests.rs @@ -0,0 +1,393 @@ +use alloc::{vec, vec::Vec}; + +use unshell::protocol::{Endpoint, Leaf, Packet}; + +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, +}; + +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 + }) +} + +#[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_frame(&packets[0], hook_id, OP_OPENED, false, &[]); +} + +#[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_retries_without_losing_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, + ); + 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); + + 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_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); +} + +#[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.connections.insert((ENDPOINT_A, true)); + + endpoint + .add_inbound_from( + ENDPOINT_A, + pty_open_packet(vec![ENDPOINT_A, ENDPOINT_B], 7, &[ENDPOINT_A]), + ) + .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/src/lib.rs b/unshell-leaves/src/lib.rs deleted file mode 100644 index 8b13789..0000000 --- a/unshell-leaves/src/lib.rs +++ /dev/null @@ -1 +0,0 @@ - diff --git a/unshell-macros-core/Cargo.toml b/unshell-macros-core/Cargo.toml new file mode 100644 index 0000000..0bd9c20 --- /dev/null +++ b/unshell-macros-core/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "unshell-macros-core" +version.workspace = true +edition.workspace = true +description = "Parser and code generator for UnShell procedural macros" + +[dependencies] +proc-macro2 = { workspace = true } +quote = { workspace = true } +syn = { workspace = true, features = ["full", "extra-traits"] } + +[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-macros-core/src/leaf/args.rs b/unshell-macros-core/src/leaf/args.rs new file mode 100644 index 0000000..357644f --- /dev/null +++ b/unshell-macros-core/src/leaf/args.rs @@ -0,0 +1,78 @@ +use syn::{ + Expr, Ident, Result, Token, Type, + parse::{Parse, ParseStream}, +}; + +/// Parsed arguments from `#[unshell_leaf(...)]`. +#[derive(Debug)] +pub(crate) struct UnshellLeafArgs { + pub(crate) leaf: Ident, + pub(crate) id: Expr, + pub(crate) sessions: Vec, + pub(crate) procedures: Vec, +} + +impl Parse for UnshellLeafArgs { + fn parse(input: ParseStream<'_>) -> Result { + let mut leaf = None; + let mut id = None; + let mut sessions = Vec::new(); + let mut procedures = Vec::new(); + + while !input.is_empty() { + let key: Ident = input.parse()?; + match key.to_string().as_str() { + "leaf" => { + reject_duplicate(&leaf, &key)?; + input.parse::()?; + leaf = Some(input.parse()?); + } + "id" => { + reject_duplicate(&id, &key)?; + input.parse::()?; + id = Some(input.parse()?); + } + "sessions" => { + sessions = parse_type_list(input)?; + } + "procedures" => { + procedures = parse_type_list(input)?; + } + _ => { + return Err(syn::Error::new( + key.span(), + "expected `leaf`, `id`, `sessions`, or `procedures`", + )); + } + } + + if input.peek(Token![,]) { + input.parse::()?; + } + } + + Ok(Self { + leaf: leaf.ok_or_else(|| input.error("missing `leaf = WrapperName`"))?, + id: id.ok_or_else(|| input.error("missing `id = LEAF_ID`"))?, + sessions, + procedures, + }) + } +} + +/// Rejects repeated scalar keys while keeping repeated list keys additive by design. +fn reject_duplicate(slot: &Option, key: &Ident) -> Result<()> { + if slot.is_some() { + Err(syn::Error::new(key.span(), "duplicate key")) + } else { + Ok(()) + } +} + +/// Parses `name(Type, Type)` argument payloads. +fn parse_type_list(input: ParseStream<'_>) -> Result> { + let content; + syn::parenthesized!(content in input); + let parsed = content.parse_terminated(Type::parse, Token![,])?; + Ok(parsed.into_iter().collect()) +} diff --git a/unshell-macros-core/src/leaf/generator.rs b/unshell-macros-core/src/leaf/generator.rs new file mode 100644 index 0000000..d237e55 --- /dev/null +++ b/unshell-macros-core/src/leaf/generator.rs @@ -0,0 +1,434 @@ +use proc_macro2::TokenStream; +use quote::{format_ident, quote}; +use syn::{Ident, ItemStruct, Result, Type}; + +use super::{ + UnshellLeafArgs, + names::{last_type_ident, to_snake_case}, +}; + +/// Code generator state for one `#[unshell_leaf]` expansion. +pub(crate) struct LeafGenerator { + args: UnshellLeafArgs, + state: ItemStruct, +} + +impl LeafGenerator { + /// Creates a generator for one parsed state struct. + pub(crate) fn new(args: UnshellLeafArgs, state: ItemStruct) -> Self { + Self { args, state } + } + + /// Emits the original state struct plus the generated wrapper leaf. + pub(crate) fn expand(self) -> Result { + let state = &self.state; + let state_ident = &state.ident; + let leaf_ident = &self.args.leaf; + let leaf_id = &self.args.id; + let vis = &state.vis; + let generics = &state.generics; + let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); + let state_type = quote!(#state_ident #ty_generics); + + let session_stores = self.session_stores()?; + let fields = self.store_fields(&session_stores, &state_type); + let initializers = self.store_initializers(&session_stores); + let packet_predicates = self.packet_predicates(&state_type); + let dispatch_arms = self.dispatch_arms(&session_stores, &state_type); + let session_updates = self.session_updates(&session_stores, &state_type); + let session_flushes = self.session_flushes(&session_stores); + let session_retains = self.session_retains(&session_stores); + let active_count_terms = self.active_count_terms(&session_stores); + let pending_count_terms = self.pending_count_terms(&session_stores); + let id_checks = self.id_checks(&state_type); + + Ok(quote! { + #state + + #vis struct #leaf_ident #generics #where_clause { + state: #state_type, + __unshell_procedure_outbox: ::unshell::protocol::PacketQueue, + #(#fields,)* + } + + impl #impl_generics #leaf_ident #ty_generics #where_clause { + const __UNSHELL_PROCEDURE_ID_CHECKS: () = { + #(#id_checks)* + }; + + /// Creates the generated leaf wrapper around user-owned state. + pub fn new(state: #state_type) -> Self { + Self { + state, + __unshell_procedure_outbox: ::unshell::protocol::PacketQueue::new(), + #(#initializers,)* + } + } + + /// Returns immutable access to the user-owned leaf state. + pub fn state(&self) -> &#state_type { + &self.state + } + + /// Returns mutable access to the user-owned leaf state. + pub fn state_mut(&mut self) -> &mut #state_type { + &mut self.state + } + + /// Returns the number of active session entries across all session families. + pub fn active_session_count(&self) -> usize { + 0usize #(+ #active_count_terms)* + } + + /// Returns queued inbound and outbound packets owned by this generated leaf. + pub fn pending_packet_count(&self) -> usize { + let mut __unshell_count = self.__unshell_procedure_outbox.len(); + #(#pending_count_terms)* + __unshell_count + } + + fn __unshell_packet_is_owned(packet: &::unshell::protocol::Packet) -> bool { + false #(|| #packet_predicates)* + } + + fn __unshell_dispatch( + &mut self, + endpoint: &mut ::unshell::protocol::Endpoint, + packet: ::unshell::protocol::Packet, + ) { + #(#dispatch_arms)* + } + + fn __unshell_update_sessions(&mut self) { + #(#session_updates)* + } + + fn __unshell_flush_all(&mut self, endpoint: &mut ::unshell::protocol::Endpoint) { + ::unshell::protocol::flush_packet_queue( + endpoint, + &mut self.__unshell_procedure_outbox, + ); + #(#session_flushes)* + #(#session_retains)* + } + + fn __unshell_parent_reply_path( + endpoint: &::unshell::protocol::Endpoint, + ) -> ::unshell::protocol::alloc::vec::Vec { + if endpoint.path.len() > 1 { + endpoint.path[..endpoint.path.len() - 1].to_vec() + } else { + endpoint.path.clone() + } + } + } + + impl #impl_generics ::unshell::protocol::Leaf for #leaf_ident #ty_generics #where_clause { + fn get_id(&self) -> u32 { + #leaf_id + } + + fn update(&mut self, endpoint: &mut ::unshell::protocol::Endpoint) { + self.__unshell_flush_all(endpoint); + + let Some(__unshell_local_id) = endpoint.path.last().copied() else { + return; + }; + + let mut __unshell_packets = ::unshell::protocol::alloc::vec::Vec::new(); + endpoint.take_inbound_matching( + __unshell_local_id, + Self::__unshell_packet_is_owned, + |packet| __unshell_packets.push(packet), + ); + + for __unshell_packet in __unshell_packets { + self.__unshell_dispatch(endpoint, __unshell_packet); + } + + self.__unshell_update_sessions(); + self.__unshell_flush_all(endpoint); + } + } + }) + } + + /// Computes one generated store name per session type. + fn session_stores(&self) -> Result> { + self.args + .sessions + .iter() + .map(|session| { + let suffix = last_type_ident(session)?; + let field_suffix = to_snake_case(&suffix.to_string()); + Ok(SessionStore { + ty: session.clone(), + field: format_ident!("__unshell_{}_sessions", field_suffix), + }) + }) + .collect() + } + + /// Emits wrapper fields for session stores. + fn store_fields(&self, stores: &[SessionStore], state_type: &TokenStream) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + let session_ty = &store.ty; + quote! { + #field: ::unshell::protocol::alloc::vec::Vec< + ::unshell::protocol::SessionEntry< + <#session_ty as ::unshell::protocol::Session<#state_type>>::State + > + > + } + }) + .collect() + } + + /// Emits constructor field initializers for session stores. + fn store_initializers(&self, stores: &[SessionStore]) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + quote!(#field: ::unshell::protocol::alloc::vec::Vec::new()) + }) + .collect() + } + + /// Emits boolean procedure-id ownership checks for the filtered endpoint drain. + fn packet_predicates(&self, state_type: &TokenStream) -> Vec { + let session_checks = self.args.sessions.iter().map(|session_ty| { + quote! { + packet.procedure_id + == <#session_ty as ::unshell::protocol::Session<#state_type>>::PROCEDURE_ID + } + }); + let procedure_checks = self.args.procedures.iter().map(|procedure_ty| { + quote! { + packet.procedure_id + == <#procedure_ty as ::unshell::protocol::Procedure<#state_type>>::PROCEDURE_ID + } + }); + + session_checks.chain(procedure_checks).collect() + } + + /// Emits static dispatch branches for every session and procedure type. + fn dispatch_arms(&self, stores: &[SessionStore], state_type: &TokenStream) -> Vec { + let mut arms = Vec::new(); + + for store in stores { + let field = &store.field; + let session_ty = &store.ty; + arms.push(quote! { + if packet.procedure_id + == <#session_ty as ::unshell::protocol::Session<#state_type>>::PROCEDURE_ID + { + if let Some(__unshell_entry) = self + .#field + .iter_mut() + .find(|entry| entry.hook_id == packet.hook_id) + { + __unshell_entry.inbox.push_back(packet); + } else { + let __unshell_hook_id = packet.hook_id; + let __unshell_packet_path = packet.path.clone(); + let mut __unshell_init = ::unshell::protocol::SessionInit::new( + __unshell_hook_id, + __unshell_packet_path, + ); + + match <#session_ty as ::unshell::protocol::Session<#state_type>>::init( + &mut self.state, + packet, + &mut __unshell_init, + ) { + ::unshell::protocol::SessionInitResult::Created(__unshell_state) => { + self.#field.push(::unshell::protocol::SessionEntry::new( + __unshell_hook_id, + __unshell_state, + )); + } + ::unshell::protocol::SessionInitResult::Rejected => {} + ::unshell::protocol::SessionInitResult::RejectedWith(__unshell_packet) => { + self.__unshell_procedure_outbox.push_back(__unshell_packet); + } + } + } + return; + } + }); + } + + for procedure_ty in &self.args.procedures { + arms.push(quote! { + if packet.procedure_id + == <#procedure_ty as ::unshell::protocol::Procedure<#state_type>>::PROCEDURE_ID + { + let mut __unshell_out = ::unshell::protocol::ProcedureOut::new( + packet.hook_id, + Self::__unshell_parent_reply_path(endpoint), + <#procedure_ty as ::unshell::protocol::Procedure<#state_type>>::PROCEDURE_ID, + ); + <#procedure_ty as ::unshell::protocol::Procedure<#state_type>>::handle( + &mut self.state, + endpoint, + packet, + &mut __unshell_out, + ); + self.__unshell_procedure_outbox.extend(__unshell_out.into_packets()); + return; + } + }); + } + + arms + } + + /// Emits the per-session update loop for every session family. + fn session_updates( + &self, + stores: &[SessionStore], + state_type: &TokenStream, + ) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + let session_ty = &store.ty; + quote! { + for __unshell_entry in &mut self.#field { + if __unshell_entry.closed { + continue; + } + + let __unshell_reply_path = + <#session_ty as ::unshell::protocol::Session<#state_type>>::reply_path( + &__unshell_entry.state, + ) + .to_vec(); + let mut __unshell_ctx = ::unshell::protocol::SessionCtx::new( + __unshell_entry.hook_id, + __unshell_reply_path, + <#session_ty as ::unshell::protocol::Session<#state_type>>::PROCEDURE_ID, + &mut __unshell_entry.outbox, + ); + let __unshell_status = + <#session_ty as ::unshell::protocol::Session<#state_type>>::update( + &mut self.state, + &mut __unshell_entry.state, + &mut __unshell_entry.inbox, + &mut __unshell_ctx, + ); + + if ::core::matches!( + __unshell_status, + ::unshell::protocol::SessionStatus::Closed + ) { + __unshell_entry.closed = true; + } + } + } + }) + .collect() + } + + /// Emits retry flushing for every session outbox. + fn session_flushes(&self, stores: &[SessionStore]) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + quote! { + for __unshell_entry in &mut self.#field { + ::unshell::protocol::flush_packet_queue( + endpoint, + &mut __unshell_entry.outbox, + ); + } + } + }) + .collect() + } + + /// Emits removal of closed sessions whose final packets have routed. + fn session_retains(&self, stores: &[SessionStore]) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + quote! { + self.#field + .retain(|entry| !entry.closed || !entry.outbox.is_empty()); + } + }) + .collect() + } + + /// Emits additive terms for active session counts. + fn active_count_terms(&self, stores: &[SessionStore]) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + quote!(self.#field.len()) + }) + .collect() + } + + /// Emits statements that accumulate pending packet counts. + fn pending_count_terms(&self, stores: &[SessionStore]) -> Vec { + stores + .iter() + .map(|store| { + let field = &store.field; + quote! { + for __unshell_entry in &self.#field { + __unshell_count += + __unshell_entry.inbox.len() + __unshell_entry.outbox.len(); + } + } + }) + .collect() + } + + /// Emits pairwise const assertions for procedure-id uniqueness. + fn id_checks(&self, state_type: &TokenStream) -> Vec { + let mut ids = Vec::new(); + for session_ty in &self.args.sessions { + ids.push( + quote!(<#session_ty as ::unshell::protocol::Session<#state_type>>::PROCEDURE_ID), + ); + } + for procedure_ty in &self.args.procedures { + ids.push( + quote!(<#procedure_ty as ::unshell::protocol::Procedure<#state_type>>::PROCEDURE_ID), + ); + } + + let mut checks = Vec::new(); + for left in 0..ids.len() { + for right in (left + 1)..ids.len() { + let left_id = &ids[left]; + let right_id = &ids[right]; + checks.push(quote! { + assert!( + #left_id != #right_id, + "duplicate unshell procedure id in #[unshell_leaf]" + ); + }); + } + } + + checks + } +} + +/// Generated storage metadata for one session family. +struct SessionStore { + ty: Type, + field: Ident, +} diff --git a/unshell-macros-core/src/leaf/mod.rs b/unshell-macros-core/src/leaf/mod.rs new file mode 100644 index 0000000..9f267ea --- /dev/null +++ b/unshell-macros-core/src/leaf/mod.rs @@ -0,0 +1,76 @@ +//! Leaf wrapper macro implementation. +//! +//! Everything in this module is specific to `#[unshell_leaf]`: argument parsing, +//! generated wrapper storage, static dispatch, and retry-safe session output. Future +//! macro families should be added as sibling modules instead of sharing this internal +//! structure. + +mod args; +mod generator; +mod names; + +use proc_macro2::TokenStream; +use syn::{ItemStruct, Result, parse2}; + +pub(crate) use args::UnshellLeafArgs; +pub(crate) use generator::LeafGenerator; + +/// Expands `#[unshell_leaf(...)]` into a wrapper leaf and `Leaf` implementation. +/// +/// Errors are returned as tokenized `compile_error!` output so the proc-macro shim can +/// stay a thin transport layer from compiler tokens to this core implementation. +pub fn expand_unshell_leaf(attr: TokenStream, item: TokenStream) -> TokenStream { + match expand_unshell_leaf_result(attr, item) { + Ok(tokens) => tokens, + Err(error) => error.to_compile_error(), + } +} + +/// Fallible expansion path used by unit tests. +pub fn expand_unshell_leaf_result(attr: TokenStream, item: TokenStream) -> Result { + let args = parse2::(attr)?; + let state = parse2::(item)?; + LeafGenerator::new(args, state).expand() +} + +#[cfg(test)] +mod tests { + use super::*; + use quote::quote; + + #[test] + fn parses_leaf_arguments() { + let args = parse2::(quote! { + leaf = DemoLeaf, + id = 42, + sessions(DemoSession), + procedures(PingProcedure) + }) + .unwrap(); + + assert_eq!(args.leaf, "DemoLeaf"); + assert_eq!(args.sessions.len(), 1); + assert_eq!(args.procedures.len(), 1); + } + + #[test] + fn missing_leaf_is_rejected() { + let error = parse2::(quote! { id = 42 }).unwrap_err(); + + assert!(error.to_string().contains("missing `leaf")); + } + + #[test] + fn expansion_contains_static_dispatch() { + let expanded = expand_unshell_leaf_result( + quote! { leaf = DemoLeaf, id = 9, sessions(DemoSession) }, + quote! { pub struct DemoState; }, + ) + .unwrap() + .to_string(); + + assert!(expanded.contains("struct DemoLeaf")); + assert!(expanded.contains("impl :: unshell :: protocol :: Leaf for DemoLeaf")); + assert!(expanded.contains("DemoSession")); + } +} diff --git a/unshell-macros-core/src/leaf/names.rs b/unshell-macros-core/src/leaf/names.rs new file mode 100644 index 0000000..d242fb9 --- /dev/null +++ b/unshell-macros-core/src/leaf/names.rs @@ -0,0 +1,58 @@ +use syn::{Ident, Result, Type}; + +/// Returns the final path segment for a session type. +pub(crate) fn last_type_ident(ty: &Type) -> Result { + let Type::Path(path) = ty else { + return Err(syn::Error::new_spanned( + ty, + "session types must be named paths", + )); + }; + let Some(segment) = path.path.segments.last() else { + return Err(syn::Error::new_spanned(ty, "session type path is empty")); + }; + + Ok(segment.ident.clone()) +} + +/// Converts a Rust type name into a snake-case fragment for generated private fields. +pub(crate) fn to_snake_case(name: &str) -> String { + let mut output = String::with_capacity(name.len()); + let chars: Vec = name.chars().collect(); + + for (index, character) in chars.iter().copied().enumerate() { + if character.is_ascii_uppercase() { + let previous = index + .checked_sub(1) + .and_then(|previous| chars.get(previous)); + let next = chars.get(index + 1); + let previous_needs_boundary = previous + .map(|previous| previous.is_ascii_lowercase() || previous.is_ascii_digit()) + .unwrap_or(false); + let acronym_needs_boundary = previous + .map(|previous| previous.is_ascii_uppercase()) + .unwrap_or(false) + && next.map(|next| next.is_ascii_lowercase()).unwrap_or(false); + + if previous_needs_boundary || acronym_needs_boundary { + output.push('_'); + } + output.push(character.to_ascii_lowercase()); + } else { + output.push(character); + } + } + + output +} + +#[cfg(test)] +mod tests { + use super::to_snake_case; + + #[test] + fn session_store_fields_are_snake_case() { + assert_eq!(to_snake_case("PtySession"), "pty_session"); + assert_eq!(to_snake_case("HTTPServer"), "http_server"); + } +} diff --git a/unshell-macros-core/src/lib.rs b/unshell-macros-core/src/lib.rs new file mode 100644 index 0000000..f7b8f9a --- /dev/null +++ b/unshell-macros-core/src/lib.rs @@ -0,0 +1,9 @@ +//! Parser and code generator for UnShell procedural macros. +//! +//! This crate is intentionally not a proc-macro crate. Keeping each macro family's +//! parser and code generator here makes them unit-testable and prevents parsing +//! dependencies from leaking into runtime crates. + +mod leaf; + +pub use leaf::{expand_unshell_leaf, expand_unshell_leaf_result}; diff --git a/unshell-leaves/Cargo.toml b/unshell-macros/Cargo.toml similarity index 57% rename from unshell-leaves/Cargo.toml rename to unshell-macros/Cargo.toml index f27c2e1..1fb533e 100644 --- a/unshell-leaves/Cargo.toml +++ b/unshell-macros/Cargo.toml @@ -1,20 +1,14 @@ [package] -name = "unshell-leaves" +name = "unshell-macros" version.workspace = true edition.workspace = true -description = "Application-layer UnShell leaves and client surfaces" +description = "Procedural macros for UnShell leaves" -[features] -default = [] -leaf_endpoint = [] -leaf_tui = [] +[lib] +proc-macro = true [dependencies] -rkyv = { workspace = true } -portable-pty = { workspace = true } -crossbeam-channel = { workspace = true } -unshell-macros = { workspace = true } -unshell-protocol = { workspace = true } +unshell-macros-core = { workspace = true } [lints.rust] elided_lifetimes_in_paths = "warn" @@ -27,4 +21,3 @@ unsafe_op_in_unsafe_fn = "warn" unused_import_braces = "warn" unused_lifetimes = "warn" trivial_casts = "allow" -missing_docs = "warn" diff --git a/unshell-macros/src/lib.rs b/unshell-macros/src/lib.rs new file mode 100644 index 0000000..9751ccc --- /dev/null +++ b/unshell-macros/src/lib.rs @@ -0,0 +1,15 @@ +//! Procedural macro shim for UnShell. +//! +//! The real parser and code generator live in `unshell-macros-core` so they can be +//! tested as ordinary Rust. This crate only adapts compiler `TokenStream`s. + +use proc_macro::TokenStream; + +/// Generates an `unshell_protocol::Leaf` wrapper for a user-owned state struct. +/// +/// See `LEAF_MACRO_INTERFACE.md` for the design contract. The generated wrapper owns +/// session stores, retry queues, filtered packet dispatch, and final-frame cleanup. +#[proc_macro_attribute] +pub fn unshell_leaf(attr: TokenStream, item: TokenStream) -> TokenStream { + unshell_macros_core::expand_unshell_leaf(attr.into(), item.into()).into() +} diff --git a/unshell-protocol/src/endpoint/mod.rs b/unshell-protocol/src/endpoint/mod.rs index a9868a9..c57772b 100644 --- a/unshell-protocol/src/endpoint/mod.rs +++ b/unshell-protocol/src/endpoint/mod.rs @@ -69,6 +69,37 @@ impl Endpoint { 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.inbound.remove(&path) 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.inbound.entry(path).or_default().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 diff --git a/unshell-protocol/src/leaf.rs b/unshell-protocol/src/leaf.rs new file mode 100644 index 0000000..aa4491b --- /dev/null +++ b/unshell-protocol/src/leaf.rs @@ -0,0 +1,359 @@ +use crate::{Endpoint, HookID, Packet, PacketQueue}; + +use alloc::vec::Vec; + +/// Application extension point hosted by an [`Endpoint`]. +/// +/// A leaf owns product-specific state and reacts to packets that endpoint routing has +/// already delivered locally. The trait intentionally stays small so handwritten +/// leaves, generated leaves, and test leaves can all share the same endpoint loop. +pub trait Leaf { + /// Returns the stable local identifier for this leaf implementation. + fn get_id(&self) -> u32; + + /// Advances the leaf by one endpoint update tick. + /// + /// Implementations normally drain matching inbound packets, mutate leaf-owned + /// state, then enqueue outbound packets with [`Endpoint::add_outbound`]. + fn update(&mut self, _: &mut Endpoint); +} + +/// 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 implementation owns only application behavior. +/// +/// # Example +/// +/// ```rust,ignore +/// impl Session for MySession { +/// 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)) +/// } +/// +/// fn update( +/// leaf: &mut MyLeafState, +/// session: &mut Self::State, +/// incoming: &mut PacketQueue, +/// ctx: &mut SessionCtx<'_>, +/// ) -> SessionStatus { +/// while let Some(packet) = incoming.pop_front() { +/// session.apply(leaf, packet, ctx); +/// } +/// SessionStatus::Running +/// } +/// } +/// ``` +pub trait Session { + /// 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; + + /// 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; + + /// 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. + fn update( + leaf: &mut L, + session: &mut Self::State, + incoming: &mut PacketQueue, + ctx: &mut SessionCtx<'_>, + ) -> SessionStatus; +} + +/// 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); +} + +/// 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 + } +} + +/// 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. + Rejected, + + /// The packet was rejected with a response that the generated leaf must route. + RejectedWith(Packet), +} + +/// 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 still retains the entry until every queued packet routes + /// successfully, which prevents a failed final frame from losing session cleanup. + 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, + reply_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, + reply_path: Vec, + procedure_id: u32, + outbox: &'a mut PacketQueue, + ) -> Self { + Self { + hook_id, + reply_path, + procedure_id, + outbox, + } + } + + /// Returns the hook id used for packets emitted through this context. + pub fn hook_id(&self) -> HookID { + 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); + } + + /// 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.reply_path.clone(), + procedure_id: self.procedure_id, + data, + }); + } +} + +/// Output accumulator passed to [`Procedure::handle`]. +pub struct ProcedureOut { + hook_id: HookID, + reply_path: Vec, + procedure_id: u32, + outbox: PacketQueue, +} + +impl ProcedureOut { + /// Creates an empty procedure output queue. + pub fn new(hook_id: HookID, reply_path: Vec, procedure_id: u32) -> Self { + Self { + hook_id, + reply_path, + procedure_id, + outbox: PacketQueue::new(), + } + } + + /// Replaces the response path used by later [`Self::send`] calls. + pub fn set_reply_path(&mut self, reply_path: Vec) { + self.reply_path = reply_path; + } + + /// Queues raw response data without closing the hook. + pub fn send(&mut self, data: &[u8]) { + self.send_with_end(data, false); + } + + /// Queues raw response data that closes the hook after successful routing. + pub fn send_final(&mut self, data: &[u8]) { + self.send_with_end(data, true); + } + + /// Consumes the output accumulator and returns packets for generated retry logic. + pub fn into_packets(self) -> PacketQueue { + self.outbox + } + + fn send_with_end(&mut self, data: &[u8], end_hook: bool) { + self.outbox.push_back(Packet { + hook_id: self.hook_id, + end_hook, + path: self.reply_path.clone(), + procedure_id: self.procedure_id, + data: data.to_vec(), + }); + } +} + +/// 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, + + /// 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. + pub closed: bool, +} + +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(), + 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-protocol/src/lib.rs b/unshell-protocol/src/lib.rs index e8ff7c6..1ec332e 100644 --- a/unshell-protocol/src/lib.rs +++ b/unshell-protocol/src/lib.rs @@ -1,23 +1,17 @@ #![no_std] -extern crate alloc; +pub extern crate alloc; mod endpoint; mod error; +mod leaf; mod packet; pub use endpoint::{Endpoint, HookID}; pub use error::*; +pub use leaf::*; pub use packet::Packet; -pub trait Leaf { - // Identifier for this leaf - fn get_id(&self) -> u32; - - // Gets called every program loop - fn update(&mut self, _: &mut Endpoint); -} - // Various named types used for brevity use alloc::{ collections::{btree_map::BTreeMap, btree_set::BTreeSet, vec_deque::VecDeque}, @@ -28,7 +22,7 @@ type Path = Vec; type EndpointName = u32; type ConnectionSet = BTreeSet<(EndpointName, bool)>; type HookMap = BTreeMap; -type PacketQueue = VecDeque; +pub type PacketQueue = VecDeque; type RouteMap = BTreeMap; #[cfg(test)] diff --git a/unshell-protocol/src/packet.rs b/unshell-protocol/src/packet.rs index 37d6699..f908f3f 100644 --- a/unshell-protocol/src/packet.rs +++ b/unshell-protocol/src/packet.rs @@ -10,7 +10,7 @@ use crate::{DeserializeError, SerializeError}; /// path. `procedure_id` is therefore a compact numeric contract id instead of a /// string label; application code can maintain its own id-to-name table outside the /// hot packet path if it needs human-readable names. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Packet { pub hook_id: u16, pub end_hook: bool,