mirror of
https://github.com/Astatin3/unshell.git
synced 2026-06-08 22:38:01 -06:00
Optimize tcp_simple endpoint integration
This commit is contained in:
Generated
+1
@@ -517,6 +517,7 @@ name = "endpoint_test"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"leaf-shell",
|
||||
"tcp_simple",
|
||||
"unshell",
|
||||
]
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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)]
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user