mirror of
https://github.com/Astatin3/unshell.git
synced 2026-06-08 22:38:01 -06:00
Move manager into it's own project
This commit is contained in:
Michael Mikovsky
Generated
-11
@@ -317,16 +317,6 @@ version = "0.2.177"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976"
|
||||
|
||||
[[package]]
|
||||
name = "libloading"
|
||||
version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lock_api"
|
||||
version = "0.4.14"
|
||||
@@ -654,7 +644,6 @@ dependencies = [
|
||||
"chrono",
|
||||
"crossbeam-channel",
|
||||
"libc",
|
||||
"libloading",
|
||||
"rand",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
||||
@@ -16,7 +16,7 @@ obfuscate = ["unshell-obfuscate/obfuscate"]
|
||||
|
||||
[dependencies]
|
||||
# Base dependencies
|
||||
libloading = {version = "0.8.9"}
|
||||
# libloading = {version = "0.8.9"}
|
||||
bincode = "2.0.1"
|
||||
unshell-obfuscate = {path = "../unshell-obfuscate"}
|
||||
chrono = "0.4.42"
|
||||
|
||||
@@ -1,25 +1,4 @@
|
||||
use std::{collections::HashMap, fmt::Debug};
|
||||
|
||||
// use bincode::{Decode, Encode};
|
||||
// use serde::{Deserialize, Serialize};
|
||||
|
||||
// use bincode::{Decode, Encode};
|
||||
|
||||
use crate::{ModuleError, ModuleRuntime};
|
||||
|
||||
// /// Payload config that is instantiated
|
||||
// #[derive(Serialize, Deserialize)]
|
||||
// pub struct Config {
|
||||
// pub id: String,
|
||||
// pub key: String,
|
||||
// pub components: Vec<String>,
|
||||
// }
|
||||
|
||||
pub struct PayloadConfig {
|
||||
pub id: &'static str,
|
||||
pub components: Vec<NamedComponent>,
|
||||
pub runtime_config: Vec<RuntimeConfig>,
|
||||
}
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RuntimeConfig {
|
||||
@@ -27,41 +6,3 @@ pub struct RuntimeConfig {
|
||||
pub name: String,
|
||||
pub config: HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct NamedComponent {
|
||||
pub name: &'static str,
|
||||
|
||||
// + Sync + Sync + Sync + Sync + Sync + Sync + Sync + Sync
|
||||
pub get_interface: &'static (dyn Fn() -> Option<&'static (dyn InterfaceWrapper + Sync)> + Sync),
|
||||
pub start_runtime: &'static (
|
||||
dyn Fn(&'static RuntimeConfig) -> Result<Box<dyn ModuleRuntime>, ModuleError>
|
||||
+ Sync
|
||||
),
|
||||
}
|
||||
|
||||
impl Debug for NamedComponent {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("NamedComponent")
|
||||
.field("name", &self.name)
|
||||
// .field("get_interface", &self.get_interface)
|
||||
// .field("start_runtime", &self.start_runtime)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait that wraps the get_interface<T>() function inside of components
|
||||
pub trait InterfaceWrapper: Send + Sync {
|
||||
fn get_interface<T: 'static>(&self) -> Option<T>
|
||||
where
|
||||
Self: Sized;
|
||||
}
|
||||
|
||||
// impl<T: 'static> InterfaceWrapper for T {
|
||||
// default fn get_interface<T>() -> Option<T>
|
||||
// where
|
||||
// Self: Sized,
|
||||
// {
|
||||
// None
|
||||
// }
|
||||
// }
|
||||
|
||||
+3
-19
@@ -2,23 +2,18 @@
|
||||
|
||||
pub mod config;
|
||||
pub mod logger;
|
||||
pub mod module;
|
||||
pub mod network;
|
||||
|
||||
mod announcement;
|
||||
use std::{
|
||||
fmt::{self, Debug},
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
use std::fmt::{self, Debug};
|
||||
|
||||
pub use announcement::Announcement;
|
||||
|
||||
use crate::module::Manager;
|
||||
pub type Result<T> = std::result::Result<T, ModuleError>;
|
||||
|
||||
///Generic error type for module-related operations.
|
||||
#[derive(Debug)]
|
||||
pub enum ModuleError {
|
||||
LibLoadingError(libloading::Error),
|
||||
LibLoadingError(String),
|
||||
// LogError(log::SetLoggerError),
|
||||
LinkError(String),
|
||||
CryptError(String),
|
||||
@@ -45,17 +40,6 @@ impl fmt::Display for ModuleError {
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for defining modules that have a runtime.
|
||||
pub trait ModuleRuntime: Send + Sync {
|
||||
fn init(&mut self, manager: Arc<Mutex<Manager>>) -> Result<(), ModuleError>;
|
||||
|
||||
/// Returns true if the module is running.
|
||||
/// After returning false, the module will be dropped.
|
||||
fn is_running(&self) -> bool;
|
||||
/// Consumes the module, implementation should kill whatever is running.
|
||||
fn kill(self: Box<Self>);
|
||||
}
|
||||
|
||||
// pub trait Component {
|
||||
// fn name(&self) -> &'static str;
|
||||
// // fn start_runtime(&self, manager: Arc<Mutex<Manager>>) -> Option<Box<dyn ModuleRuntime>>;
|
||||
|
||||
@@ -1,196 +0,0 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex},
|
||||
thread::{self, JoinHandle},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
config::{NamedComponent, PayloadConfig, RuntimeConfig},
|
||||
network::Stream,
|
||||
*,
|
||||
};
|
||||
use module::Module;
|
||||
use unshell_obfuscate::symbol;
|
||||
|
||||
// #[derive(Debug)]
|
||||
pub struct Manager {
|
||||
id: &'static str,
|
||||
|
||||
handle: Option<JoinHandle<()>>,
|
||||
|
||||
pub modules: Vec<Module>,
|
||||
|
||||
components: HashMap<String, NamedComponent>,
|
||||
active_runtimes: Vec<Box<dyn ModuleRuntime>>,
|
||||
|
||||
pub connections: Vec<Box<dyn Stream<Announcement>>>,
|
||||
}
|
||||
|
||||
// static mut MANAGER_RUNTIME: Option<Arc<Mutex<Manager>>> = None;
|
||||
|
||||
impl Manager {
|
||||
fn new(id: &'static str, components: Vec<NamedComponent>, modules: Vec<Module>) -> Self {
|
||||
Self {
|
||||
id,
|
||||
handle: None,
|
||||
|
||||
modules,
|
||||
components: components
|
||||
.into_iter()
|
||||
.map(|c| (c.name.to_string(), c))
|
||||
.collect(),
|
||||
active_runtimes: Vec::new(),
|
||||
|
||||
connections: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create Manager, and run initialization for each Module
|
||||
#[allow(static_mut_refs)]
|
||||
pub fn start(config: &'static PayloadConfig, modules: Vec<Module>) -> Arc<Mutex<Self>> {
|
||||
// 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();
|
||||
|
||||
let this = Arc::new(Mutex::new(this));
|
||||
|
||||
debug!("Creating runtimes...");
|
||||
for runtime in &config.runtime_config {
|
||||
Self::create_runtime(this.clone(), runtime);
|
||||
}
|
||||
|
||||
debug!("Starting runtimes...");
|
||||
for runtime in &mut this.lock().unwrap().active_runtimes {
|
||||
if let Err(e) = runtime.init(this.clone()) {
|
||||
warn!("Failed to start runtime: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
this.lock().unwrap().handle = Some(Self::start_thread(this.clone()));
|
||||
|
||||
this
|
||||
}
|
||||
|
||||
fn load_components(&mut self) {
|
||||
for module in &self.modules {
|
||||
// Load get_components function from shared object library
|
||||
let component_func = match module
|
||||
.get_symbol::<fn() -> Vec<NamedComponent>>(symbol!("get_components").as_bytes())
|
||||
{
|
||||
Ok(func) => func,
|
||||
Err(_) => {
|
||||
warn!("get_components function not found");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let components = component_func();
|
||||
let component_name = "TODO"; //TODO: Make this actually load component name
|
||||
|
||||
debug!("{} - Retrieved payload metadata", component_name);
|
||||
|
||||
// Add each component into self
|
||||
for c in components {
|
||||
debug!("{} - Found component '{}'", "TODO", c.name);
|
||||
self.components.insert(c.name.to_owned(), c);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The manager thread. receives announcements, and kills runtimes.
|
||||
fn start_thread(this: Arc<Mutex<Self>>) -> JoinHandle<()> {
|
||||
thread::spawn(move || {
|
||||
loop {
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
|
||||
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| {
|
||||
if runtime.is_running() {
|
||||
true
|
||||
} else {
|
||||
debug!("Runtime exited!"); //TODO: Make this better
|
||||
false
|
||||
}
|
||||
});
|
||||
|
||||
// Read announcements
|
||||
this_lock.recv_connection_announcements();
|
||||
|
||||
// Prune dead connections
|
||||
this_lock.prune_connections();
|
||||
|
||||
drop(this_lock)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Wait for manager thread to finish.
|
||||
pub fn join(this: Arc<Mutex<Self>>) {
|
||||
loop {
|
||||
if this.lock().unwrap().handle.as_ref().unwrap().is_finished() {
|
||||
break;
|
||||
}
|
||||
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
}
|
||||
|
||||
/// Start a runtime
|
||||
fn create_runtime<'a>(this: Arc<Mutex<Self>>, runtime: &'static RuntimeConfig) {
|
||||
let mut this_lock = this.lock().unwrap();
|
||||
|
||||
let component = match this_lock.components.get(&runtime.parent_component) {
|
||||
Some(component) => component,
|
||||
None => {
|
||||
warn!(
|
||||
"Could not find component '{}' which is referenced by runtime: {}",
|
||||
runtime.parent_component, runtime.name
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
debug!("Starting runtime: {}", runtime.name);
|
||||
|
||||
let runtime = match (*component.start_runtime)(runtime) {
|
||||
Ok(runtime) => runtime,
|
||||
Err(e) => {
|
||||
warn!("Failed to start runtime: {:?}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
this_lock.active_runtimes.push(runtime);
|
||||
}
|
||||
|
||||
pub fn add_runtime(
|
||||
this: Arc<Mutex<Self>>,
|
||||
runtime: &'static RuntimeConfig,
|
||||
) -> Result<(), ModuleError> {
|
||||
Self::create_runtime(this.clone(), runtime);
|
||||
|
||||
this.lock()
|
||||
.unwrap()
|
||||
.active_runtimes
|
||||
.iter_mut()
|
||||
.last()
|
||||
.unwrap()
|
||||
.init(this.clone())
|
||||
}
|
||||
|
||||
pub fn get_name(&self) -> &str {
|
||||
self.id
|
||||
}
|
||||
}
|
||||
@@ -1,15 +0,0 @@
|
||||
use crate::{Announcement, module::Manager};
|
||||
|
||||
impl Manager {
|
||||
pub fn recv_announcement(&mut self, announcement: &Announcement) {
|
||||
match announcement {
|
||||
Announcement::TestAnnouncement(str) => {
|
||||
println!("Got test announcement: {}", str)
|
||||
} // Announcement::GetRuntimes => todo!(),
|
||||
// Announcement::GetRuntimesAck(_) => todo!(),
|
||||
// Announcement::StartRuntime(runtime_config) => todo!(),
|
||||
// Announcement::StartRuntimeAck(_) => todo!(),
|
||||
// _ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,32 +0,0 @@
|
||||
use crate::{Announcement, ModuleError, module::Manager, network::Stream};
|
||||
|
||||
impl Manager {
|
||||
pub fn add_connection(&mut self, connection: Box<dyn Stream<Announcement>>) {
|
||||
self.connections.push(connection);
|
||||
}
|
||||
|
||||
pub fn prune_connections(&mut self) {
|
||||
self.connections.retain(|c| c.is_alive());
|
||||
}
|
||||
|
||||
pub fn recv_connection_announcements(&mut self) {
|
||||
// Collect all incoming announcements
|
||||
let announcements = self
|
||||
.connections
|
||||
.iter_mut()
|
||||
.map(|c| c.try_read())
|
||||
.flat_map(|array| array)
|
||||
.collect::<Vec<Announcement>>();
|
||||
|
||||
for announcement in announcements {
|
||||
self.recv_announcement(&announcement)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn broadcast(&mut self, announcement: Announcement) -> Result<(), ModuleError> {
|
||||
for connection in &mut self.connections {
|
||||
connection.write(announcement.clone())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,68 +0,0 @@
|
||||
mod manager;
|
||||
mod manager_announcement;
|
||||
mod manager_connection;
|
||||
|
||||
mod module;
|
||||
|
||||
mod proc_load;
|
||||
|
||||
// use std::any::Any;
|
||||
|
||||
// pub use logger::setup_logger;
|
||||
pub use manager::Manager;
|
||||
pub use module::Module;
|
||||
|
||||
/// "Module Interface" helper macro that creates a struct with function pointers
|
||||
/// Useful for defining and requiring modules' functions accross FFI boundry.
|
||||
#[macro_export]
|
||||
macro_rules! module_interface {
|
||||
($(#[$struct_meta:meta])* $interface_name:ident { $($(#[$fn_meta:meta])* fn $fn_name:ident $(<$($gen:ident),+ $(,)?>)?($($arg:ident : $ty:ty),* $(,)?) $(-> $ret:ty)? $(where $($where_clause:tt)*)?);* $(;)? }) => {
|
||||
|
||||
#[repr(C)]
|
||||
#[allow(non_camel_case_types)]
|
||||
#[derive(Clone, Copy)]
|
||||
#[allow(improper_ctypes_definitions)]
|
||||
$(#[$struct_meta])*
|
||||
pub struct $interface_name {
|
||||
$(
|
||||
// This line will FAIL TO COMPILE if you use generics in the macro input.
|
||||
// You MUST use concrete types like *mut c_void for "generic" data.
|
||||
$fn_name: extern "C" fn($($ty),*) $(-> $ret)?,
|
||||
)*
|
||||
}
|
||||
|
||||
impl $interface_name {
|
||||
$(
|
||||
#[inline(always)]
|
||||
$(#[$fn_meta])* // Propagate function attributes
|
||||
// This is the fix for the `impl` block.
|
||||
// It adds the captured generics and where-clause to the wrapper function.
|
||||
pub fn $fn_name $(<$($gen),+>)? (&self, $($arg: $ty),*) $(-> $ret)?
|
||||
$(where $($where_clause)*)?
|
||||
{
|
||||
(self.$fn_name)($($arg),*)
|
||||
}
|
||||
)*
|
||||
|
||||
/// Create from raw function pointers
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// The caller must ensure all function pointers are valid and have
|
||||
/// the correct signatures
|
||||
pub fn from_raw(
|
||||
$($fn_name: extern "C" fn($($ty),*) $(-> $ret)?),*
|
||||
) -> Self {
|
||||
Self {
|
||||
$($fn_name),*
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// impl crate::module::Interface for $interface_name {
|
||||
// fn as_any(self: Box<Self>) -> Box<dyn std::any::Any> {
|
||||
// self
|
||||
// }
|
||||
// }
|
||||
};
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
use libloading::{Library, Symbol};
|
||||
|
||||
use crate::module::proc_load::memfd_create_dlopen;
|
||||
use crate::{ModuleError, logger::SetupLogger, logger::logger};
|
||||
|
||||
use crate::*;
|
||||
|
||||
pub struct Module {
|
||||
lib: Library,
|
||||
}
|
||||
|
||||
impl Module {
|
||||
pub fn new(path: &str) -> Result<Self, ModuleError> {
|
||||
let lib = unsafe { Library::new(&path) }.map_err(|e| ModuleError::LibLoadingError(e))?;
|
||||
|
||||
let this = Self { lib };
|
||||
|
||||
if let Ok(setup_logger) = this.get_symbol::<SetupLogger>(b"setup_logger") {
|
||||
setup_logger(logger());
|
||||
} else {
|
||||
warn!("setup_logger not found");
|
||||
}
|
||||
|
||||
Ok(this)
|
||||
}
|
||||
|
||||
// TODO: Implement actual reflective ELF loading (possibly even custom format)
|
||||
// Look at https://github.com/weizhiao/rust-elfloader
|
||||
pub fn new_bytes(bytes: &[u8]) -> Result<Self, ModuleError> {
|
||||
let lib =
|
||||
memfd_create_dlopen(bytes).map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||
|
||||
let this = Self { lib };
|
||||
|
||||
if let Ok(setup_logger) = this.get_symbol::<SetupLogger>(b"setup_logger") {
|
||||
setup_logger(logger());
|
||||
} else {
|
||||
warn!("setup_logger not found");
|
||||
}
|
||||
|
||||
Ok(this)
|
||||
}
|
||||
pub fn get_symbol<T>(&self, symbol: &[u8]) -> Result<Symbol<'_, T>, ModuleError> {
|
||||
let symbol = unsafe { self.lib.get::<T>(symbol) }
|
||||
.map_err(|e| ModuleError::LinkError(format!("Failed to load symbol: {}", e)))?;
|
||||
|
||||
Ok(symbol)
|
||||
}
|
||||
}
|
||||
@@ -1,56 +0,0 @@
|
||||
// Load a shared object by saving bytes to a filesystem in /proc
|
||||
|
||||
use std::{error::Error, ffi::CString, io}; // 0.8
|
||||
|
||||
use libloading::Library;
|
||||
|
||||
use crate::debug;
|
||||
|
||||
// The `memfd_create` syscall flags (MFD_CLOEXEC is common and good practice)
|
||||
const MFD_CLOEXEC: u32 = 0x0001;
|
||||
const MFD_ALLOW_SEALING: u32 = 0x0002;
|
||||
|
||||
pub fn memfd_create_dlopen(payload: &[u8]) -> Result<Library, Box<dyn Error>> {
|
||||
use rand::distr::{Alphanumeric, SampleString};
|
||||
|
||||
let string = Alphanumeric.sample_string(&mut rand::rng(), 16);
|
||||
|
||||
// 1. Create the anonymous in-memory file descriptor using the raw syscall
|
||||
let c_name = CString::new(string).expect("CString conversion failed");
|
||||
|
||||
let fd = unsafe { libc::memfd_create(c_name.as_ptr(), MFD_CLOEXEC | MFD_ALLOW_SEALING) };
|
||||
|
||||
if fd < 0 {
|
||||
return Err(io::Error::last_os_error().to_string().into());
|
||||
}
|
||||
|
||||
// 2. Write the payload bytes to the in-memory file
|
||||
let bytes_written =
|
||||
unsafe { libc::write(fd, payload.as_ptr() as *const libc::c_void, payload.len()) };
|
||||
|
||||
if bytes_written != payload.len() as isize {
|
||||
// If write fails or is incomplete, clean up the file descriptor
|
||||
unsafe {
|
||||
libc::close(fd);
|
||||
}
|
||||
return Err("Failed to write full payload to memfd".into());
|
||||
}
|
||||
|
||||
// Optional: Seal the file to prevent modification, common for security/integrity
|
||||
// Note: The MFD_ALLOW_SEALING flag must be set during creation for this to work.
|
||||
let seals = libc::F_SEAL_GROW | libc::F_SEAL_SHRINK | libc::F_SEAL_WRITE;
|
||||
if unsafe { libc::fcntl(fd, libc::F_ADD_SEALS, seals) } == -1 {
|
||||
// Log a warning but continue if sealing fails (e.g., due to permissions)
|
||||
debug!(
|
||||
"memfd_create_dlopen: Failed to apply seals. Error: {}",
|
||||
io::Error::last_os_error()
|
||||
);
|
||||
}
|
||||
|
||||
// 3. Construct the virtual path to the in-memory file
|
||||
// This path is necessary for dlopen to work, as dlopen expects a filesystem path.
|
||||
let dl_path = format!("/proc/self/fd/{}", fd);
|
||||
|
||||
// 4. Use dlopen (via libloading) on the virtual path
|
||||
Ok(unsafe { Library::new(&dl_path)? })
|
||||
}
|
||||
@@ -1,31 +0,0 @@
|
||||
// mod connection;
|
||||
mod tcp_stream;
|
||||
pub use tcp_stream::TcpStream;
|
||||
|
||||
// pub use connection::Connection;
|
||||
|
||||
use crate::ModuleError;
|
||||
|
||||
/// This is the data transmission type
|
||||
pub trait Stream<T>: Send + Sync {
|
||||
// fn get_info(&self) -> String;
|
||||
fn is_alive(&self) -> bool;
|
||||
|
||||
fn has_recv(&self) -> bool;
|
||||
|
||||
/// Possibly blocking stream read function
|
||||
fn read(&mut self) -> Vec<T>;
|
||||
|
||||
/// Non-blocking read function
|
||||
fn try_read(&mut self) -> Vec<T> {
|
||||
if self.has_recv() {
|
||||
self.read()
|
||||
} else {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&mut self, data: T) -> Result<(), ModuleError>;
|
||||
|
||||
fn try_clone(&self) -> Result<Box<dyn Stream<T> + Send + Sync>, ModuleError>;
|
||||
}
|
||||
@@ -1,126 +0,0 @@
|
||||
use std::{
|
||||
io::{Read, Write},
|
||||
net,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
};
|
||||
|
||||
use crate::{Announcement, ModuleError, debug, network::Stream};
|
||||
|
||||
pub struct TcpStream(Arc<AtomicBool>, net::TcpStream);
|
||||
|
||||
impl TcpStream {
|
||||
pub fn new(stream: net::TcpStream) -> Self {
|
||||
stream.set_nonblocking(true).unwrap();
|
||||
Self(Arc::new(AtomicBool::new(true)), stream)
|
||||
}
|
||||
|
||||
// Call this when the stream ends
|
||||
fn disconnected(&mut self) {
|
||||
self.0.store(false, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream<Announcement> for TcpStream {
|
||||
fn is_alive(&self) -> bool {
|
||||
// if self.1.take_error().unwrap_or(None).is_some() {
|
||||
// // self.1.pe
|
||||
// warn!("Disconnected #################");
|
||||
// return true;
|
||||
// } else {
|
||||
// return false;
|
||||
// }
|
||||
|
||||
// let mut buf = [0u8; 1];
|
||||
// match self.1.peek(&mut buf) {
|
||||
// Ok(n) => n == 1,
|
||||
// Err(_) => false,
|
||||
// }
|
||||
|
||||
let mut buf = [0u8; 1];
|
||||
match self.1.peek(&mut buf) {
|
||||
Ok(0) => false, // Connection closed (EOF)
|
||||
Ok(_) => true, // Data available or connection alive
|
||||
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => true, // No data but alive
|
||||
Err(_) => false, // Connection error
|
||||
}
|
||||
|
||||
// true
|
||||
|
||||
// self.0.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn has_recv(&self) -> bool {
|
||||
let mut buf = [0u8; 1];
|
||||
match self.1.peek(&mut buf) {
|
||||
Ok(n) if n > 0 => true, // Data is available
|
||||
Ok(_) => false, // EOF (connection closed)
|
||||
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => false, // No data
|
||||
Err(_) => false,
|
||||
}
|
||||
// false
|
||||
}
|
||||
|
||||
fn read(&mut self) -> Vec<Announcement> {
|
||||
let mut ret = Vec::new();
|
||||
|
||||
while self.has_recv() {
|
||||
let mut size_buf = [0u8; 4];
|
||||
match self.1.read_exact(&mut size_buf) {
|
||||
Ok(()) => {}
|
||||
Err(_) => {
|
||||
self.disconnected();
|
||||
break;
|
||||
}
|
||||
};
|
||||
let size = u32::from_be_bytes(size_buf);
|
||||
|
||||
let mut buf = vec![0u8; size as usize];
|
||||
|
||||
match self.1.read_exact(&mut buf) {
|
||||
Ok(()) => {}
|
||||
Err(_) => {
|
||||
self.disconnected();
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(a) = Announcement::decode(&buf) {
|
||||
ret.push(a);
|
||||
} else {
|
||||
debug!("Malformed data");
|
||||
}
|
||||
}
|
||||
|
||||
ret
|
||||
}
|
||||
|
||||
fn write(&mut self, announcement: Announcement) -> Result<(), crate::ModuleError> {
|
||||
let bytes = announcement.encode();
|
||||
|
||||
// Write length of bytes
|
||||
self.1
|
||||
.write_all(&u32::to_be_bytes(bytes.len() as u32))
|
||||
.map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||
// Write data
|
||||
self.1
|
||||
.write_all(&bytes)
|
||||
.map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||
// Flush data
|
||||
self.1
|
||||
.flush()
|
||||
.map_err(|e| ModuleError::Error(e.to_string().into()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn try_clone(&self) -> Result<Box<dyn Stream<Announcement> + Send + Sync>, crate::ModuleError> {
|
||||
Ok(Box::new(Self(
|
||||
self.0.clone(),
|
||||
self.1
|
||||
.try_clone()
|
||||
.map_err(|e| ModuleError::Error(e.to_string().into()))?,
|
||||
)))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user