diff --git a/Cargo.lock b/Cargo.lock index 3b201aa..64482e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -517,6 +517,7 @@ name = "endpoint_test" version = "0.1.0" dependencies = [ "leaf-shell", + "tcp_simple", "unshell", ] diff --git a/examples/endpoint_test/Cargo.toml b/examples/endpoint_test/Cargo.toml index b6a7a76..5d9e5ba 100644 --- a/examples/endpoint_test/Cargo.toml +++ b/examples/endpoint_test/Cargo.toml @@ -10,6 +10,7 @@ include.workspace = true [dependencies] unshell = { workspace = true } leaf-shell = { path = "../../unshell-leaves/leaf-shell" } +tcp_simple = { path = "../../unshell-leaves/tcp_simple" } [[bin]] name = "endpoint_test" diff --git a/examples/endpoint_test/src/main.rs b/examples/endpoint_test/src/main.rs index 7d4f826..1aca7cb 100644 --- a/examples/endpoint_test/src/main.rs +++ b/examples/endpoint_test/src/main.rs @@ -1,19 +1,27 @@ #![no_std] #![no_main] -extern crate alloc; +use core::net::Ipv4Addr; use leaf_shell::{ShellLeaf, ShellState}; +use tcp_simple::TCPServerLeaf; use unshell::protocol::{Endpoint, Leaf}; const ID: u32 = 0x12345678; +const CHILD_ID: u32 = 0x87654321; #[unsafe(no_mangle)] pub fn main(_argc: i32, _argv: *const *const u8) { let mut endpoint = Endpoint::new(ID); + endpoint.path.push(ID); + let mut shell = ShellLeaf::new(ShellState::new()); + let mut tcp = TCPServerLeaf::bind_ipv4(Ipv4Addr::LOCALHOST, 1337, CHILD_ID).unwrap(); loop { + // One transport tick keeps the minimized binary smaller. Packets read from TCP + // are processed by the shell on the next loop, then flushed by this same tick. shell.update(&mut endpoint); + tcp.update(&mut endpoint); } } diff --git a/src/protocol/endpoint/queues.rs b/src/protocol/endpoint/queues.rs index 0f17d5b..383298f 100644 --- a/src/protocol/endpoint/queues.rs +++ b/src/protocol/endpoint/queues.rs @@ -29,7 +29,8 @@ impl Endpoint { let mut unmatched = Vec::new(); - while let Some(packet) = queue.pop_front() { + while !queue.is_empty() { + let packet = queue.remove(0); if predicate(&packet) { f(packet); } else { @@ -74,7 +75,7 @@ impl Endpoint { /// 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); + Self::route_queue_mut(endpoint, routes).push(packet); } /// Returns the route queue for `endpoint` if one exists. @@ -87,6 +88,9 @@ impl Endpoint { } /// Removes and returns the queue for `endpoint`. + /// + /// Route map entry order is intentionally not observable; each route's packet + /// queue preserves FIFO ordering internally, so `swap_remove` keeps removal small. pub(crate) fn route_remove( endpoint: EndpointName, routes: &mut RouteMap, @@ -95,7 +99,7 @@ impl Endpoint { .iter() .position(|(queued_endpoint, _)| *queued_endpoint == endpoint)?; - Some(routes.remove(index).1) + Some(routes.swap_remove(index).1) } /// Returns whether a route queue exists for `endpoint`. @@ -118,8 +122,9 @@ impl Endpoint { { &mut routes[index].1 } else { + let index = routes.len(); routes.push((endpoint, PacketQueue::new())); - &mut routes.last_mut().unwrap().1 + &mut routes[index].1 } } diff --git a/src/protocol/leaf_template.rs b/src/protocol/leaf_template.rs index eac74fe..9eefad3 100644 --- a/src/protocol/leaf_template.rs +++ b/src/protocol/leaf_template.rs @@ -76,6 +76,7 @@ macro_rules! unshell_leaf { )* } + #[inline(never)] fn __unshell_update_inner( &mut self, endpoint: &mut $crate::protocol::Endpoint, @@ -269,7 +270,7 @@ macro_rules! unshell_leaf { $id } - #[inline(never)] + #[inline(always)] fn update(&mut self, endpoint: &mut $crate::protocol::Endpoint) { self.__unshell_update_inner(endpoint); } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index edc92c7..ad09820 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -22,13 +22,20 @@ pub use session::*; pub use ratatui; // Various named types used for brevity -use alloc::{collections::vec_deque::VecDeque, vec::Vec}; +use alloc::vec::Vec; type Path = Vec; type EndpointName = u32; type ConnectionSet = Vec<(EndpointName, bool)>; type HookMap = Vec<(HookID, EndpointName)>; -pub type PacketQueue = VecDeque; + +/// FIFO packet storage for generated leaves and endpoint route queues. +/// +/// A compact `Vec` is smaller than `VecDeque` in minimized endpoint binaries. Callers +/// that drain the front should use `remove(0)` to preserve protocol packet order; the +/// expected per-hook queues are short enough that the O(n) shift is preferable to the +/// larger ring-buffer machinery in implant-sized builds. +pub type PacketQueue = Vec; type RouteMap = Vec<(EndpointName, PacketQueue)>; #[cfg(test)] diff --git a/src/protocol/procedure/out.rs b/src/protocol/procedure/out.rs index 9c13f53..0c8348c 100644 --- a/src/protocol/procedure/out.rs +++ b/src/protocol/procedure/out.rs @@ -42,7 +42,7 @@ impl ProcedureOut { } fn send_with_end(&mut self, data: &[u8], end_hook: bool) { - self.outbox.push_back(Packet { + self.outbox.push(Packet { hook_id: self.hook_id, end_hook, path: self.reply_path.clone(), diff --git a/src/protocol/runtime/interface.rs b/src/protocol/runtime/interface.rs index 5da610c..6c0e026 100644 --- a/src/protocol/runtime/interface.rs +++ b/src/protocol/runtime/interface.rs @@ -37,7 +37,7 @@ pub fn dispatch_session_interface( .iter_mut() .find(|entry| entry.hook_id == hook_id) { - entry.inbox.push_back(packet); + entry.inbox.push(packet); interface.record_for( target, diff --git a/src/protocol/runtime/session.rs b/src/protocol/runtime/session.rs index 5dda6ae..2813863 100644 --- a/src/protocol/runtime/session.rs +++ b/src/protocol/runtime/session.rs @@ -23,7 +23,7 @@ pub fn dispatch_session( .iter_mut() .find(|entry| entry.hook_id == hook_id) { - entry.inbox.push_back(packet); + entry.inbox.push(packet); return; } diff --git a/src/protocol/session/contract.rs b/src/protocol/session/contract.rs index 97dcf81..aea7cba 100644 --- a/src/protocol/session/contract.rs +++ b/src/protocol/session/contract.rs @@ -28,7 +28,8 @@ use crate::interface::SessionView; /// incoming: &mut PacketQueue, /// endpoint: &mut Endpoint, /// ) -> SessionStatus { -/// while let Some(packet) = incoming.pop_front() { +/// while !incoming.is_empty() { +/// let packet = incoming.remove(0); /// session.apply(leaf, packet, endpoint); /// } /// SessionStatus::Running diff --git a/src/protocol/tests/integration/oneshot.rs b/src/protocol/tests/integration/oneshot.rs index e356aa9..991a224 100644 --- a/src/protocol/tests/integration/oneshot.rs +++ b/src/protocol/tests/integration/oneshot.rs @@ -141,7 +141,7 @@ fn request_response_round_trip_over_mock_transport() { ); let response = &Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) .unwrap() - .front() + .first() .unwrap(); assert!(response.end_hook); assert_eq!(response.data, "ABC123".as_bytes()); diff --git a/src/protocol/tests/support/endpoints.rs b/src/protocol/tests/support/endpoints.rs index ae45ede..13f30ea 100644 --- a/src/protocol/tests/support/endpoints.rs +++ b/src/protocol/tests/support/endpoints.rs @@ -26,7 +26,7 @@ pub(crate) fn single_outbound_packet(endpoint: &Endpoint, next_hop: u32) -> &Pac 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() + queue.first().unwrap() } /// Returns the only inbound packet delivered to `local_id`. @@ -38,5 +38,5 @@ pub(crate) fn single_inbound_packet(endpoint: &Endpoint, local_id: u32) -> &Pack 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() + queue.first().unwrap() } diff --git a/unshell-leaves/leaf-pty/src/session.rs b/unshell-leaves/leaf-pty/src/session.rs index 66a5e75..28b187f 100644 --- a/unshell-leaves/leaf-pty/src/session.rs +++ b/unshell-leaves/leaf-pty/src/session.rs @@ -59,7 +59,8 @@ impl Session for PtySessionState { session.opened_pending = false; } - while let Some(packet) = incoming.pop_front() { + while !incoming.is_empty() { + let packet = incoming.remove(0); match frame_opcode(&packet) { Some(OP_INPUT) => { let _ = endpoint.send_hook_frame( diff --git a/unshell-leaves/leaf-shell/src/shell/mod.rs b/unshell-leaves/leaf-shell/src/shell/mod.rs index c5d2a7e..1ede796 100644 --- a/unshell-leaves/leaf-shell/src/shell/mod.rs +++ b/unshell-leaves/leaf-shell/src/shell/mod.rs @@ -74,7 +74,6 @@ impl ShellSession { fn spawn(hook_id: HookID) -> Result { let child = Command::new("/bin/bash") .stdin(Stdio::piped()) - .stdout(Stdio::piped()) .spawn() .map_err(|_| SessionInitError::rejected())?; @@ -116,7 +115,8 @@ impl Session for ShellSession { incoming: &mut PacketQueue, _endpoint: &mut Endpoint, ) -> SessionStatus { - while let Some(packet) = incoming.pop_front() { + while !incoming.is_empty() { + let packet = incoming.remove(0); if packet.end_hook { session.close_stdin(); }