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

feat: complete protocol spec and initial implementation

Browse Source 
- Write PROTOCOL.md with full wire format spec and 8 real-world scenario
  analyses (reconnect, multi-operator, large files, AV evasion, router crash,
  malformed packets, future pivoting)

- Rewrite workspace structure:
  - unshell lib: protocol types (PacketHeader, TreeRequest/Response,
    HandshakeMessage/Ack), Transport trait, TcpTransport, Tree routing
  - ush-router: router binary with per-node threads, NodeRegistry with
    longest-prefix path matching, packet relay
  - ush-payload: implant binary with reconnect loop, module tree, InfoModule
  - ush-cli: operator REPL with rustyline, session management, command parser

- Protocol design: two-part rkyv frame [header][payload]; router reads only
  header for routing, payload bytes forwarded opaque

- All code documented with doc comments and examples
- Zero warnings, zero errors across entire workspace
- 32 tests pass (unit tests for tree routing, TCP transport, framing,
  command parsing, node registry)
This commit is contained in:
Michael Mikovsky
2026-04-20 23:38:02 -06:00
parent 959ea469a8
commit fcb3b2be17
30 changed files with 4623 additions and 658 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+46 -4
View File
@@ -1,11 +1,53 @@
#![no_main]
#![no_std]
//! # UnShell Core Library
//!
//! This crate provides the core building blocks for the UnShell C2 framework:
//!
//! - **[`protocol`]** — wire types: `PacketHeader`, `TreeRequest`, `TreeResponse`,
//! `HandshakeMessage`, `HandshakeAck`, and associated enums.
//! - **[`transport`]** — the `Transport` trait and its TCP implementation.
//! - **[`tree`]** — the `Tree` and `Endpoint` abstractions for module dispatch.
//! - **[`logger`]** — lightweight logging (no dependency on `std::io`).
//!
//! ## `no_std` Compatibility
//!
//! This crate is `no_std` but requires `alloc`. It can be used in the payload
//! binary which runs without a full standard library.
//!
//! Binaries that have `std` available (the router, the CLI) can also use this
//! crate; they simply get `alloc` types backed by the system allocator.
//!
//! ## Architecture
//!
//! ```text
//! ┌────────────────────────────────────────────────────────────────┐
//! │ Router / Relay │
//! │ Reads PacketHeader → longest-prefix routes to node │
//! │ Payload bytes forwarded opaque │
//! └───────────┬─────────────────────────┬──────────────────────────┘
//! │ TCP │ TCP
//! ┌────────▼────────┐ ┌─────────▼──────────────────────────┐
//! │ Operator Node │ │ Payload Node(s) │
//! │ (ush-cli) │ │ Local Tree + Endpoint modules │
//! │ Interactive │ │ Reverse-connects to router │
//! │ REPL │ │ Recv loop → dispatch → respond │
//! └─────────────────┘ └─────────────────────────────────────┘
//! ```
//!
//! For the full protocol specification, see `PROTOCOL.md` in the repository root.
// Enable std when the `tcp` feature is active (TCP transport requires it).
// Without tcp, we stay fully no_std for bare-metal payload targets.
#![cfg_attr(not(feature = "tcp"), no_std)]
// no_main is only applied in non-test builds.
// The test harness generates its own main function, so we must NOT suppress it.
#![cfg_attr(not(test), no_main)]
extern crate alloc;
pub mod logger;
pub mod protocol;
pub mod transport;
pub mod tree;
// Re-exports
// pub use serde_json::{Value, json};
// Re-export the obfuscation crate so payloads only need to depend on `unshell`.
pub use ush_obfuscate as obfuscate;
src/logger/log_disabled.rs
-6
View File
@@ -1,6 +0,0 @@
// Macros that are used that just drop the inside variables
#[macro_export]
macro_rules! log {
($level:expr, $fmt:tt) => {{}};
($level:expr, $fmt:tt, $($arg:expr),*) => {{}};
}
src/logger/log_enabled.rs
-44
View File
@@ -1,44 +0,0 @@
#[macro_export]
macro_rules! log {
($level:expr, $fmt:tt) => {{
use $crate::obfuscate;
let log_result = obfuscate::sym_format!($fmt);
$crate::logger::add_record(
$level,
#[cfg(feature = "log_debug")]
Some(String::from(obfuscate::file_symbol!())),
#[cfg(not(feature = "log_debug"))]
None,
#[cfg(feature = "log_debug")]
Some(std::time::SystemTime::now()),
#[cfg(not(feature = "log_debug"))]
None,
log_result
);
}};
($level:expr, $fmt:tt, $($arg:expr),*) => {{
use $crate::obfuscate;
let log_result = obfuscate::sym_format!($fmt, $($arg),*);
$crate::logger::add_record(
$level,
#[cfg(feature = "log_debug")]
Some(String::from(obfuscate::file_symbol!())),
#[cfg(not(feature = "log_debug"))]
None,
#[cfg(feature = "log_debug")]
Some(std::time::SystemTime::now()),
#[cfg(not(feature = "log_debug"))]
None,
log_result
);
}};
}
src/logger/mod.rs
+283 -67
View File
@@ -1,115 +1,331 @@
// Choose if the macros are enabled based on the feature setting
#[cfg(feature = "log")]
mod log_enabled;
//! # Logger Module
//!
//! A lightweight, no_std-compatible logging system.
//!
//! ## Usage
//!
//! ```rust
//! use unshell::{info, warn, error};
//! use unshell::logger::Logger;
//!
//! // Uses the default (no-op) logger until one is installed.
//! info!("Starting up");
//! warn!("Something is off");
//! error!("Critical failure");
//! ```
//!
//! ## Installing a logger
//!
//! Call [`set_logger`] with any type that implements [`Logger`]:
//!
//! ```rust,no_run
//! use unshell::logger::{Logger, LogLevel, Record, set_logger};
//!
//! struct StdoutLogger;
//! impl Logger for StdoutLogger {
//! fn log(&self, record: &Record<'_>) {
//! // In a no_std environment you would use the `unix-print` crate
//! // or write to a pre-opened file descriptor.
//! let _ = record; // placeholder
//! }
//! }
//!
//! static MY_LOGGER: StdoutLogger = StdoutLogger;
//! set_logger(&MY_LOGGER);
//! ```
//!
//! ## Thread safety
//!
//! The global logger pointer is set **once at startup**, before any threads
//! are spawned. After that, it is only read (never written). This is safe
//! because:
//!
//! 1. The payload is single-threaded.
//! 2. The router and CLI set the logger before spawning node threads.
//!
//! If you need to change the logger after threads start, synchronise access
//! with a `Mutex` or an atomic pointer in your logger implementation.
#[cfg(not(feature = "log"))]
mod log_disabled;
// ---------------------------------------------------------------------------
// Log levels
// ---------------------------------------------------------------------------
mod pretty_logger;
use alloc::boxed::Box;
use alloc::string::String;
pub use pretty_logger::PrettyLogger;
pub use pretty_logger::log;
pub static mut IS_DEFAULT_LOGGER: bool = true;
static mut LOGGER: &dyn Logger = &DefaultLogger;
#[derive(Debug)]
/// The severity level of a log record.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum LogLevel {
/// Verbose diagnostic information.
Debug,
/// Normal operational messages.
Info,
/// Something unexpected happened but execution can continue.
Warn,
/// A serious error occurred.
Error,
}
#[derive(Debug)]
pub struct Record {
log_level: LogLevel,
location: Option<String>,
// line: u32,
time: Option<u64>,
message: String,
}
pub trait Logger {
fn log(&self, log: Record);
}
struct DefaultLogger;
impl Logger for DefaultLogger {
fn log(&self, _: Record) {}
}
#[allow(unused_variables)]
pub fn set_logger_box(logger: Box<dyn Logger>) {
#[cfg(feature = "log")]
unsafe {
LOGGER = Box::leak(logger);
IS_DEFAULT_LOGGER = false;
impl LogLevel {
/// Short uppercase label, suitable for log line prefixes.
///
/// # Example
///
/// ```rust
/// use unshell::logger::LogLevel;
/// assert_eq!(LogLevel::Info.as_str(), "INFO");
/// ```
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Debug => "DEBUG",
Self::Info => "INFO",
Self::Warn => "WARN",
Self::Error => "ERROR",
}
}
}
// ---------------------------------------------------------------------------
// Log record
// ---------------------------------------------------------------------------
/// A single log entry passed to a [`Logger`].
///
/// Borrows from the call site to avoid heap allocation on the hot path.
pub struct Record<'a> {
/// Severity level.
pub level: LogLevel,
/// The log message.
pub message: &'a str,
/// Source file, if available (e.g. `file!()`).
pub file: Option<&'static str>,
/// Source line number, if available (e.g. `line!()`).
pub line: Option<u32>,
}
// ---------------------------------------------------------------------------
// Logger trait
// ---------------------------------------------------------------------------
/// A sink for log records.
///
/// Implement this to direct log output wherever you want (stdout, a file,
/// a TCP connection, a memory buffer for tests).
pub trait Logger: Sync {
/// Receive and process a log record.
fn log(&self, record: &Record<'_>);
}
// ---------------------------------------------------------------------------
// Global logger state
// ---------------------------------------------------------------------------
/// The no-op logger used before any logger is installed.
struct NullLogger;
impl Logger for NullLogger {
fn log(&self, _record: &Record<'_>) {}
}
/// The global logger pointer.
///
/// Written once at startup via [`set_logger`], then only read.
/// # Safety
/// This is `static mut` to avoid a dependency on synchronisation primitives
/// in a no_std context. It is safe as long as `set_logger` is called before
/// any threads are spawned (see module-level docs).
static mut GLOBAL_LOGGER: &dyn Logger = &NullLogger;
/// Install a new global logger.
///
/// Must be called **before** spawning any threads. After this call, all
/// `info!`, `warn!`, `error!`, and `debug!` macros route to this logger.
///
/// # Safety
///
/// This function writes to a `static mut`. It is safe when called exactly
/// once at program startup before any other threads exist.
///
/// # Example
///
/// ```rust,no_run
/// use unshell::logger::{Logger, Record, set_logger};
///
/// static MY_LOGGER: MyLogger = MyLogger;
/// set_logger(&MY_LOGGER);
///
/// # struct MyLogger;
/// # impl Logger for MyLogger { fn log(&self, _: &Record<'_>) {} }
/// ```
pub fn set_logger(logger: &'static dyn Logger) {
// SAFETY: called once at startup before any threads are spawned.
#[allow(static_mut_refs)]
unsafe {
LOGGER = logger;
IS_DEFAULT_LOGGER = false;
GLOBAL_LOGGER = logger;
}
}
pub fn add_record(
log_level: LogLevel,
location: Option<String>,
time: Option<u64>,
message: String,
) {
logger().log(Record {
log_level,
location,
time,
/// Return a reference to the currently installed logger.
///
/// Used internally by the logging macros.
#[must_use]
pub fn global_logger() -> &'static dyn Logger {
// SAFETY: GLOBAL_LOGGER is only written once (at startup) and is
// read-only thereafter. No data race is possible.
#[allow(static_mut_refs)]
unsafe {
GLOBAL_LOGGER
}
}
/// Log a record through the global logger.
///
/// This is the low-level function called by the macros. Prefer using the
/// `info!`, `warn!`, `error!`, and `debug!` macros directly.
pub fn log(level: LogLevel, message: &str, file: Option<&'static str>, line: Option<u32>) {
global_logger().log(&Record {
level,
message,
file,
line,
});
}
pub fn logger() -> &'static dyn Logger {
unsafe { LOGGER }
// ---------------------------------------------------------------------------
// A minimal stdout logger for use in std binaries (router, CLI)
// ---------------------------------------------------------------------------
/// A simple logger that prints to stderr.
///
/// Suitable for the router and operator CLI binaries.
/// Do not use in the payload binary (which may not have stderr available).
///
/// # Example
///
/// ```rust,no_run
/// use unshell::logger::{StderrLogger, set_logger};
///
/// static LOGGER: StderrLogger = StderrLogger::new(unshell::logger::LogLevel::Info);
/// set_logger(&LOGGER);
/// ```
pub struct StderrLogger {
/// Minimum level to log. Records below this level are discarded.
min_level: LogLevel,
}
#[allow(dead_code, improper_ctypes_definitions)]
pub type SetupLogger = extern "C" fn(logger: &'static dyn Logger);
#[unsafe(no_mangle)]
#[allow(improper_ctypes_definitions)]
pub extern "C" fn setup_logger(logger: &'static dyn Logger) {
set_logger(logger);
impl StderrLogger {
/// Create a new `StderrLogger` that logs records at `min_level` and above.
///
/// # Example
///
/// ```rust
/// use unshell::logger::{StderrLogger, LogLevel};
/// let logger = StderrLogger::new(LogLevel::Info);
/// ```
#[must_use]
pub const fn new(min_level: LogLevel) -> Self {
Self { min_level }
}
}
// Macro Definitions
impl Logger for StderrLogger {
fn log(&self, record: &Record<'_>) {
if record.level < self.min_level {
return;
}
// eprintln! and String require std (available only with the `tcp` feature).
// In no_std builds this method is a no-op. The payload uses a different
// logger (or the null logger) in no_std contexts.
#[cfg(feature = "tcp")]
{
use alloc::string::String;
let location = match (record.file, record.line) {
(Some(f), Some(l)) => {
let mut s = String::from(f);
s.push(':');
s.push_str(&format!("{l}"));
s
}
_ => String::new(),
};
if location.is_empty() {
eprintln!("[{}] {}", record.level.as_str(), record.message);
} else {
eprintln!("[{}] {} - {}", record.level.as_str(), record.message, location);
}
}
}
}
// ---------------------------------------------------------------------------
// Logging macros
// ---------------------------------------------------------------------------
/// Log at [`LogLevel::Debug`] level.
///
/// ```rust
/// use unshell::debug;
/// debug!("loop iteration {}", 42);
/// ```
#[macro_export]
macro_rules! debug {
($($arg:tt)*) => {
$crate::log!($crate::logger::LogLevel::Debug, $($arg)*)
$crate::logger::log(
$crate::logger::LogLevel::Debug,
&format!($($arg)*),
Some(file!()),
Some(line!()),
)
};
}
/// Log at [`LogLevel::Info`] level.
///
/// ```rust
/// use unshell::info;
/// info!("server started on port {}", 9000);
/// ```
#[macro_export]
macro_rules! info {
($($arg:tt)*) => {
$crate::log!($crate::logger::LogLevel::Info, $($arg)*)
$crate::logger::log(
$crate::logger::LogLevel::Info,
&format!($($arg)*),
Some(file!()),
Some(line!()),
)
};
}
/// Log at [`LogLevel::Warn`] level.
///
/// ```rust
/// use unshell::warn;
/// warn!("unexpected path: {}", "/unknown");
/// ```
#[macro_export]
macro_rules! warn {
($($arg:tt)*) => {
$crate::log!($crate::logger::LogLevel::Warn, $($arg)*)
$crate::logger::log(
$crate::logger::LogLevel::Warn,
&format!($($arg)*),
Some(file!()),
Some(line!()),
)
};
}
/// Log at [`LogLevel::Error`] level.
///
/// ```rust
/// use unshell::error;
/// error!("connection failed: {}", "timeout");
/// ```
#[macro_export]
macro_rules! error {
($($arg:tt)*) => {
$crate::log!($crate::logger::LogLevel::Error, $($arg)*)
$crate::logger::log(
$crate::logger::LogLevel::Error,
&format!($($arg)*),
Some(file!()),
Some(line!()),
)
};
}
src/logger/pretty_logger.rs
-80
View File
@@ -1,80 +0,0 @@
use alloc::{boxed::Box, format};
use crate::logger::{LogLevel, Logger, Record};
pub struct PrettyLogger {
output: Option<Box<dyn Fn(&Record)>>,
}
impl Logger for PrettyLogger {
fn log(&self, message: Record) {
if let Some(ref func) = self.output {
(*func)(&message)
}
log(&message);
}
}
pub fn log(message: &Record) {
static DEBUG_COLOR: &str = "\x1b[36m";
static INFO_COLOR: &str = "\x1b[32m";
static WARN_COLOR: &str = "\x1b[33m";
static ERROR_COLOR: &str = "\x1b[31m";
let log_level = match message.log_level {
LogLevel::Debug => format!("{DEBUG_COLOR}DBUG"),
LogLevel::Info => format!("{INFO_COLOR}INFO"),
LogLevel::Warn => format!("{WARN_COLOR}WARN"),
LogLevel::Error => format!("{ERROR_COLOR}ERR!"),
};
match (message.time, &message.location) {
(None, None) => {
static WHITE: &str = "\x1b[97m";
unix_print::unix_println!("{} {WHITE}{}", log_level, message.message);
}
#[cfg(feature = "log_debug")]
(Some(time), Some(location)) => {
use chrono::{DateTime, Utc};
let date: DateTime<Utc> = time.into();
static WHITE: &str = "\x1b[97m";
static OFF_WHITE: &str = "\x1b[37m";
static TIME_COLOR: &str = "\x1b[36m";
static GREY: &str = "\x1b[90m";
unix_print::unix_println!(
"{OFF_WHITE}[{TIME_COLOR}{}{OFF_WHITE}] {} {WHITE}{} {GREY}{}{WHITE}",
date,
log_level,
message.message,
location
);
}
_ => unreachable!("Invalid log configuration"),
}
}
impl PrettyLogger {
pub fn init() {
if unsafe { crate::logger::IS_DEFAULT_LOGGER } {
crate::logger::set_logger_box(Box::new(PrettyLogger { output: None }));
}
}
pub fn init_output<T>(output: T)
where
T: Fn(&Record) + 'static,
{
if !unsafe { crate::logger::IS_DEFAULT_LOGGER } {
crate::logger::set_logger_box(Box::new(PrettyLogger {
output: Some(Box::new(output)),
}));
}
}
}
src/protocol/content.rs
+59
View File
@@ -0,0 +1,59 @@
//! # Content Type Constants
//!
//! Content types describe how to interpret the `data` field of a
//! [`TreeRequest`](super::TreeRequest) or [`TreeResponse`](super::TreeResponse).
//!
//! They follow a `"namespace/TypeName"` convention, similar to MIME types.
//!
//! ## Built-in types
//!
//! | Constant | Value | Meaning |
//! |---|---|---|
//! | [`NONE`] | `"core/None"` | No data (empty payload) |
//! | [`UTF8_STRING`] | `"core/Utf8String"` | Raw UTF-8 string |
//! | [`BYTES`] | `"core/Bytes"` | Raw bytes (no specific interpretation) |
//! | [`PROCEDURE_LIST`] | `"core/ProcedureList"` | rkyv-serialised `Vec<ProcedureDescriptor>` |
//!
//! ## Custom types
//!
//! Module authors should prefix with their module name:
//!
//! ```rust
//! const MY_TYPE: &str = "mymodule/MyType";
//! ```
/// No data. Use for requests/responses that carry no payload.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{TreeRequest, RequestType, content};
///
/// // A ping-style read with no payload
/// let req = TreeRequest {
/// request_id: 1,
/// request_type: RequestType::Read,
/// content_type: content::NONE.into(),
/// data: Vec::new(),
/// };
/// ```
pub const NONE: &str = "core/None";
/// A raw UTF-8 string.
///
/// The `data` field contains the string's bytes (no null terminator, no length prefix).
pub const UTF8_STRING: &str = "core/Utf8String";
/// Raw bytes with no specific interpretation.
pub const BYTES: &str = "core/Bytes";
/// A rkyv-serialised `Vec<ProcedureDescriptor>`.
///
/// Used in responses to [`RequestType::GetProcedures`](super::RequestType::GetProcedures).
pub const PROCEDURE_LIST: &str = "core/ProcedureList";
/// Shell command output: UTF-8 stdout and stderr combined.
pub const SHELL_OUTPUT: &str = "shell/Output";
/// Raw file contents as bytes.
pub const FILE_BYTES: &str = "files/Bytes";
src/protocol/mod.rs
+40
View File
@@ -0,0 +1,40 @@
//! # Protocol Module
//!
//! All wire types used by the UnShell protocol.
//!
//! ## Module layout
//!
//! ```text
//! protocol/
//! mod.rs ← you are here; re-exports everything
//! types.rs ← PacketHeader, TreeRequest, TreeResponse, Handshake*
//! content.rs ← content-type string constants
//! ```
//!
//! ## Quick start
//!
//! ```rust
//! use unshell::protocol::{
//! PacketHeader, PacketType,
//! TreeRequest, RequestType,
//! content,
//! };
//!
//! let header = PacketHeader {
//! dst_path: "/agents/abc123/shell/exec".into(),
//! src_path: "/operator/sess1".into(),
//! packet_type: PacketType::Request,
//! };
//!
//! let request = TreeRequest {
//! request_id: 1,
//! request_type: RequestType::CallProcedure,
//! content_type: content::UTF8_STRING.into(),
//! data: b"ls -la".to_vec(),
//! };
//! ```
pub mod content;
mod types;
pub use types::*;
src/protocol/types.rs
+314
View File
@@ -0,0 +1,314 @@
//! # Protocol Wire Types
//!
//! All structs and enums that appear on the wire.
//!
//! ## Serialisation
//!
//! Every type here derives rkyv's `Archive`, `Serialize`, and `Deserialize`.
//! This means they can be serialised to a byte slice and deserialised back
//! with zero copying — the deserialised view (`Archived<T>`) reads directly
//! from the byte slice without allocating.
//!
//! ## Wire Frame Format
//!
//! Every packet on the wire uses a two-part frame:
//!
//! ```text
//! ┌──────────────────────────────────────────────────────────────────────┐
//! │ Part 1: Header │ Part 2: Payload │
//! │ [u32 big-endian length] │ [u32 big-endian length] │
//! │ [rkyv-serialised PacketHeader bytes] │ [rkyv payload bytes] │
//! └──────────────────────────────────────────┴───────────────────────────┘
//! ```
//!
//! The router reads only Part 1 to determine where to route the packet.
//! Part 2 is forwarded opaque (the router does not deserialise it).
use alloc::string::String;
use alloc::vec::Vec;
use rkyv::{Archive, Deserialize, Serialize};
// ---------------------------------------------------------------------------
// PacketHeader
// ---------------------------------------------------------------------------
/// The header prefixed to every packet on the wire.
///
/// The router reads ONLY this field to determine routing.
/// The payload body is opaque to the router.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{PacketHeader, PacketType};
///
/// let header = PacketHeader {
/// dst_path: "/agents/abc123/shell/exec".into(),
/// src_path: "/operator/sess1".into(),
/// packet_type: PacketType::Request,
/// };
/// ```
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
#[rkyv(derive(Debug))]
pub struct PacketHeader {
/// Destination path in the global tree.
///
/// The router does a longest-prefix match against registered node paths.
/// Example: `"/agents/abc123/shell/exec"`.
pub dst_path: String,
/// Source path of the sending node.
///
/// Used by the destination to route the response back.
/// Example: `"/operator/sess1"`.
pub src_path: String,
/// Discriminates between handshake messages and protocol messages.
pub packet_type: PacketType,
}
/// Discriminates the payload type.
///
/// The receiver uses this to know which type to deserialise the payload as.
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
#[rkyv(derive(Debug, PartialEq))]
pub enum PacketType {
/// Sent by a newly-connected node to register with the router.
Handshake,
/// Sent by the router acknowledging (or rejecting) a handshake.
HandshakeAck,
/// An application-level request (the primary protocol message).
Request,
/// An application-level response.
Response,
}
// ---------------------------------------------------------------------------
// Handshake
// ---------------------------------------------------------------------------
/// Sent by a node immediately after connecting to the router.
///
/// The router reads this to register the node in its routing table.
///
/// # Wire format
///
/// This struct is the payload part of a frame whose header has
/// `packet_type = PacketType::Handshake`. The `dst_path` in the header is
/// `"/router"` (the router's own registration endpoint).
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{HandshakeMessage, NodeType};
///
/// let msg = HandshakeMessage {
/// node_id: "abc123".into(),
/// node_type: NodeType::Payload,
/// registered_paths: vec!["/agents/abc123".into()],
/// platform: "linux-x86_64".into(),
/// };
/// ```
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
#[rkyv(derive(Debug))]
pub struct HandshakeMessage {
/// Node identifier.
///
/// For payloads: a base62 string baked at compile time.
/// For operator sessions: a random string generated on startup.
pub node_id: String,
/// Whether this node is a payload or an operator shell.
pub node_type: NodeType,
/// The path prefixes this node claims ownership of.
///
/// All sub-paths under these prefixes are owned by this node.
/// The router uses these for longest-prefix route matching.
///
/// Example: `["/agents/abc123"]`
pub registered_paths: Vec<String>,
/// Human-readable platform identifier for operator visibility.
///
/// Example: `"linux-x86_64"`, `"windows-x86_64"`, `"operator"`.
pub platform: String,
}
/// Sent by the router in response to a `HandshakeMessage`.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::HandshakeAck;
///
/// // Successful registration
/// let ack = HandshakeAck {
/// accepted: true,
/// assigned_base_path: "/agents/abc123".into(),
/// rejection_reason: None,
/// };
///
/// // Rejection (duplicate node ID)
/// let nack = HandshakeAck {
/// accepted: false,
/// assigned_base_path: String::new(),
/// rejection_reason: Some("duplicate_node_id".into()),
/// };
/// ```
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
#[rkyv(derive(Debug))]
pub struct HandshakeAck {
/// Whether the router accepted the registration.
pub accepted: bool,
/// The canonical base path assigned by the router.
///
/// Typically matches the first entry in `HandshakeMessage::registered_paths`.
/// Empty string if `accepted == false`.
pub assigned_base_path: String,
/// Human-readable rejection reason when `accepted == false`.
///
/// Known values: `"duplicate_node_id"`, `"invalid_path"`.
pub rejection_reason: Option<String>,
}
/// The type of node connecting to the router.
///
/// The `Router` variant is reserved for future multi-hop/pivoting support
/// and is not used in v1.
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
#[rkyv(derive(Debug, PartialEq))]
pub enum NodeType {
/// An implant running on a target machine.
Payload,
/// An operator's interactive shell session.
Operator,
// Router variant will be added when multi-hop/pivoting is implemented.
// Router,
}
// ---------------------------------------------------------------------------
// TreeRequest / TreeResponse
// ---------------------------------------------------------------------------
/// An application-level request sent from an operator to a payload module.
///
/// The request travels: operator → router → destination node.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{TreeRequest, RequestType, content};
///
/// // Ask a shell module to execute a command
/// let req = TreeRequest {
/// request_id: 42,
/// request_type: RequestType::CallProcedure,
/// content_type: content::UTF8_STRING.into(),
/// data: b"ls -la /tmp".to_vec(),
/// };
/// ```
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
#[rkyv(derive(Debug))]
pub struct TreeRequest {
/// Unique request ID generated by the sender.
///
/// The responder echoes this back in [`TreeResponse::request_id`].
/// This allows the sender to match responses to outstanding requests,
/// which matters when multiple requests are in-flight concurrently
/// (e.g., background sessions in the operator CLI).
pub request_id: u64,
/// The operation type.
pub request_type: RequestType,
/// Content-type describing how to interpret [`data`](Self::data).
///
/// Use the constants in [`content`](super::content) for the built-in types.
/// Custom module types should use the module name as namespace:
/// `"mymodule/MyType"`.
pub content_type: String,
/// Operation payload. Interpretation depends on `content_type`.
pub data: Vec<u8>,
}
/// The type of operation being requested.
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
#[rkyv(derive(Debug, PartialEq))]
pub enum RequestType {
/// Read a value at the target path.
Read = 0,
/// List available sub-paths and callable procedures at the target path.
GetProcedures = 1,
/// Write a value to the target path.
Write = 2,
/// Invoke a named procedure at the target path.
CallProcedure = 3,
}
/// An application-level response from a payload module back to the operator.
///
/// The response travels: payload → router → requesting operator.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{TreeResponse, ResponseStatus, content};
///
/// let resp = TreeResponse {
/// request_id: 42, // echoed from the corresponding TreeRequest
/// status: ResponseStatus::Ok,
/// content_type: content::UTF8_STRING.into(),
/// data: b"file1.txt\nfile2.txt\n".to_vec(),
/// };
/// ```
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
#[rkyv(derive(Debug))]
pub struct TreeResponse {
/// Echoed from the corresponding [`TreeRequest::request_id`].
pub request_id: u64,
/// Whether the operation succeeded.
pub status: ResponseStatus,
/// Content-type of the response data.
pub content_type: String,
/// Response payload. Empty if `status` is an error variant.
pub data: Vec<u8>,
}
/// Indicates the outcome of a [`TreeRequest`].
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
#[rkyv(derive(Debug, PartialEq))]
pub enum ResponseStatus {
/// The operation completed successfully.
Ok = 0,
/// The requested path does not exist at the destination node.
NoBranchError = 1,
/// The requested operation is not supported at this path.
UnsupportedOperation = 2,
/// The destination node encountered an internal error.
ExecutionError = 3,
/// The request payload was malformed or could not be deserialised.
ProtocolError = 4,
}
/// A descriptor for a callable procedure, returned by [`RequestType::GetProcedures`].
///
/// This is what fills the `data` field of a `TreeResponse` when the
/// request type is `GetProcedures` and `content_type` is `content::PROCEDURE_LIST`.
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
#[rkyv(derive(Debug))]
pub struct ProcedureDescriptor {
/// The name of the procedure (the path component after the module path).
///
/// Example: `"exec"` for the module at `/agents/abc123/shell/exec`.
pub name: String,
/// Human-readable description of what this procedure does.
pub description: String,
}
src/transport/mod.rs
+304
View File
@@ -0,0 +1,304 @@
//! # Transport Module
//!
//! The transport layer abstracts the network connection used to carry protocol packets.
//!
//! ## Module layout
//!
//! ```text
//! transport/
//! mod.rs ← you are here; Transport trait, TransportError, frame encoding
//! tcp.rs ← TcpTransport: Transport implemented for std::net::TcpStream
//! ```
//!
//! ## Design
//!
//! A `Transport` sends and receives complete logical packets. Each packet is
//! one `PacketHeader` + one opaque payload byte slice.
//!
//! Internally, implementations must use the two-part framing format:
//!
//! ```text
//! ┌──────────────────────────────────────────────────────────────────────┐
//! │ [u32 big-endian header_len][header bytes][u32 big-endian pay_len] │
//! │ [payload bytes] │
//! └──────────────────────────────────────────────────────────────────────┘
//! ```
//!
//! **IMPORTANT:** TCP is a stream protocol. A single `read()` call may return
//! fewer bytes than requested. All receive operations MUST loop until the
//! exact number of bytes has been read. The standard pattern is `read_exact()`.
//!
//! ## Size limits
//!
//! | Limit | Value | Reason |
//! |---|---|---|
//! | Max header bytes | 64 KB | Headers are always small; larger = bug or attack |
//! | Max payload bytes | 64 MB | Sufficient for most file transfers |
//!
//! ## Transport implementations
//!
//! | Type | Where | Description |
//! |---|---|---|
//! | [`tcp::TcpTransport`] | `transport/tcp.rs` | Standard TCP socket |
//!
//! Future additions: `HttpsTransport`, `IcmpTransport`, `OpenVpnTransport`.
extern crate alloc;
use alloc::vec::Vec;
#[allow(unused_imports)]
use alloc::vec;
use crate::protocol::PacketHeader;
/// TCP transport implementation.
///
/// Only available when the `tcp` feature is enabled (requires `std`).
/// Enable with `unshell = { features = ["tcp"] }` in your `Cargo.toml`.
#[cfg(feature = "tcp")]
pub mod tcp;
// ---------------------------------------------------------------------------
// Frame size limits
// ---------------------------------------------------------------------------
/// Maximum allowed size for a serialised `PacketHeader` (64 KB).
///
/// Headers should be tiny (< 200 bytes in practice). Anything larger suggests
/// either a bug in the sender or a malformed/malicious frame.
pub const MAX_HEADER_BYTES: usize = 64 * 1024;
/// Maximum allowed size for a packet payload (64 MB).
///
/// Sufficient for most file transfers without chunking.
/// Larger transfers will require the (not-yet-implemented) streaming extension.
pub const MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024;
// ---------------------------------------------------------------------------
// TransportError
// ---------------------------------------------------------------------------
/// Errors that can occur during [`Transport`] operations.
///
/// # Reconnect policy
///
/// When a payload receives [`TransportError::Disconnected`] or
/// [`TransportError::Io`], it should:
/// 1. Close the current transport.
/// 2. Wait 5 seconds.
/// 3. Attempt to create a new transport connection.
/// 4. Repeat indefinitely on failure.
///
/// The operator CLI exits on disconnect (the user restarts it manually).
#[derive(Debug)]
pub enum TransportError {
/// An I/O error from the underlying stream.
///
/// This includes partial writes, socket errors, and OS-level failures.
/// Only available when the `tcp` feature is enabled (requires std).
#[cfg(feature = "tcp")]
Io(std::io::Error),
/// The announced frame header length exceeds [`MAX_HEADER_BYTES`].
///
/// The connection should be closed immediately — the remote end is either
/// buggy or malicious. Do not allocate a buffer of the announced size.
///
/// Fields: `(announced_size, limit)`.
HeaderTooLarge(usize, usize),
/// The announced frame payload length exceeds [`MAX_PAYLOAD_BYTES`].
///
/// Fields: `(announced_size, limit)`.
PayloadTooLarge(usize, usize),
/// The remote end closed the connection cleanly (EOF).
///
/// This is not an error in the traditional sense. It means the other side
/// disconnected intentionally (e.g., payload restarted, operator exited).
Disconnected,
/// The received bytes could not be deserialised as a `PacketHeader`.
///
/// This indicates a protocol version mismatch or data corruption.
DeserialiseError,
}
#[cfg(feature = "tcp")]
impl core::fmt::Display for TransportError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Io(e) => write!(f, "transport I/O error: {e}"),
Self::HeaderTooLarge(got, max) => {
write!(f, "frame header too large: {got} bytes (limit: {max})")
}
Self::PayloadTooLarge(got, max) => {
write!(f, "frame payload too large: {got} bytes (limit: {max})")
}
Self::Disconnected => write!(f, "connection closed by remote"),
Self::DeserialiseError => write!(f, "failed to deserialise packet header"),
}
}
}
#[cfg(not(feature = "tcp"))]
impl core::fmt::Display for TransportError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::HeaderTooLarge(got, max) => {
write!(f, "frame header too large: {got} bytes (limit: {max})")
}
Self::PayloadTooLarge(got, max) => {
write!(f, "frame payload too large: {got} bytes (limit: {max})")
}
Self::Disconnected => write!(f, "connection closed by remote"),
Self::DeserialiseError => write!(f, "failed to deserialise packet header"),
}
}
}
#[cfg(feature = "tcp")]
impl From<std::io::Error> for TransportError {
fn from(e: std::io::Error) -> Self {
Self::Io(e)
}
}
// Implement std::error::Error so TransportError works with `?` in Box<dyn Error> contexts.
#[cfg(feature = "tcp")]
impl std::error::Error for TransportError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(e) => Some(e),
_ => None,
}
}
}
// ---------------------------------------------------------------------------
// Transport trait
// ---------------------------------------------------------------------------
/// A bidirectional framed transport.
///
/// Implementors handle the low-level byte transfer, including framing,
/// length prefixes, and the `read_exact` loop. The protocol layer above
/// sees complete logical packets (header + payload pairs).
///
/// # Contract
///
/// - `send` must write all bytes before returning `Ok(())`.
/// - `recv` must block until a complete header+payload pair is available.
/// - Both methods must use `read_exact`-style loops (never a single `read`).
/// - Frame size checks must be performed before any allocation.
///
/// # Example: implementing a custom transport
///
/// ```rust,no_run
/// use unshell::transport::{Transport, TransportError};
/// use unshell::protocol::PacketHeader;
///
/// struct MyTransport { /* ... */ }
///
/// impl Transport for MyTransport {
/// fn send(&mut self, header: &PacketHeader, payload: &[u8])
/// -> Result<(), TransportError>
/// {
/// // 1. Serialise header with rkyv
/// // 2. Write [u32 header_len][header bytes][u32 payload_len][payload bytes]
/// // 3. Use write_all() — never plain write()
/// todo!()
/// }
///
/// fn recv(&mut self) -> Result<(PacketHeader, Vec<u8>), TransportError> {
/// // 1. read_exact 4 bytes → header_len
/// // 2. Check header_len <= MAX_HEADER_BYTES before allocating
/// // 3. read_exact header_len bytes
/// // 4. Deserialise header
/// // 5. read_exact 4 bytes → payload_len
/// // 6. Check payload_len <= MAX_PAYLOAD_BYTES before allocating
/// // 7. read_exact payload_len bytes
/// // 8. Return (header, payload)
/// todo!()
/// }
/// }
///
/// // SAFETY: MyTransport owns its stream exclusively and does not share it.
/// unsafe impl Send for MyTransport {}
/// ```
pub trait Transport: Send {
/// Send one complete packet over this transport.
///
/// Blocks until all bytes have been written.
///
/// # Errors
///
/// Returns [`TransportError::Io`] if the write fails partway through,
/// or [`TransportError::Disconnected`] if the remote end is closed.
fn send(&mut self, header: &PacketHeader, payload: &[u8]) -> Result<(), TransportError>;
/// Receive one complete packet from this transport.
///
/// Blocks until a full header+payload pair is available.
///
/// # Errors
///
/// Returns [`TransportError::Disconnected`] if the remote closes cleanly,
/// [`TransportError::Io`] on I/O errors, [`TransportError::HeaderTooLarge`]
/// or [`TransportError::PayloadTooLarge`] if a size limit is exceeded,
/// and [`TransportError::DeserialiseError`] if the header cannot be decoded.
fn recv(&mut self) -> Result<(PacketHeader, Vec<u8>), TransportError>;
}
// ---------------------------------------------------------------------------
// Frame encoding helpers (shared by all transport implementations)
// ---------------------------------------------------------------------------
/// Encode a `PacketHeader` to bytes using rkyv.
///
/// Returns the serialised byte vector, or `None` if serialisation fails.
///
/// This is a low-level helper; transport implementations call it in `send()`.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{PacketHeader, PacketType};
/// use unshell::transport::encode_header;
///
/// let header = PacketHeader {
/// dst_path: "/router".into(),
/// src_path: "/agents/abc123".into(),
/// packet_type: PacketType::Handshake,
/// };
/// let bytes = encode_header(&header).expect("serialisation should not fail");
/// assert!(!bytes.is_empty());
/// ```
pub fn encode_header(header: &PacketHeader) -> Option<Vec<u8>> {
rkyv::to_bytes::<rkyv::rancor::Error>(header).ok().map(|b| b.to_vec())
}
/// Decode a `PacketHeader` from rkyv bytes.
///
/// Returns `Err(TransportError::DeserialiseError)` if the bytes are invalid.
///
/// This is a low-level helper; transport implementations call it in `recv()`.
///
/// # Example
///
/// ```rust
/// use unshell::protocol::{PacketHeader, PacketType};
/// use unshell::transport::{encode_header, decode_header};
///
/// let header = PacketHeader {
/// dst_path: "/router".into(),
/// src_path: "/agents/abc123".into(),
/// packet_type: PacketType::Handshake,
/// };
/// let bytes = encode_header(&header).unwrap();
/// let decoded = decode_header(&bytes).unwrap();
/// assert_eq!(decoded.dst_path, "/router");
/// ```
pub fn decode_header(bytes: &[u8]) -> Result<PacketHeader, TransportError> {
rkyv::from_bytes::<PacketHeader, rkyv::rancor::Error>(bytes)
.map_err(|_| TransportError::DeserialiseError)
}
src/transport/tcp.rs
+390
View File
@@ -0,0 +1,390 @@
//! # TCP Transport
//!
//! Only available when the `tcp` feature is enabled (requires `std`).
//! This file is only included in the module tree when `cfg(feature = "tcp")`,
//! as declared in `transport/mod.rs`.
//!
//! [`TcpTransport`] implements [`Transport`](super::Transport) over a
//! `std::net::TcpStream`.
//!
//! ## Framing
//!
//! Each `send` call writes:
//!
//! ```text
//! [u32 big-endian header_len] [header bytes]
//! [u32 big-endian payload_len] [payload bytes]
//! ```
//!
//! Each `recv` call:
//! 1. Reads exactly 4 bytes → `header_len`.
//! 2. Checks `header_len <= MAX_HEADER_BYTES`.
//! 3. Reads exactly `header_len` bytes.
//! 4. Deserialises the `PacketHeader`.
//! 5. Reads exactly 4 bytes → `payload_len`.
//! 6. Checks `payload_len <= MAX_PAYLOAD_BYTES`.
//! 7. Reads exactly `payload_len` bytes.
//! 8. Returns `(header, payload)`.
//!
//! **All reads use `read_exact`.** TCP is a stream protocol; a single `read`
//! may return fewer bytes than requested. `read_exact` loops until it has
//! the full count or the stream ends.
//!
//! ## Reconnection
//!
//! `TcpTransport` does not handle reconnection internally. The caller (the
//! payload's main loop or the operator CLI) is responsible for catching
//! [`TransportError::Disconnected`] and [`TransportError::Io`], then
//! creating a new `TcpTransport` to the router address.
extern crate alloc;
use alloc::vec;
use alloc::vec::Vec;
use std::io::{Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use super::{
decode_header, encode_header, TransportError, Transport, MAX_HEADER_BYTES, MAX_PAYLOAD_BYTES,
};
use crate::protocol::PacketHeader;
/// A framed TCP transport wrapping a `TcpStream`.
///
/// # Example: connecting as a payload
///
/// ```rust,no_run
/// use unshell::transport::tcp::TcpTransport;
///
/// // Connect to the router
/// let transport = TcpTransport::connect("127.0.0.1:9000").expect("connection failed");
/// ```
///
/// # Example: accepting a connection on the router
///
/// ```rust,no_run
/// use std::net::TcpListener;
/// use unshell::transport::tcp::TcpTransport;
///
/// let listener = TcpListener::bind("0.0.0.0:9000").unwrap();
/// for stream in listener.incoming() {
/// let transport = TcpTransport::from_stream(stream.unwrap());
/// // hand off to a node thread
/// }
/// ```
pub struct TcpTransport {
stream: TcpStream,
}
impl TcpTransport {
/// Connect to a remote address and return a transport wrapping that connection.
///
/// # Errors
///
/// Returns [`TransportError::Io`] if the connection fails.
///
/// # Example
///
/// ```rust,no_run
/// use unshell::transport::tcp::TcpTransport;
/// let t = TcpTransport::connect("127.0.0.1:9000").unwrap();
/// ```
pub fn connect<A: ToSocketAddrs>(addr: A) -> Result<Self, TransportError> {
let stream = TcpStream::connect(addr)?;
Ok(Self { stream })
}
/// Wrap an already-connected `TcpStream`.
///
/// Used by the router's accept loop, which creates streams via
/// `TcpListener::incoming()`.
///
/// # Example
///
/// ```rust,no_run
/// use std::net::TcpListener;
/// use unshell::transport::tcp::TcpTransport;
///
/// let listener = TcpListener::bind("0.0.0.0:9000").unwrap();
/// let (stream, _addr) = listener.accept().unwrap();
/// let transport = TcpTransport::from_stream(stream);
/// ```
pub fn from_stream(stream: TcpStream) -> Self {
Self { stream }
}
/// Access the underlying `TcpStream` for configuration (e.g., timeouts).
///
/// # Example
///
/// ```rust,no_run
/// use unshell::transport::tcp::TcpTransport;
/// use std::time::Duration;
///
/// let t = TcpTransport::connect("127.0.0.1:9000").unwrap();
/// t.stream_ref().set_read_timeout(Some(Duration::from_secs(5))).unwrap();
/// ```
pub fn stream_ref(&self) -> &TcpStream {
&self.stream
}
}
impl Transport for TcpTransport {
/// Send a packet (header + payload) over the TCP stream.
///
/// Writes the two-part frame atomically from the caller's perspective:
/// this call does not return until all bytes have been written or an
/// error occurs.
///
/// # Errors
///
/// - [`TransportError::Io`] on write failure or partial write.
/// - [`TransportError::Disconnected`] if the remote closed the connection.
fn send(&mut self, header: &PacketHeader, payload: &[u8]) -> Result<(), TransportError> {
// Serialise the header
let header_bytes =
encode_header(header).ok_or(TransportError::DeserialiseError)?;
// Build the full frame in one allocation so we can use a single
// write_all() call, reducing the chance of partial writes causing
// the remote to see a split frame.
//
// Frame layout:
// [u32 header_len][header bytes][u32 payload_len][payload bytes]
let header_len = header_bytes.len() as u32;
let payload_len = payload.len() as u32;
let mut frame =
Vec::with_capacity(8 + header_bytes.len() + payload.len());
frame.extend_from_slice(&header_len.to_be_bytes());
frame.extend_from_slice(&header_bytes);
frame.extend_from_slice(&payload_len.to_be_bytes());
frame.extend_from_slice(payload);
self.stream.write_all(&frame).map_err(|e| {
if e.kind() == std::io::ErrorKind::BrokenPipe
|| e.kind() == std::io::ErrorKind::ConnectionReset
|| e.kind() == std::io::ErrorKind::UnexpectedEof
{
TransportError::Disconnected
} else {
TransportError::Io(e)
}
})
}
/// Receive one complete packet from the TCP stream.
///
/// Blocks until a full header+payload pair is available.
///
/// # Errors
///
/// - [`TransportError::Disconnected`] if the remote closed cleanly (EOF).
/// - [`TransportError::Io`] on I/O errors.
/// - [`TransportError::HeaderTooLarge`] if the announced header size
/// exceeds [`MAX_HEADER_BYTES`].
/// - [`TransportError::PayloadTooLarge`] if the announced payload size
/// exceeds [`MAX_PAYLOAD_BYTES`].
/// - [`TransportError::DeserialiseError`] if the header bytes are invalid.
fn recv(&mut self) -> Result<(PacketHeader, Vec<u8>), TransportError> {
// --- Step 1: Read header length (4 bytes) ---
let header_len = read_u32(&mut self.stream)?;
if header_len > MAX_HEADER_BYTES {
return Err(TransportError::HeaderTooLarge(header_len, MAX_HEADER_BYTES));
}
// --- Step 2: Read header bytes ---
let mut header_buf = vec![0u8; header_len];
read_exact(&mut self.stream, &mut header_buf)?;
// --- Step 3: Deserialise header ---
let header = decode_header(&header_buf)?;
// --- Step 4: Read payload length (4 bytes) ---
let payload_len = read_u32(&mut self.stream)?;
if payload_len > MAX_PAYLOAD_BYTES {
return Err(TransportError::PayloadTooLarge(payload_len, MAX_PAYLOAD_BYTES));
}
// --- Step 5: Read payload bytes ---
let mut payload = vec![0u8; payload_len];
read_exact(&mut self.stream, &mut payload)?;
Ok((header, payload))
}
}
// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------
/// Read exactly 4 bytes from `stream` and interpret them as a big-endian `u32`.
///
/// Returns [`TransportError::Disconnected`] on clean EOF (zero bytes read),
/// or [`TransportError::Io`] on other errors.
fn read_u32(stream: &mut TcpStream) -> Result<usize, TransportError> {
let mut buf = [0u8; 4];
read_exact(stream, &mut buf)?;
Ok(u32::from_be_bytes(buf) as usize)
}
/// Read exactly `buf.len()` bytes from `stream`.
///
/// Unlike `stream.read()`, this function loops until the buffer is full or
/// an error occurs. This is essential for TCP, which may deliver data in
/// smaller chunks than requested.
///
/// Returns [`TransportError::Disconnected`] on clean EOF,
/// [`TransportError::Io`] on I/O errors.
fn read_exact(stream: &mut TcpStream, buf: &mut [u8]) -> Result<(), TransportError> {
stream.read_exact(buf).map_err(|e| {
if e.kind() == std::io::ErrorKind::UnexpectedEof
|| e.kind() == std::io::ErrorKind::ConnectionReset
{
TransportError::Disconnected
} else {
TransportError::Io(e)
}
})
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use crate::protocol::PacketType;
use std::net::TcpListener;
use std::thread;
/// Test that a packet sent through a real TcpStream arrives intact.
///
/// This test spins up a local listener on an ephemeral port, sends one
/// packet from one thread, and verifies the other receives it correctly.
#[test]
fn roundtrip_over_real_tcp() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind failed");
let addr = listener.local_addr().expect("local_addr failed");
let header_sent = PacketHeader {
dst_path: "/agents/test/shell".into(),
src_path: "/operator/sess1".into(),
packet_type: PacketType::Request,
};
let payload_sent = b"hello world".to_vec();
let header_clone = header_sent.clone();
let payload_clone = payload_sent.clone();
// Sender thread
let sender = thread::spawn(move || {
let stream = TcpStream::connect(addr).expect("connect failed");
let mut transport = TcpTransport::from_stream(stream);
transport
.send(&header_clone, &payload_clone)
.expect("send failed");
});
// Receiver (main thread)
let (stream, _) = listener.accept().expect("accept failed");
let mut transport = TcpTransport::from_stream(stream);
let (header_recv, payload_recv) = transport.recv().expect("recv failed");
sender.join().expect("sender thread panicked");
assert_eq!(header_recv.dst_path, header_sent.dst_path);
assert_eq!(header_recv.src_path, header_sent.src_path);
assert_eq!(header_recv.packet_type, header_sent.packet_type);
assert_eq!(payload_recv, payload_sent);
}
/// Test that an empty payload round-trips correctly.
#[test]
fn roundtrip_empty_payload() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind failed");
let addr = listener.local_addr().expect("local_addr failed");
let header = PacketHeader {
dst_path: "/router/ping".into(),
src_path: "/operator/sess1".into(),
packet_type: PacketType::Request,
};
let header_clone = header.clone();
let sender = thread::spawn(move || {
let stream = TcpStream::connect(addr).expect("connect failed");
let mut t = TcpTransport::from_stream(stream);
t.send(&header_clone, &[]).expect("send failed");
});
let (stream, _) = listener.accept().expect("accept failed");
let mut t = TcpTransport::from_stream(stream);
let (recv_header, recv_payload) = t.recv().expect("recv failed");
sender.join().expect("sender thread panicked");
assert_eq!(recv_header.dst_path, "/router/ping");
assert!(recv_payload.is_empty());
}
/// Test that a large payload (1 MB) survives the TCP framing.
#[test]
fn roundtrip_large_payload() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind failed");
let addr = listener.local_addr().expect("local_addr failed");
let payload: Vec<u8> = (0..1_000_000u32).map(|i| (i % 256) as u8).collect();
let payload_clone = payload.clone();
let header = PacketHeader {
dst_path: "/agents/x/files/read".into(),
src_path: "/operator/sess1".into(),
packet_type: PacketType::Response,
};
let header_clone = header.clone();
let sender = thread::spawn(move || {
let stream = TcpStream::connect(addr).expect("connect failed");
let mut t = TcpTransport::from_stream(stream);
t.send(&header_clone, &payload_clone).expect("send failed");
});
let (stream, _) = listener.accept().expect("accept failed");
let mut t = TcpTransport::from_stream(stream);
let (_, recv_payload) = t.recv().expect("recv failed");
sender.join().expect("sender thread panicked");
assert_eq!(recv_payload, payload);
}
/// Test that a frame whose announced header size exceeds the limit is rejected
/// without allocating the full buffer.
#[test]
fn rejects_oversized_header() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind failed");
let addr = listener.local_addr().expect("local_addr failed");
let sender = thread::spawn(move || {
let mut stream = TcpStream::connect(addr).expect("connect failed");
// Write an enormous header length
let huge_len = (MAX_HEADER_BYTES + 1) as u32;
stream
.write_all(&huge_len.to_be_bytes())
.expect("write failed");
});
let (stream, _) = listener.accept().expect("accept failed");
let mut t = TcpTransport::from_stream(stream);
let result = t.recv();
sender.join().expect("sender panicked");
assert!(
matches!(result, Err(TransportError::HeaderTooLarge(_, _))),
"expected HeaderTooLarge, got: {result:?}"
);
}
}
src/tree/mod.rs
+504 -40
View File
@@ -1,56 +1,520 @@
use alloc::{boxed::Box, string::String, vec::Vec};
//! # Tree Module
//!
//! The `Tree` dispatches incoming [`TreeRequest`]s to registered [`Endpoint`]s
//! by matching the request's destination path.
//!
//! ## Path matching
//!
//! Paths are `/`-delimited strings. An `Endpoint` is registered at a path prefix.
//! A request matches an endpoint if the endpoint's path is a prefix of the request path.
//! When multiple endpoints match, the one with the **longest** prefix wins.
//!
//! ```text
//! Registered endpoints: Request path:
//! /shell ← prefix /shell/exec → matches /shell
//! /files ← prefix /files/read → matches /files
//! /shell/exec ← more specific /shell/exec → matches /shell/exec (longer)
//! ```
//!
//! ## Usage
//!
//! ```rust
//! use unshell::tree::{Tree, Endpoint};
//! use unshell::protocol::{
//! TreeRequest, TreeResponse, RequestType, ResponseStatus, content,
//! };
//!
//! /// A simple echo endpoint that reflects the request data back.
//! struct EchoEndpoint;
//!
//! impl Endpoint for EchoEndpoint {
//! fn handle(&mut self, request: TreeRequest) -> TreeResponse {
//! TreeResponse {
//! request_id: request.request_id,
//! status: ResponseStatus::Ok,
//! content_type: request.content_type.clone(),
//! data: request.data.clone(),
//! }
//! }
//! }
//!
//! let mut tree = Tree::new();
//! tree.register("/echo", EchoEndpoint);
//!
//! let response = tree.dispatch(TreeRequest {
//! request_id: 1,
//! request_type: RequestType::Read,
//! content_type: content::UTF8_STRING.into(),
//! data: b"hello".to_vec(),
//! }, "/echo/anything");
//!
//! assert_eq!(response.status, ResponseStatus::Ok);
//! assert_eq!(response.data, b"hello");
//! ```
mod request;
extern crate alloc;
use alloc::borrow::ToOwned;
use alloc::boxed::Box;
use alloc::string::String;
use alloc::vec::Vec;
pub use request::{TreeRequest, TreeRequestType};
use crate::protocol::{
content, ResponseStatus, TreeRequest, TreeResponse,
};
pub mod types;
// ---------------------------------------------------------------------------
// Endpoint trait
// ---------------------------------------------------------------------------
#[derive(Default)]
/// A module that handles [`TreeRequest`]s at a registered path.
///
/// Implement this trait to add capabilities to a payload. The `Tree` calls
/// `handle` when a request's path matches this endpoint's registration prefix.
///
/// # Example
///
/// ```rust
/// use unshell::tree::Endpoint;
/// use unshell::protocol::{TreeRequest, TreeResponse, ResponseStatus, content};
///
/// struct PingEndpoint;
///
/// impl Endpoint for PingEndpoint {
/// fn handle(&mut self, request: TreeRequest) -> TreeResponse {
/// TreeResponse {
/// request_id: request.request_id,
/// status: ResponseStatus::Ok,
/// content_type: content::UTF8_STRING.into(),
/// data: b"pong".to_vec(),
/// }
/// }
/// }
/// ```
pub trait Endpoint: Send {
/// Handle a request and return a response.
///
/// This method is called synchronously on the recv loop thread. It should
/// not block for extended periods. For long-running operations, spawn a
/// background thread and return immediately with a `pending` response
/// (streaming responses are a future protocol feature).
fn handle(&mut self, request: TreeRequest) -> TreeResponse;
}
// ---------------------------------------------------------------------------
// Tree
// ---------------------------------------------------------------------------
/// A path-addressed dispatcher that routes [`TreeRequest`]s to [`Endpoint`]s.
///
/// # Path matching algorithm
///
/// The tree uses **longest-prefix matching**:
/// 1. Split the request path by `/`.
/// 2. For each registered endpoint, check if the endpoint's path components
/// are a prefix of the request path components.
/// 3. Among all matching endpoints, return the one with the most components
/// (the most specific match).
/// 4. If no match: return a [`ResponseStatus::NoBranchError`] response.
///
/// # Example
///
/// ```rust
/// use unshell::tree::{Tree, Endpoint};
/// use unshell::protocol::{TreeRequest, TreeResponse, RequestType, ResponseStatus, content};
///
/// struct Shell;
///
/// impl Endpoint for Shell {
/// fn handle(&mut self, req: TreeRequest) -> TreeResponse {
/// TreeResponse {
/// request_id: req.request_id,
/// status: ResponseStatus::Ok,
/// content_type: content::UTF8_STRING.into(),
/// data: b"shell output".to_vec(),
/// }
/// }
/// }
///
/// let mut tree = Tree::new();
/// tree.register("/shell", Shell);
///
/// // A request to /shell/exec/anything matches /shell (the registered prefix).
/// let resp = tree.dispatch(
/// TreeRequest {
/// request_id: 1,
/// request_type: RequestType::CallProcedure,
/// content_type: content::NONE.into(),
/// data: Vec::new(),
/// },
/// "/shell/exec",
/// );
/// assert_eq!(resp.status, ResponseStatus::Ok);
/// ```
pub struct Tree {
endpoints: Vec<(Box<dyn Endpoint>, Vec<String>)>,
/// Registered endpoints with their path prefixes.
///
/// The path is stored as a `Vec<String>` of components (split on `/`,
/// empty leading component from the leading `/` is discarded).
endpoints: Vec<(Vec<String>, Box<dyn Endpoint>)>,
}
impl Tree {
pub fn add_endpoint<T: Endpoint + 'static>(&mut self, endpoint: T, path: Vec<String>) {
self.add_endpoint_box(Box::new(endpoint), path);
}
pub fn add_endpoint_box(&mut self, endpoint: Box<dyn Endpoint>, path: Vec<String>) {
self.endpoints.push((endpoint, path));
/// Create an empty tree with no registered endpoints.
#[must_use]
pub fn new() -> Self {
Self {
endpoints: Vec::new(),
}
}
pub fn get_endpoint(&mut self, search_path: &Vec<String>) -> Option<&mut Box<dyn Endpoint>> {
for (endpoint, endpoint_path) in &mut self.endpoints {
if search_path.len() < endpoint_path.len() {
return None;
}
/// Register an endpoint at the given path prefix.
///
/// # Arguments
///
/// * `path` — the path prefix this endpoint owns, e.g. `"/shell"`.
/// Leading `/` is stripped; components are split on `/`.
/// * `endpoint` — the handler that will receive matching requests.
///
/// # Panics
///
/// Does not panic. Registering the same path twice is allowed; the second
/// registration shadows the first for that exact path (longest-prefix
/// matching still applies for sub-paths).
///
/// # Example
///
/// ```rust
/// use unshell::tree::{Tree, Endpoint};
/// use unshell::protocol::{TreeRequest, TreeResponse, ResponseStatus, content};
///
/// struct Noop;
/// impl Endpoint for Noop {
/// fn handle(&mut self, req: TreeRequest) -> TreeResponse {
/// TreeResponse {
/// request_id: req.request_id,
/// status: ResponseStatus::Ok,
/// content_type: content::NONE.into(),
/// data: Vec::new(),
/// }
/// }
/// }
///
/// let mut tree = Tree::new();
/// tree.register("/shell", Noop);
/// ```
pub fn register<E: Endpoint + 'static>(&mut self, path: &str, endpoint: E) {
let components = split_path(path);
self.endpoints.push((components, Box::new(endpoint)));
}
for i in 0..endpoint_path.len() {
if search_path[i] != endpoint_path[i] {
return None;
/// Dispatch a request to the best-matching endpoint.
///
/// Returns a [`TreeResponse`] with [`ResponseStatus::NoBranchError`]
/// if no registered endpoint matches the request path.
///
/// # Arguments
///
/// * `request` — the incoming request.
/// * `dst_path` — the destination path from the packet header.
///
/// # Example
///
/// ```rust
/// use unshell::tree::Tree;
/// use unshell::protocol::{TreeRequest, RequestType, ResponseStatus, content};
///
/// let mut tree = Tree::new();
/// // (register some endpoints here)
///
/// let resp = tree.dispatch(
/// TreeRequest {
/// request_id: 99,
/// request_type: RequestType::Read,
/// content_type: content::NONE.into(),
/// data: Vec::new(),
/// },
/// "/unknown/path",
/// );
/// assert_eq!(resp.status, ResponseStatus::NoBranchError);
/// ```
pub fn dispatch(&mut self, request: TreeRequest, dst_path: &str) -> TreeResponse {
let path_components = split_path(dst_path);
// Find the endpoint with the longest matching prefix.
let best = self
.endpoints
.iter_mut()
.filter(|(ep_path, _)| is_prefix(ep_path, &path_components))
.max_by_key(|(ep_path, _)| ep_path.len());
match best {
Some((_, endpoint)) => endpoint.handle(request),
None => TreeResponse {
request_id: request.request_id,
status: ResponseStatus::NoBranchError,
content_type: content::NONE.into(),
data: Vec::new(),
},
}
}
/// Return the list of registered path prefixes.
///
/// Used during handshake to tell the router which paths this tree owns.
///
/// # Example
///
/// ```rust
/// use unshell::tree::{Tree, Endpoint};
/// use unshell::protocol::{TreeRequest, TreeResponse, ResponseStatus, content};
///
/// struct Noop;
/// impl Endpoint for Noop {
/// fn handle(&mut self, req: TreeRequest) -> TreeResponse {
/// TreeResponse {
/// request_id: req.request_id,
/// status: ResponseStatus::Ok,
/// content_type: content::NONE.into(),
/// data: Vec::new(),
/// }
/// }
/// }
///
/// let mut tree = Tree::new();
/// tree.register("/shell", Noop);
/// tree.register("/files", Noop);
///
/// let paths = tree.registered_paths("/agents/abc123");
/// assert!(paths.contains(&"/agents/abc123/shell".to_string()));
/// assert!(paths.contains(&"/agents/abc123/files".to_string()));
/// ```
#[must_use]
pub fn registered_paths(&self, base_prefix: &str) -> Vec<String> {
let base = base_prefix.trim_end_matches('/');
self.endpoints
.iter()
.map(|(components, _)| {
let sub = components.join("/");
if sub.is_empty() {
base.to_owned()
} else {
alloc::format!("{base}/{sub}")
}
}
return Some(endpoint);
}
return None;
}
pub fn request(&mut self, request: TreeRequest) -> TreeRequest {
if let Some(endpoint) = self.get_endpoint(&request.path) {
endpoint.request(request)
} else {
TreeRequest {
path: request.path,
request_type: TreeRequestType::NoBranchError,
content_type: types::TYPE_NONE.into(),
data: Vec::with_capacity(0),
}
}
})
.collect()
}
}
pub trait Endpoint {
fn request(&mut self, request: TreeRequest) -> TreeRequest;
impl Default for Tree {
fn default() -> Self {
Self::new()
}
}
// ---------------------------------------------------------------------------
// Path utilities
// ---------------------------------------------------------------------------
/// Split a path string into its components.
///
/// Leading `/` and empty segments are discarded.
///
/// ```text
/// "/shell/exec" → ["shell", "exec"]
/// "/shell/" → ["shell"]
/// "shell" → ["shell"]
/// "/" → []
/// ```
fn split_path(path: &str) -> Vec<String> {
path.split('/')
.filter(|s| !s.is_empty())
.map(String::from)
.collect()
}
/// Returns `true` if `prefix` is a prefix of (or equal to) `path`.
///
/// Both are slices of path components (already split on `/`).
///
/// ```text
/// prefix = ["shell"] path = ["shell", "exec"] → true
/// prefix = ["shell", "exec"] path = ["shell", "exec"] → true (exact match)
/// prefix = ["shell", "exec"] path = ["shell"] → false (prefix longer)
/// prefix = ["files"] path = ["shell", "exec"] → false (different root)
/// ```
fn is_prefix(prefix: &[String], path: &[String]) -> bool {
if prefix.len() > path.len() {
return false;
}
prefix.iter().zip(path.iter()).all(|(a, b)| a == b)
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use crate::protocol::{RequestType, ResponseStatus, content};
// A minimal endpoint that echoes the request data.
struct Echo;
impl Endpoint for Echo {
fn handle(&mut self, req: TreeRequest) -> TreeResponse {
TreeResponse {
request_id: req.request_id,
status: ResponseStatus::Ok,
content_type: req.content_type,
data: req.data,
}
}
}
// A minimal endpoint that always returns a fixed string.
struct Fixed(&'static str);
impl Endpoint for Fixed {
fn handle(&mut self, req: TreeRequest) -> TreeResponse {
TreeResponse {
request_id: req.request_id,
status: ResponseStatus::Ok,
content_type: content::UTF8_STRING.into(),
data: self.0.as_bytes().to_vec(),
}
}
}
fn make_req(id: u64) -> TreeRequest {
TreeRequest {
request_id: id,
request_type: RequestType::Read,
content_type: content::NONE.into(),
data: Vec::new(),
}
}
/// A single endpoint is matched correctly.
#[test]
fn single_endpoint_match() {
let mut tree = Tree::new();
tree.register("/shell", Echo);
let resp = tree.dispatch(make_req(1), "/shell/exec");
assert_eq!(resp.status, ResponseStatus::Ok, "expected Ok for /shell/exec");
assert_eq!(resp.request_id, 1);
}
/// When two endpoints are registered, the second one is also reachable.
///
/// This test specifically catches the old `return None` bug in `get_endpoint`:
/// the first endpoint (/files) doesn't match /shell/exec, so the tree must
/// continue to the second entry (/shell).
#[test]
fn second_endpoint_match() {
let mut tree = Tree::new();
tree.register("/files", Fixed("files"));
tree.register("/shell", Fixed("shell"));
let resp = tree.dispatch(make_req(2), "/shell/exec");
assert_eq!(resp.status, ResponseStatus::Ok);
assert_eq!(resp.data, b"shell");
}
/// No matching endpoint returns NoBranchError.
#[test]
fn no_match_returns_no_branch_error() {
let mut tree = Tree::new();
tree.register("/shell", Echo);
let resp = tree.dispatch(make_req(3), "/nonexistent/path");
assert_eq!(resp.status, ResponseStatus::NoBranchError);
assert_eq!(resp.request_id, 3);
}
/// Longer (more specific) prefix wins over shorter prefix.
#[test]
fn longer_prefix_wins() {
let mut tree = Tree::new();
tree.register("/shell", Fixed("short"));
tree.register("/shell/exec", Fixed("long"));
let resp = tree.dispatch(make_req(4), "/shell/exec/anything");
assert_eq!(resp.data, b"long", "longer prefix should win");
}
/// A request path that is shorter than the registered prefix does not match.
#[test]
fn prefix_does_not_overmatch() {
let mut tree = Tree::new();
tree.register("/shell/exec/something", Echo);
// /shell/exec is shorter than the registered path — should NOT match
let resp = tree.dispatch(make_req(5), "/shell/exec");
assert_eq!(resp.status, ResponseStatus::NoBranchError);
}
/// `registered_paths` returns all prefixes with the base path prepended.
#[test]
fn registered_paths_prepends_base() {
let mut tree = Tree::new();
tree.register("/shell", Echo);
tree.register("/files", Echo);
let paths = tree.registered_paths("/agents/abc123");
assert!(paths.contains(&"/agents/abc123/shell".to_string()));
assert!(paths.contains(&"/agents/abc123/files".to_string()));
assert_eq!(paths.len(), 2);
}
// -----------------------------------------------------------------------
// Path utility tests
// -----------------------------------------------------------------------
#[test]
fn split_path_leading_slash() {
assert_eq!(split_path("/shell/exec"), vec!["shell", "exec"]);
}
#[test]
fn split_path_no_leading_slash() {
assert_eq!(split_path("shell/exec"), vec!["shell", "exec"]);
}
#[test]
fn split_path_trailing_slash() {
assert_eq!(split_path("/shell/"), vec!["shell"]);
}
#[test]
fn split_path_root() {
let result: Vec<String> = split_path("/");
assert!(result.is_empty());
}
#[test]
fn is_prefix_exact_match() {
let p = split_path("/shell/exec");
assert!(is_prefix(&p, &p));
}
#[test]
fn is_prefix_valid() {
let prefix = split_path("/shell");
let path = split_path("/shell/exec");
assert!(is_prefix(&prefix, &path));
}
#[test]
fn is_prefix_prefix_too_long() {
let prefix = split_path("/shell/exec");
let path = split_path("/shell");
assert!(!is_prefix(&prefix, &path));
}
#[test]
fn is_prefix_different_root() {
let prefix = split_path("/files");
let path = split_path("/shell/exec");
assert!(!is_prefix(&prefix, &path));
}
}
src/tree/request.rs
-39
View File
@@ -1,39 +0,0 @@
// use std::collections::VecDeque;
use alloc::{string::String, vec::Vec};
use rkyv::{Archive, Deserialize, Serialize};
#[derive(Archive, Deserialize, Serialize)]
#[rkyv(compare(PartialEq), derive(Debug))]
pub struct TreeRequest {
// The exact path that this packet should be heading down to
pub path: Vec<String>,
// // The list of previous paths that this packet came from
// // This is the destination path added in reverse order
// pub source_path: VecDeque<String>,
pub request_type: TreeRequestType,
// The data type of the payload, to determine how to deserialize and interpret it on the other side
// This is equivalent to HTTP's content-type header
pub content_type: String,
// The payload of the packet
pub data: Vec<u8>,
}
#[derive(Archive, Deserialize, Serialize)]
#[rkyv(compare(PartialEq), derive(Debug))]
pub enum TreeRequestType {
Return = 0,
Read = 1,
GetProcedures = 2,
Write = 11,
CallProcedure = 12,
UnnamedError = 100,
NoBranchError = 101,
ProtocolError = 102,
ExecutionError = 103,
}
src/tree/types.rs
-16
View File
@@ -1,16 +0,0 @@
use alloc::{string::String, vec::Vec};
use crate::obfuscate::sym;
pub const TYPE_NONE: &'static str = sym!("core/None");
pub const TYPE_PROCEDURE_CALL_DESCRIPTOR: &'static str = sym!("core/Procedure_call_descriptor");
pub struct ProcedureCallDescriptor {
name: String,
}
pub const TYPE_PROCEDURE_CALL_DESCRIPTOR_LIST: &'static str =
sym!("core/Procedure_call_descriptor_list");
pub type ProcedureCallDescriptorList = Vec<ProcedureCallDescriptor>;
pub const TYPE_PROCEDURE_CALL_ARGUMENTS: &'static str = sym!("core/Procedure_call_arguments");