mirror of
https://github.com/Astatin3/unshell.git
synced 2026-06-09 06:47:59 -06:00
Move protocol to workspace root.
This commit is contained in:
Michael Mikovsky
@@ -0,0 +1,111 @@
|
||||
use crate::protocol::{Endpoint, EndpointError, EndpointName};
|
||||
|
||||
/// Compact identifier for one routed return channel.
|
||||
///
|
||||
/// Hook ids are local endpoint state, not globally unique session ids. A downward
|
||||
/// packet with `end_hook = false` reserves the id at each endpoint it crosses so
|
||||
/// later upward packets can prove that the route was paved by trusted downward
|
||||
/// traffic first.
|
||||
pub type HookID = u16;
|
||||
|
||||
impl Endpoint {
|
||||
/// Allocates a hook id that is not currently active on this endpoint.
|
||||
///
|
||||
/// The first id is still deterministic (`0`) for the protocol tests, but the
|
||||
/// allocator now skips active hooks so long-lived streams cannot accidentally
|
||||
/// reuse an id before the previous route has closed. If every `u16` id is active
|
||||
/// the function panics; that is a hard local resource exhaustion condition, not a
|
||||
/// recoverable packet error.
|
||||
pub fn allocate_hook_id(&mut self) -> HookID {
|
||||
for _ in 0..=HookID::MAX {
|
||||
let candidate = self.last_hook;
|
||||
self.last_hook = self.last_hook.wrapping_add(1);
|
||||
|
||||
if !self.hooks.contains_key(&candidate) {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
|
||||
// Avoid a panic message here: this crate is optimized for small binaries,
|
||||
// and exhausting every `u16` hook id is unrecoverable local state corruption.
|
||||
panic!();
|
||||
}
|
||||
|
||||
/// Backwards-compatible name for [`Self::allocate_hook_id`].
|
||||
///
|
||||
/// Existing leaves and tests still call `get_hook_id`; new code should prefer
|
||||
/// `allocate_hook_id` because it describes the reservation semantics more clearly.
|
||||
pub fn get_hook_id(&mut self) -> HookID {
|
||||
self.allocate_hook_id()
|
||||
}
|
||||
|
||||
/// Explicitly records that `peer` may use `hook_id` as this endpoint's return channel.
|
||||
///
|
||||
/// Routing calls this automatically for successful downward packets whose
|
||||
/// `end_hook` flag is false. The public method exists for trusted local setup and
|
||||
/// tests; ordinary leaf procedures should usually let packet routing pave hooks
|
||||
/// instead of mutating hook state by hand.
|
||||
pub fn accept_hook(&mut self, hook_id: HookID, peer: u32) -> Option<u32> {
|
||||
self.hooks.insert(hook_id, peer)
|
||||
}
|
||||
|
||||
/// Returns true when `hook_id` is currently active.
|
||||
pub fn has_hook(&self, hook_id: HookID) -> bool {
|
||||
self.hooks.contains_key(&hook_id)
|
||||
}
|
||||
|
||||
/// Returns the adjacent peer currently associated with `hook_id`.
|
||||
///
|
||||
/// The peer is the next endpoint expected to participate in the return channel:
|
||||
/// a child for downward calls that will reply upward, or a parent for a local
|
||||
/// callee that will emit an upward response.
|
||||
pub fn hook_peer(&self, hook_id: HookID) -> Option<u32> {
|
||||
self.hooks.get(&hook_id).copied()
|
||||
}
|
||||
|
||||
/// Returns the number of active hooks on this endpoint.
|
||||
pub fn hook_count(&self) -> usize {
|
||||
self.hooks.len()
|
||||
}
|
||||
|
||||
/// Locally forgets a hook without sending protocol traffic.
|
||||
///
|
||||
/// Graceful shutdown should use a packet with `end_hook = true` so every endpoint
|
||||
/// along the route cleans up after successful delivery. This method is for local
|
||||
/// emergency cleanup such as a crashed PTY process, a timed-out stream, or a lost
|
||||
/// transport where no final packet can be delivered.
|
||||
pub fn forget_hook(&mut self, hook_id: HookID) -> bool {
|
||||
self.close_hook(hook_id)
|
||||
}
|
||||
|
||||
/// Validates that `actual_peer` is the peer allowed to use `hook_id`.
|
||||
pub(crate) fn ensure_hook_peer(
|
||||
&self,
|
||||
hook_id: HookID,
|
||||
actual_peer: EndpointName,
|
||||
) -> Result<(), EndpointError> {
|
||||
let expected_peer = self
|
||||
.hook_peer(hook_id)
|
||||
.ok_or(EndpointError::UnknownHook { hook_id })?;
|
||||
|
||||
if expected_peer == actual_peer {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(EndpointError::HookPeerMismatch {
|
||||
hook_id,
|
||||
expected_peer,
|
||||
actual_peer,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens or refreshes `hook_id` for the adjacent `peer` after downward routing succeeds.
|
||||
pub(crate) fn open_hook(&mut self, hook_id: HookID, peer: EndpointName) {
|
||||
self.hooks.insert(hook_id, peer);
|
||||
}
|
||||
|
||||
/// Removes `hook_id` and reports whether it existed.
|
||||
pub(crate) fn close_hook(&mut self, hook_id: HookID) -> bool {
|
||||
self.hooks.remove(&hook_id).is_some()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
mod hooks;
|
||||
mod routing;
|
||||
|
||||
pub use hooks::HookID;
|
||||
|
||||
use alloc::{boxed::Box, vec::Vec};
|
||||
|
||||
use crate::protocol::{ConnectionSet, HookMap, Leaf, Packet, Path, RouteMap};
|
||||
|
||||
pub struct Endpoint {
|
||||
// This endpoint's identifier
|
||||
pub id: u32,
|
||||
|
||||
// A counter that creates unique hook IDs.
|
||||
// TODO: Randomize the hooks for more obfuscation
|
||||
pub(crate) last_hook: u16,
|
||||
|
||||
// Absolute path for this node. Must be set by some leaf
|
||||
pub path: Path,
|
||||
pub leaves: Vec<Box<dyn Leaf>>,
|
||||
|
||||
// Map of connections so that we can know what is connected
|
||||
// and which endpoints are authorities
|
||||
pub connections: ConnectionSet,
|
||||
|
||||
// Local list of hooks.
|
||||
pub(crate) hooks: HookMap,
|
||||
|
||||
// Map of endpoints to packet queues
|
||||
pub(crate) inbound: RouteMap,
|
||||
pub(crate) outbound: RouteMap,
|
||||
}
|
||||
|
||||
impl Endpoint {
|
||||
pub fn new(id: u32, leaves: Vec<Box<dyn Leaf>>) -> Self {
|
||||
Self {
|
||||
id,
|
||||
// Init the hook at 0, which will increment
|
||||
last_hook: 0,
|
||||
|
||||
// Set the current path as an empty vec
|
||||
path: Vec::new(),
|
||||
leaves,
|
||||
hooks: HookMap::new(),
|
||||
connections: ConnectionSet::new(),
|
||||
inbound: RouteMap::new(),
|
||||
outbound: RouteMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Pass the endpoint state into all of the leaves
|
||||
pub fn update(&mut self) {
|
||||
// Grab the leaf vec temporarily so that we can iter over self
|
||||
// Apparently this only swaps out pointers
|
||||
let mut leaves = core::mem::take(&mut self.leaves);
|
||||
|
||||
for leaf in leaves.iter_mut() {
|
||||
leaf.update(self);
|
||||
}
|
||||
|
||||
self.leaves = leaves;
|
||||
}
|
||||
|
||||
/// Run a function over all inbound packets with some ID then clear it.
|
||||
pub fn take_inbound_clear<F>(&mut self, path: u32, f: F)
|
||||
where
|
||||
F: FnMut(&Packet),
|
||||
{
|
||||
Self::take_clear(path, f, &mut self.inbound);
|
||||
}
|
||||
|
||||
/// Drain inbound packets for `path` that match `predicate` and preserve the rest.
|
||||
///
|
||||
/// Generated leaf dispatch uses this instead of [`Self::take_inbound_clear`] so
|
||||
/// one leaf can consume only its procedure or session packets without stealing
|
||||
/// traffic intended for another leaf. Matching packets are passed by value because
|
||||
/// most handlers need to move payload bytes into application state; unmatched
|
||||
/// packets are reinserted in their original FIFO order.
|
||||
pub fn take_inbound_matching<P, F>(&mut self, path: u32, mut predicate: P, mut f: F)
|
||||
where
|
||||
P: FnMut(&Packet) -> bool,
|
||||
F: FnMut(Packet),
|
||||
{
|
||||
let Some(mut queue) = self.inbound.remove(&path) else {
|
||||
return;
|
||||
};
|
||||
|
||||
let mut unmatched = Vec::new();
|
||||
|
||||
while let Some(packet) = queue.pop_front() {
|
||||
if predicate(&packet) {
|
||||
f(packet);
|
||||
} else {
|
||||
unmatched.push(packet);
|
||||
}
|
||||
}
|
||||
|
||||
if !unmatched.is_empty() {
|
||||
self.inbound.entry(path).or_default().extend(unmatched);
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a function over all outbound packets with some ID then clear it.
|
||||
pub fn take_outbound_clear<F>(&mut self, path: u32, f: F)
|
||||
where
|
||||
F: FnMut(&Packet),
|
||||
{
|
||||
Self::take_clear(path, f, &mut self.outbound);
|
||||
}
|
||||
|
||||
fn take_clear<F>(path: u32, mut f: F, queue: &mut RouteMap)
|
||||
where
|
||||
F: FnMut(&Packet),
|
||||
{
|
||||
if let Some(queue) = queue.get_mut(&path) {
|
||||
for packet in queue.iter() {
|
||||
f(packet);
|
||||
}
|
||||
|
||||
queue.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
use crate::protocol::{Endpoint, EndpointError, Packet, RouteDirection};
|
||||
|
||||
impl Endpoint {
|
||||
/// Register an inbound packet from legacy trusted code.
|
||||
///
|
||||
/// Transports should prefer [`Self::add_inbound_from`] because peer-bound hook
|
||||
/// validation needs to know which adjacent endpoint supplied the bytes. This
|
||||
/// method keeps the old trusted in-process path small: it derives path direction,
|
||||
/// forwards or delivers the packet, and only checks that upward hooks exist.
|
||||
pub fn add_inbound(&mut self, packet: Packet) -> Result<(), EndpointError> {
|
||||
self.route_trusted_packet(packet)
|
||||
}
|
||||
|
||||
/// Register an inbound packet received from `remote_id` and route it locally.
|
||||
///
|
||||
/// Packets from a parent are downward traffic and pave return hooks when
|
||||
/// `end_hook` is false. Packets from a child are upward traffic and must match an
|
||||
/// already-paved hook for that exact child before they can move farther upward.
|
||||
pub fn add_inbound_from(
|
||||
&mut self,
|
||||
remote_id: u32,
|
||||
packet: Packet,
|
||||
) -> Result<(), EndpointError> {
|
||||
self.ensure_path_is_set()?;
|
||||
|
||||
let inbound_direction = self.inbound_direction_from_peer(remote_id)?;
|
||||
|
||||
if packet.path == self.path {
|
||||
return match inbound_direction {
|
||||
RouteDirection::Downward => self.deliver_local_downward(packet, remote_id),
|
||||
RouteDirection::Upward => self.deliver_local_upward(packet, remote_id),
|
||||
};
|
||||
}
|
||||
|
||||
if packet.path.starts_with(&self.path) {
|
||||
self.ensure_inbound_direction(remote_id, inbound_direction, RouteDirection::Downward)?;
|
||||
let next_hop = self.immediate_child_hop(&packet)?;
|
||||
return self.route_downward(packet, next_hop);
|
||||
}
|
||||
|
||||
if self.path.starts_with(&packet.path) {
|
||||
self.ensure_inbound_direction(remote_id, inbound_direction, RouteDirection::Upward)?;
|
||||
let next_hop = self.parent_hop()?;
|
||||
return self.route_upward(packet, next_hop, Some(remote_id));
|
||||
}
|
||||
|
||||
Err(EndpointError::DestinationOutsideLocalTree)
|
||||
}
|
||||
|
||||
/// Register an outbound packet produced locally and route it to the next queue.
|
||||
pub fn add_outbound(&mut self, packet: Packet) -> Result<(), EndpointError> {
|
||||
self.ensure_path_is_set()?;
|
||||
|
||||
if packet.path == self.path {
|
||||
return self.deliver_local(packet);
|
||||
}
|
||||
|
||||
if packet.path.starts_with(&self.path) {
|
||||
let next_hop = self.immediate_child_hop(&packet)?;
|
||||
return self.route_downward(packet, next_hop);
|
||||
}
|
||||
|
||||
if self.path.starts_with(&packet.path) {
|
||||
let next_hop = self.parent_hop()?;
|
||||
return self.route_upward(packet, next_hop, Some(next_hop));
|
||||
}
|
||||
|
||||
Err(EndpointError::DestinationOutsideLocalTree)
|
||||
}
|
||||
|
||||
/// Routes a trusted packet without transport-peer direction metadata.
|
||||
///
|
||||
/// This intentionally does not create local hooks on local delivery because the
|
||||
/// endpoint cannot know whether the packet came from a parent or child. Transit
|
||||
/// routing still maintains hook state where path direction is unambiguous.
|
||||
fn route_trusted_packet(&mut self, packet: Packet) -> Result<(), EndpointError> {
|
||||
self.ensure_path_is_set()?;
|
||||
|
||||
if packet.path == self.path {
|
||||
return self.deliver_local(packet);
|
||||
}
|
||||
|
||||
if packet.path.starts_with(&self.path) {
|
||||
let next_hop = self.immediate_child_hop(&packet)?;
|
||||
return self.route_downward(packet, next_hop);
|
||||
}
|
||||
|
||||
if self.path.starts_with(&packet.path) {
|
||||
let next_hop = self.parent_hop()?;
|
||||
return self.route_upward(packet, next_hop, None);
|
||||
}
|
||||
|
||||
Err(EndpointError::DestinationOutsideLocalTree)
|
||||
}
|
||||
|
||||
/// Delivers a packet to local leaves without changing hook state.
|
||||
fn deliver_local(&mut self, packet: Packet) -> Result<(), EndpointError> {
|
||||
let local_id = self.local_id()?;
|
||||
self.inbound.entry(local_id).or_default().push_back(packet);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Delivers parent-originated traffic locally and applies downward hook policy.
|
||||
fn deliver_local_downward(&mut self, packet: Packet, peer: u32) -> Result<(), EndpointError> {
|
||||
let hook_id = packet.hook_id;
|
||||
let end_hook = packet.end_hook;
|
||||
|
||||
self.deliver_local(packet)?;
|
||||
self.apply_downward_hook_lifecycle(hook_id, end_hook, peer);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Delivers child-originated traffic locally after validating its return hook.
|
||||
fn deliver_local_upward(&mut self, packet: Packet, peer: u32) -> Result<(), EndpointError> {
|
||||
let hook_id = packet.hook_id;
|
||||
let end_hook = packet.end_hook;
|
||||
|
||||
self.ensure_hook_peer(hook_id, peer)?;
|
||||
self.deliver_local(packet)?;
|
||||
self.apply_upward_hook_lifecycle(hook_id, end_hook);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Forwards a packet to a child and applies downward hook lifecycle rules.
|
||||
fn route_downward(&mut self, packet: Packet, next_hop: u32) -> Result<(), EndpointError> {
|
||||
let hook_id = packet.hook_id;
|
||||
let end_hook = packet.end_hook;
|
||||
|
||||
self.ensure_registered_connection(next_hop, RouteDirection::Downward)?;
|
||||
self.outbound.entry(next_hop).or_default().push_back(packet);
|
||||
self.apply_downward_hook_lifecycle(hook_id, end_hook, next_hop);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Forwards a packet toward the parent after validating hook state.
|
||||
///
|
||||
/// `actual_peer` is `None` only for legacy trusted inbound routing where the
|
||||
/// transport source is unknown; in that mode the endpoint can check that a hook
|
||||
/// exists but cannot enforce peer ownership.
|
||||
fn route_upward(
|
||||
&mut self,
|
||||
packet: Packet,
|
||||
next_hop: u32,
|
||||
actual_peer: Option<u32>,
|
||||
) -> Result<(), EndpointError> {
|
||||
let hook_id = packet.hook_id;
|
||||
let end_hook = packet.end_hook;
|
||||
|
||||
self.ensure_upward_hook_peer(hook_id, actual_peer)?;
|
||||
self.ensure_registered_connection(next_hop, RouteDirection::Upward)?;
|
||||
self.outbound.entry(next_hop).or_default().push_back(packet);
|
||||
self.apply_upward_hook_lifecycle(hook_id, end_hook);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns this endpoint's final path segment for local queueing.
|
||||
fn local_id(&self) -> Result<u32, EndpointError> {
|
||||
self.path
|
||||
.last()
|
||||
.copied()
|
||||
.ok_or(EndpointError::EndpointPathUnset)
|
||||
}
|
||||
|
||||
/// Returns the child that should receive a downward packet next.
|
||||
fn immediate_child_hop(&self, packet: &Packet) -> Result<u32, EndpointError> {
|
||||
packet
|
||||
.path
|
||||
.get(self.path.len())
|
||||
.copied()
|
||||
.ok_or(EndpointError::DestinationOutsideLocalTree)
|
||||
}
|
||||
|
||||
/// Returns the direct parent next hop for upward routing.
|
||||
fn parent_hop(&self) -> Result<u32, EndpointError> {
|
||||
let parent_index = self
|
||||
.path
|
||||
.len()
|
||||
.checked_sub(2)
|
||||
.ok_or(EndpointError::MissingParentRoute)?;
|
||||
|
||||
Ok(self.path[parent_index])
|
||||
}
|
||||
|
||||
/// Reject routing before path-relative decisions when no absolute path is known.
|
||||
///
|
||||
/// This preserves the current runtime sentinel where an empty path means the
|
||||
/// endpoint has not been attached to the tree yet.
|
||||
fn ensure_path_is_set(&self) -> Result<(), EndpointError> {
|
||||
if self.path.is_empty() {
|
||||
Err(EndpointError::EndpointPathUnset)
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Derives packet direction from a registered inbound adjacent peer.
|
||||
fn inbound_direction_from_peer(&self, remote_id: u32) -> Result<RouteDirection, EndpointError> {
|
||||
let is_upstream = self.connections.contains(&(remote_id, true));
|
||||
let is_downstream = self.connections.contains(&(remote_id, false));
|
||||
|
||||
match (is_upstream, is_downstream) {
|
||||
(true, false) => Ok(RouteDirection::Downward),
|
||||
(false, true) => Ok(RouteDirection::Upward),
|
||||
(false, false) => Err(EndpointError::UnknownConnection { remote_id }),
|
||||
(true, true) => Err(EndpointError::AmbiguousConnection { remote_id }),
|
||||
}
|
||||
}
|
||||
|
||||
/// Rejects inbound packets whose path-derived direction contradicts the connection.
|
||||
fn ensure_inbound_direction(
|
||||
&self,
|
||||
remote_id: u32,
|
||||
expected: RouteDirection,
|
||||
actual: RouteDirection,
|
||||
) -> Result<(), EndpointError> {
|
||||
if expected == actual {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(EndpointError::InboundDirectionMismatch {
|
||||
remote_id,
|
||||
expected,
|
||||
actual,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify that the derived adjacent endpoint is registered in this direction.
|
||||
///
|
||||
/// The current connection table stores direction as a boolean. Keeping the bool
|
||||
/// conversion here confines that legacy representation to one place in routing.
|
||||
fn ensure_registered_connection(
|
||||
&self,
|
||||
next_hop: u32,
|
||||
direction: RouteDirection,
|
||||
) -> Result<(), EndpointError> {
|
||||
let is_upward = matches!(direction, RouteDirection::Upward);
|
||||
|
||||
if self.connections.contains(&(next_hop, is_upward)) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(EndpointError::MissingConnection {
|
||||
next_hop,
|
||||
direction,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Validates hook state for upward routing.
|
||||
fn ensure_upward_hook_peer(
|
||||
&self,
|
||||
hook_id: u16,
|
||||
actual_peer: Option<u32>,
|
||||
) -> Result<(), EndpointError> {
|
||||
if let Some(actual_peer) = actual_peer {
|
||||
self.ensure_hook_peer(hook_id, actual_peer)
|
||||
} else if self.has_hook(hook_id) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(EndpointError::UnknownHook { hook_id })
|
||||
}
|
||||
}
|
||||
|
||||
/// Applies hook state for successfully routed downward packets.
|
||||
fn apply_downward_hook_lifecycle(&mut self, hook_id: u16, end_hook: bool, peer: u32) {
|
||||
if end_hook {
|
||||
self.close_hook(hook_id);
|
||||
} else {
|
||||
self.open_hook(hook_id, peer);
|
||||
}
|
||||
}
|
||||
|
||||
/// Applies hook cleanup for successfully routed upward final packets.
|
||||
fn apply_upward_hook_lifecycle(&mut self, hook_id: u16, end_hook: bool) {
|
||||
if end_hook {
|
||||
self.close_hook(hook_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user