Optimize tcp_simple endpoint integration

This commit is contained in:
Michael Mikovsky
2026-06-01 15:54:17 -06:00
parent 9ab130a620
commit 641ee7682a
14 changed files with 43 additions and 18 deletions
Generated
+1
View File
@@ -517,6 +517,7 @@ name = "endpoint_test"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"leaf-shell", "leaf-shell",
"tcp_simple",
"unshell", "unshell",
] ]
+1
View File
@@ -10,6 +10,7 @@ include.workspace = true
[dependencies] [dependencies]
unshell = { workspace = true } unshell = { workspace = true }
leaf-shell = { path = "../../unshell-leaves/leaf-shell" } leaf-shell = { path = "../../unshell-leaves/leaf-shell" }
tcp_simple = { path = "../../unshell-leaves/tcp_simple" }
[[bin]] [[bin]]
name = "endpoint_test" name = "endpoint_test"
+9 -1
View File
@@ -1,19 +1,27 @@
#![no_std] #![no_std]
#![no_main] #![no_main]
extern crate alloc; use core::net::Ipv4Addr;
use leaf_shell::{ShellLeaf, ShellState}; use leaf_shell::{ShellLeaf, ShellState};
use tcp_simple::TCPServerLeaf;
use unshell::protocol::{Endpoint, Leaf}; use unshell::protocol::{Endpoint, Leaf};
const ID: u32 = 0x12345678; const ID: u32 = 0x12345678;
const CHILD_ID: u32 = 0x87654321;
#[unsafe(no_mangle)] #[unsafe(no_mangle)]
pub fn main(_argc: i32, _argv: *const *const u8) { pub fn main(_argc: i32, _argv: *const *const u8) {
let mut endpoint = Endpoint::new(ID); let mut endpoint = Endpoint::new(ID);
endpoint.path.push(ID);
let mut shell = ShellLeaf::new(ShellState::new()); let mut shell = ShellLeaf::new(ShellState::new());
let mut tcp = TCPServerLeaf::bind_ipv4(Ipv4Addr::LOCALHOST, 1337, CHILD_ID).unwrap();
loop { 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); shell.update(&mut endpoint);
tcp.update(&mut endpoint);
} }
} }
+9 -4
View File
@@ -29,7 +29,8 @@ impl Endpoint {
let mut unmatched = Vec::new(); let mut unmatched = Vec::new();
while let Some(packet) = queue.pop_front() { while !queue.is_empty() {
let packet = queue.remove(0);
if predicate(&packet) { if predicate(&packet) {
f(packet); f(packet);
} else { } else {
@@ -74,7 +75,7 @@ impl Endpoint {
/// Appends a packet to the route queue for `endpoint`. /// Appends a packet to the route queue for `endpoint`.
pub(crate) fn route_push(endpoint: EndpointName, packet: Packet, routes: &mut RouteMap) { 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. /// Returns the route queue for `endpoint` if one exists.
@@ -87,6 +88,9 @@ impl Endpoint {
} }
/// Removes and returns the queue for `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( pub(crate) fn route_remove(
endpoint: EndpointName, endpoint: EndpointName,
routes: &mut RouteMap, routes: &mut RouteMap,
@@ -95,7 +99,7 @@ impl Endpoint {
.iter() .iter()
.position(|(queued_endpoint, _)| *queued_endpoint == endpoint)?; .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`. /// Returns whether a route queue exists for `endpoint`.
@@ -118,8 +122,9 @@ impl Endpoint {
{ {
&mut routes[index].1 &mut routes[index].1
} else { } else {
let index = routes.len();
routes.push((endpoint, PacketQueue::new())); routes.push((endpoint, PacketQueue::new()));
&mut routes.last_mut().unwrap().1 &mut routes[index].1
} }
} }
+2 -1
View File
@@ -76,6 +76,7 @@ macro_rules! unshell_leaf {
)* )*
} }
#[inline(never)]
fn __unshell_update_inner( fn __unshell_update_inner(
&mut self, &mut self,
endpoint: &mut $crate::protocol::Endpoint, endpoint: &mut $crate::protocol::Endpoint,
@@ -269,7 +270,7 @@ macro_rules! unshell_leaf {
$id $id
} }
#[inline(never)] #[inline(always)]
fn update(&mut self, endpoint: &mut $crate::protocol::Endpoint) { fn update(&mut self, endpoint: &mut $crate::protocol::Endpoint) {
self.__unshell_update_inner(endpoint); self.__unshell_update_inner(endpoint);
} }
+9 -2
View File
@@ -22,13 +22,20 @@ pub use session::*;
pub use ratatui; pub use ratatui;
// Various named types used for brevity // Various named types used for brevity
use alloc::{collections::vec_deque::VecDeque, vec::Vec}; use alloc::vec::Vec;
type Path = Vec<u32>; type Path = Vec<u32>;
type EndpointName = u32; type EndpointName = u32;
type ConnectionSet = Vec<(EndpointName, bool)>; type ConnectionSet = Vec<(EndpointName, bool)>;
type HookMap = Vec<(HookID, EndpointName)>; type HookMap = Vec<(HookID, EndpointName)>;
pub type PacketQueue = VecDeque<Packet>;
/// 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<Packet>;
type RouteMap = Vec<(EndpointName, PacketQueue)>; type RouteMap = Vec<(EndpointName, PacketQueue)>;
#[cfg(test)] #[cfg(test)]
+1 -1
View File
@@ -42,7 +42,7 @@ impl ProcedureOut {
} }
fn send_with_end(&mut self, data: &[u8], end_hook: bool) { 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, hook_id: self.hook_id,
end_hook, end_hook,
path: self.reply_path.clone(), path: self.reply_path.clone(),
+1 -1
View File
@@ -37,7 +37,7 @@ pub fn dispatch_session_interface<L, S>(
.iter_mut() .iter_mut()
.find(|entry| entry.hook_id == hook_id) .find(|entry| entry.hook_id == hook_id)
{ {
entry.inbox.push_back(packet); entry.inbox.push(packet);
interface.record_for( interface.record_for(
target, target,
+1 -1
View File
@@ -23,7 +23,7 @@ pub fn dispatch_session<L, S>(
.iter_mut() .iter_mut()
.find(|entry| entry.hook_id == hook_id) .find(|entry| entry.hook_id == hook_id)
{ {
entry.inbox.push_back(packet); entry.inbox.push(packet);
return; return;
} }
+2 -1
View File
@@ -28,7 +28,8 @@ use crate::interface::SessionView;
/// incoming: &mut PacketQueue, /// incoming: &mut PacketQueue,
/// endpoint: &mut Endpoint, /// endpoint: &mut Endpoint,
/// ) -> SessionStatus { /// ) -> SessionStatus {
/// while let Some(packet) = incoming.pop_front() { /// while !incoming.is_empty() {
/// let packet = incoming.remove(0);
/// session.apply(leaf, packet, endpoint); /// session.apply(leaf, packet, endpoint);
/// } /// }
/// SessionStatus::Running /// SessionStatus::Running
+1 -1
View File
@@ -141,7 +141,7 @@ fn request_response_round_trip_over_mock_transport() {
); );
let response = &Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound) let response = &Endpoint::route_get(ENDPOINT_A, &endpoint_a.inbound)
.unwrap() .unwrap()
.front() .first()
.unwrap(); .unwrap();
assert!(response.end_hook); assert!(response.end_hook);
assert_eq!(response.data, "ABC123".as_bytes()); assert_eq!(response.data, "ABC123".as_bytes());
+2 -2
View File
@@ -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) let queue = Endpoint::route_get(next_hop, &endpoint.outbound)
.unwrap_or_else(|| panic!("expected one outbound queue for {next_hop}")); .unwrap_or_else(|| panic!("expected one outbound queue for {next_hop}"));
assert_eq!(queue.len(), 1, "expected exactly one outbound packet"); 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`. /// 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) let queue = Endpoint::route_get(local_id, &endpoint.inbound)
.unwrap_or_else(|| panic!("expected one inbound queue for {local_id}")); .unwrap_or_else(|| panic!("expected one inbound queue for {local_id}"));
assert_eq!(queue.len(), 1, "expected exactly one inbound packet"); assert_eq!(queue.len(), 1, "expected exactly one inbound packet");
queue.front().unwrap() queue.first().unwrap()
} }
+2 -1
View File
@@ -59,7 +59,8 @@ impl Session<FakePtyState> for PtySessionState {
session.opened_pending = false; session.opened_pending = false;
} }
while let Some(packet) = incoming.pop_front() { while !incoming.is_empty() {
let packet = incoming.remove(0);
match frame_opcode(&packet) { match frame_opcode(&packet) {
Some(OP_INPUT) => { Some(OP_INPUT) => {
let _ = endpoint.send_hook_frame( let _ = endpoint.send_hook_frame(
+2 -2
View File
@@ -74,7 +74,6 @@ impl ShellSession {
fn spawn(hook_id: HookID) -> Result<Self, SessionInitError> { fn spawn(hook_id: HookID) -> Result<Self, SessionInitError> {
let child = Command::new("/bin/bash") let child = Command::new("/bin/bash")
.stdin(Stdio::piped()) .stdin(Stdio::piped())
.stdout(Stdio::piped())
.spawn() .spawn()
.map_err(|_| SessionInitError::rejected())?; .map_err(|_| SessionInitError::rejected())?;
@@ -116,7 +115,8 @@ impl Session<ShellState> for ShellSession {
incoming: &mut PacketQueue, incoming: &mut PacketQueue,
_endpoint: &mut Endpoint, _endpoint: &mut Endpoint,
) -> SessionStatus { ) -> SessionStatus {
while let Some(packet) = incoming.pop_front() { while !incoming.is_empty() {
let packet = incoming.remove(0);
if packet.end_hook { if packet.end_hook {
session.close_stdin(); session.close_stdin();
} }