astatin3/unshell
1
0
Fork 0
mirror of https://github.com/Astatin3/unshell.git synced 2026-06-08 22:38:01 -06:00
Code Issues Packages Projects Releases Wiki Activity

Simplify session routing path

Browse Source
This commit is contained in:
Michael Mikovsky
2026-06-01 12:11:01 -06:00
parent 64e53c8cfe
commit 4cd496ed2b
12 changed files with 166 additions and 396 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/protocol/endpoint/hooks.rs
+47 -1
View File
@@ -1,6 +1,6 @@
use alloc::vec::Vec;
use crate::protocol::{Endpoint, EndpointError, EndpointName};
use crate::protocol::{Endpoint, EndpointError, EndpointName, Packet};
/// Compact identifier for one routed return channel.
///
@@ -105,6 +105,52 @@ impl Endpoint {
}
}
/// Routes raw response data over an existing hook immediately.
///
/// This is the compact session-output path: it avoids an intermediate context and
/// retry queue. If a final packet cannot route, the local hook is still removed so
/// an implant does not retain dead hook state forever.
pub fn send_hook_raw(
&mut self,
hook_id: HookID,
procedure_id: u32,
data: Vec<u8>,
end_hook: bool,
) -> Result<(), EndpointError> {
let path = self.hook_path(hook_id)?;
let packet = Packet {
hook_id,
end_hook,
path,
procedure_id,
data,
};
let result = self.add_outbound(packet);
if result.is_err() && end_hook {
self.close_hook(hook_id);
}
result
}
/// Routes a one-byte-opcode response frame over an existing hook immediately.
pub fn send_hook_frame(
&mut self,
hook_id: HookID,
procedure_id: u32,
opcode: u8,
payload: &[u8],
end_hook: bool,
) -> Result<(), EndpointError> {
let mut data = Vec::with_capacity(payload.len() + 1);
data.push(opcode);
data.extend_from_slice(payload);
self.send_hook_raw(hook_id, procedure_id, data, end_hook)
}
/// Validates that `actual_peer` is the peer allowed to use `hook_id`.
pub(crate) fn ensure_hook_peer(
&self,
src/protocol/leaf_template.rs
+1 -10
View File
@@ -98,6 +98,7 @@ macro_rules! unshell_leaf {
$(
$crate::protocol::update_session_family::<$State, $Session>(
endpoint,
leaf_id,
&mut self.state,
&mut self.$session_field,
@@ -126,7 +127,6 @@ macro_rules! unshell_leaf {
&mut self.state,
&mut self.$session_field,
packet,
&mut self.outbox,
interface,
);
return;
@@ -167,15 +167,6 @@ macro_rules! unshell_leaf {
&mut self.outbox,
interface,
);
$(
$crate::protocol::flush_session_family::<$State, $Session>(
endpoint,
leaf_id,
&mut self.$session_field,
interface,
);
)*
}
}
src/protocol/runtime.rs
+17 -85
View File
@@ -3,26 +3,27 @@ use alloc::collections::VecDeque;
use crate::{
interface::{InterfaceEventKind, InterfaceStore, InterfaceTarget},
protocol::{
Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionCtx, SessionEntry,
Endpoint, Packet, PacketQueue, Procedure, ProcedureOut, Session, SessionEntry,
SessionFamily, SessionInitError, SessionStatus,
},
};
/// Retry queue shared by generated leaves.
///
/// Sessions already own per-hook outboxes. This leaf-level queue is for rejected
/// session initialization responses and one-shot procedures, both of which need the
/// same retry semantics as session output without becoming separate framework types.
/// Leaf-level retry queue shared by generated leaves.
///
/// Sessions route directly through `Endpoint` to keep their runtime shape small. This
/// queue remains only for one-shot procedures, whose handlers still use `ProcedureOut`
/// and should not route while the procedure is borrowing leaf state.
pub struct LeafOutbox {
packets: VecDeque<LeafOutboxEntry>,
}
/// One packet retained by a leaf-level retry queue.
///
/// Session entry outboxes have an obvious owner from their surrounding session entry.
/// Leaf-level outboxes are mixed: rejected session initialization packets and one-shot
/// procedure responses both land here. Storing the owner beside the packet keeps route
/// logging precise without exposing another public queue type.
/// Procedure responses from different generated branches share one queue. Storing the
/// owner beside the packet keeps route logging precise without exposing another public
/// queue type.
#[derive(Clone)]
struct LeafOutboxEntry {
packet: Packet,
@@ -85,15 +86,14 @@ impl Default for LeafOutbox {
/// Dispatches one packet into a generated session family.
///
/// The macro picks `S` and the family field. This helper owns the boring details:
/// find the hook, initialize missing sessions, queue rejected responses, and update
/// find the hook, initialize missing sessions, route rejected responses, and update
/// interface state when a caller supplied one.
pub fn dispatch_session<L, S>(
endpoint: &Endpoint,
endpoint: &mut Endpoint,
leaf_id: u32,
leaf: &mut L,
family: &mut SessionFamily<S>,
packet: Packet,
outbox: &mut LeafOutbox,
interface: &mut Option<&mut InterfaceStore>,
) where
S: Session<L>,
@@ -149,7 +149,7 @@ pub fn dispatch_session<L, S>(
};
match S::init(leaf, packet) {
Ok(state) => {
family.entries.push(SessionEntry::new(hook_id, path, state));
family.entries.push(SessionEntry::new(hook_id, state));
if let Some(store) = interface.as_mut() {
store.record_for(
@@ -195,21 +195,16 @@ pub fn dispatch_session<L, S>(
finished_ns: store.now_ns(),
},
);
store.record_for(
target,
InterfaceEventKind::OutboundQueued {
packet: packet.clone(),
},
);
}
outbox.push_for_target(packet, target);
let _ = flush_packet_with_target(endpoint, target, &packet, interface);
}
}
}
/// Updates every live session in one generated session family.
pub fn update_session_family<L, S>(
endpoint: &mut Endpoint,
leaf_id: u32,
leaf: &mut L,
family: &mut SessionFamily<S>,
@@ -223,13 +218,7 @@ pub fn update_session_family<L, S>(
}
let started_ns = interface.as_ref().and_then(|store| store.now_ns());
let outbox_start = entry.outbox.len();
let path = entry.path.clone();
let status = {
let mut ctx = SessionCtx::new(entry.hook_id, path, S::PROCEDURE_ID, &mut entry.outbox);
S::update(leaf, &mut entry.state, &mut entry.inbox, &mut ctx)
};
let status = S::update(leaf, &mut entry.state, &mut entry.inbox, endpoint);
let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id);
if let Some(store) = interface.as_mut() {
@@ -243,21 +232,14 @@ pub fn update_session_family<L, S>(
finished_ns: store.now_ns(),
},
);
for packet in entry.outbox.iter().skip(outbox_start) {
store.record_for(
target,
InterfaceEventKind::OutboundQueued {
packet: packet.clone(),
},
);
}
}
if matches!(status, SessionStatus::Closed) {
entry.closed = true;
}
}
family.entries.retain(|entry| !entry.closed);
}
/// Dispatches one packet into a generated one-shot procedure.
@@ -331,56 +313,6 @@ pub fn flush_leaf_outbox(
})
}
/// Flushes and retains one generated session family.
pub fn flush_session_family<L, S>(
endpoint: &mut Endpoint,
leaf_id: u32,
family: &mut SessionFamily<S>,
interface: &mut Option<&mut InterfaceStore>,
) where
S: Session<L>,
{
for entry in &mut family.entries {
let target = InterfaceTarget::session(leaf_id, S::PROCEDURE_ID, entry.hook_id);
flush_packet_queue_with_target(endpoint, target, &mut entry.outbox, interface);
}
family
.entries
.retain(|entry| !entry.closed || !entry.outbox.is_empty());
}
/// Flushes a retry queue through [`Endpoint::add_outbound`].
///
/// This is the interface-aware version of [`crate::protocol::flush_packet_queue`]. It
/// logs route attempts before trying them, then logs either success or the route error
/// without dropping the packet on failure.
pub fn flush_packet_queue_with_interface(
endpoint: &mut Endpoint,
leaf_id: u32,
outbox: &mut PacketQueue,
interface: &mut Option<&mut InterfaceStore>,
) -> bool {
flush_outbox(endpoint, outbox, interface, |packet| {
(
InterfaceTarget::session(leaf_id, packet.procedure_id, packet.hook_id),
packet.clone(),
)
})
}
/// Flushes a packet queue whose owner is already known by the generated runtime.
fn flush_packet_queue_with_target(
endpoint: &mut Endpoint,
target: InterfaceTarget,
outbox: &mut PacketQueue,
interface: &mut Option<&mut InterfaceStore>,
) -> bool {
flush_outbox(endpoint, outbox, interface, |packet| {
(target, packet.clone())
})
}
fn flush_outbox<T>(
endpoint: &mut Endpoint,
outbox: &mut VecDeque<T>,
src/protocol/session.rs
+12 -128
View File
@@ -28,10 +28,10 @@ use crate::interface::SessionView;
/// leaf: &mut MyLeafState,
/// session: &mut Self,
/// incoming: &mut PacketQueue,
/// ctx: &mut SessionCtx<'_>,
/// endpoint: &mut Endpoint,
/// ) -> SessionStatus {
/// while let Some(packet) = incoming.pop_front() {
/// session.apply(leaf, packet, ctx);
/// session.apply(leaf, packet, endpoint);
/// }
/// SessionStatus::Running
/// }
@@ -51,14 +51,15 @@ pub trait Session<L>: Sized {
/// 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.
/// sessions can poll external workers even when no new packet arrived. Session
/// output is routed immediately through `endpoint`; callers that need retry
/// semantics should keep their own compact application state and retry on a later
/// tick.
fn update(
leaf: &mut L,
session: &mut Self,
incoming: &mut PacketQueue,
ctx: &mut SessionCtx<'_>,
endpoint: &mut Endpoint,
) -> SessionStatus;
#[cfg(feature = "interface_ratatui")]
@@ -121,99 +122,11 @@ pub enum SessionStatus {
/// 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.
/// The generated leaf removes the entry after the update tick. Final packets are
/// routed immediately by the session before returning this status.
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,
path: Vec<u32>,
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,
path: Vec<u32>,
procedure_id: u32,
outbox: &'a mut PacketQueue,
) -> Self {
Self {
hook_id,
path,
procedure_id,
outbox,
}
}
/// Returns the hook id used for packets emitted through this context.
pub fn hook_id(&self) -> HookID {
self.hook_id
}
/// 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<u8>, end_hook: bool) {
self.outbox.push_back(Packet {
hook_id: self.hook_id,
end_hook,
path: self.path.clone(),
procedure_id: self.procedure_id,
data,
});
}
}
/// Storage entry used by macro-generated session stores.
///
/// The fields are public so generated code in downstream crates can keep the update
@@ -223,23 +136,13 @@ pub struct SessionEntry<S> {
/// Hook id associated with this live session.
pub hook_id: HookID,
/// Destination path for packets emitted on this hook.
///
/// This is generated runtime state, not user session state. It is captured from
/// endpoint hook routing when the session is created so leaf sessions never have
/// to carry or understand a reply path.
pub path: Vec<u32>,
/// 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.
/// Whether application logic has finished and should be removed after update.
pub closed: bool,
}
@@ -266,7 +169,7 @@ impl<S> SessionFamily<S> {
let mut count = 0usize;
for entry in &self.entries {
count += entry.inbox.len() + entry.outbox.len();
count += entry.inbox.len();
}
count
@@ -281,31 +184,12 @@ impl<S> Default for SessionFamily<S> {
impl<S> SessionEntry<S> {
/// Creates one active session entry for `hook_id`.
pub fn new(hook_id: HookID, path: Vec<u32>, state: S) -> Self {
pub fn new(hook_id: HookID, state: S) -> Self {
Self {
hook_id,
path,
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
}