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"
dependencies = [
"leaf-shell",
"tcp_simple",
"unshell",
]
+1
View File
@@ -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"
+9 -1
View File
@@ -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);
}
}
+9 -4
View File
@@ -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
}
}
+2 -1
View File
@@ -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);
}
+9 -2
View File
@@ -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<u32>;
type EndpointName = u32;
type ConnectionSet = Vec<(EndpointName, bool)>;
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)>;
#[cfg(test)]
+1 -1
View File
@@ -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(),
+1 -1
View File
@@ -37,7 +37,7 @@ pub fn dispatch_session_interface<L, S>(
.iter_mut()
.find(|entry| entry.hook_id == hook_id)
{
entry.inbox.push_back(packet);
entry.inbox.push(packet);
interface.record_for(
target,
+1 -1
View File
@@ -23,7 +23,7 @@ pub fn dispatch_session<L, S>(
.iter_mut()
.find(|entry| entry.hook_id == hook_id)
{
entry.inbox.push_back(packet);
entry.inbox.push(packet);
return;
}
+2 -1
View File
@@ -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
+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)
.unwrap()
.front()
.first()
.unwrap();
assert!(response.end_hook);
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)
.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()
}
+2 -1
View File
@@ -59,7 +59,8 @@ impl Session<FakePtyState> 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(
+2 -2
View File
@@ -74,7 +74,6 @@ impl ShellSession {
fn spawn(hook_id: HookID) -> Result<Self, SessionInitError> {
let child = Command::new("/bin/bash")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.spawn()
.map_err(|_| SessionInitError::rejected())?;
@@ -116,7 +115,8 @@ impl Session<ShellState> 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();
}