mirror of
https://github.com/Astatin3/unshell.git
synced 2026-06-09 06:47:59 -06:00
feat: complete protocol spec and initial implementation
- 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
+504
-40
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user