diff --git a/Cargo.lock b/Cargo.lock index 4b590f5..78eba1a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -129,8 +129,9 @@ name = "agent_loader" version = "0.1.0" dependencies = [ "ctor", - "jvmti", - "libloading", + "jni", + "libc", + "libloading 0.8.6", "log", "simplelog", ] @@ -185,6 +186,15 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc7eb209b1518d6bb87b283c20095f5228ecda460da70b44f0802523dea6da04" +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anyhow" version = "1.0.95" @@ -619,6 +629,17 @@ dependencies = [ "libc", ] +[[package]] +name = "chrono" +version = "0.4.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "145052bdd345b87320e369255277e3fb5152762ad123a901ef5c262dd38fe8d2" +dependencies = [ + "iana-time-zone", + "num-traits", + "windows-link", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -627,7 +648,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" dependencies = [ "glob", "libc", - "libloading", + "libloading 0.8.6", ] [[package]] @@ -871,7 +892,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330c60081dcc4c72131f8eb70510f1ac07223e5d4163db481a04a0befcffa412" dependencies = [ - "libloading", + "libloading 0.8.6", ] [[package]] @@ -989,7 +1010,7 @@ dependencies = [ "egui", "glow", "log", - "memoffset", + "memoffset 0.9.1", "wasm-bindgen", "web-sys", ] @@ -1115,6 +1136,16 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "eyre" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd915d99f24784cdc19fd37ef22b97e3ff0ae756c7e492e9fbfe897d61e2aec" +dependencies = [ + "indenter", + "once_cell", +] + [[package]] name = "fastrand" version = "2.1.1" @@ -1269,7 +1300,7 @@ checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", "libc", - "wasi 0.11.0+wasi-snapshot-preview1", + "wasi", ] [[package]] @@ -1314,7 +1345,7 @@ dependencies = [ "glutin_egl_sys", "glutin_glx_sys", "glutin_wgl_sys", - "libloading", + "libloading 0.8.6", "objc2 0.6.2", "objc2-app-kit 0.3.1", "objc2-core-foundation", @@ -1373,6 +1404,12 @@ version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +[[package]] +name = "hermit-abi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" + [[package]] name = "hermit-abi" version = "0.4.0" @@ -1394,6 +1431,30 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "iana-time-zone" +version = "0.1.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33e57f83510bb73707521ebaffa789ec8caf86f9657cad665b092b581d40e9fb" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "idna" version = "0.5.0" @@ -1425,6 +1486,12 @@ dependencies = [ "arrayvec", ] +[[package]] +name = "indenter" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "964de6e86d545b246d84badc0fef527924ace5134f30641c203ef52ba83f58d5" + [[package]] name = "indexmap" version = "2.5.0" @@ -1444,9 +1511,21 @@ dependencies = [ "egui", "log", "proc-maps", + "ptrace-inject", "simplelog", ] +[[package]] +name = "io-lifetimes" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eae7b9aee968036d54dce06cebaefd919e4472e753296daccd6d344e3e2df0c2" +dependencies = [ + "hermit-abi 0.3.9", + "libc", + "windows-sys 0.48.0", +] + [[package]] name = "itertools" version = "0.13.0" @@ -1502,20 +1581,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "jvmti" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7ac79ec7519346b4495a694b0f2387bef5ba4c826da0dbe86900ab296c12fb2" -dependencies = [ - "lazy_static", - "libc", - "serde", - "serde_derive", - "time 0.1.45", - "toml", -] - [[package]] name = "khronos_api" version = "3.1.0" @@ -1524,9 +1589,9 @@ checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" [[package]] name = "lazy_static" -version = "0.2.11" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76f033c7ad61445c5b347c7382dd1237847eb1bce590fe50365dcb33d546be73" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" @@ -1534,6 +1599,16 @@ version = "0.2.175" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543" +[[package]] +name = "libloading" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" +dependencies = [ + "cfg-if", + "winapi", +] + [[package]] name = "libloading" version = "0.8.6" @@ -1566,6 +1641,12 @@ dependencies = [ "redox_syscall 0.4.1", ] +[[package]] +name = "linux-raw-sys" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f051f77a7c8e6957c0696eac88f26b0117e54f52d3fc682ab19397a8812846a4" + [[package]] name = "linux-raw-sys" version = "0.4.14" @@ -1624,6 +1705,15 @@ dependencies = [ "libc", ] +[[package]] +name = "memoffset" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" +dependencies = [ + "autocfg", +] + [[package]] name = "memoffset" version = "0.9.1" @@ -1666,7 +1756,7 @@ checksum = "78bed444cc8a2160f01cbcf811ef18cac863ad68ae8ca62092e8db51d51c761c" dependencies = [ "libc", "log", - "wasi 0.11.0+wasi-snapshot-preview1", + "wasi", "windows-sys 0.59.0", ] @@ -1700,6 +1790,31 @@ dependencies = [ "jni-sys", ] +[[package]] +name = "nix" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f346ff70e7dbfd675fe90590b92d59ef2de15a8779ae305ebcbfd3f0caf59be4" +dependencies = [ + "autocfg", + "bitflags 1.3.2", + "cfg-if", + "libc", + "memoffset 0.6.5", + "pin-utils", +] + +[[package]] +name = "nix" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "598beaf3cc6fdd9a5dfb1630c2800c7acd31df7aaf0f565796fba2b53ca1af1b" +dependencies = [ + "bitflags 1.3.2", + "cfg-if", + "libc", +] + [[package]] name = "nix" version = "0.29.0" @@ -1710,7 +1825,7 @@ dependencies = [ "cfg-if", "cfg_aliases", "libc", - "memoffset", + "memoffset 0.9.1", ] [[package]] @@ -2095,6 +2210,18 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "pete" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "493eb0f2954de9a30ee85b01891f6519e9d065464eaab79aff65d949f44b7850" +dependencies = [ + "libc", + "memoffset 0.6.5", + "nix 0.25.1", + "thiserror", +] + [[package]] name = "pin-project" version = "1.1.5" @@ -2165,7 +2292,7 @@ checksum = "cc2790cd301dec6cd3b7a025e4815cf825724a51c98dccfe6a3e55f05ffb6511" dependencies = [ "cfg-if", "concurrent-queue", - "hermit-abi", + "hermit-abi 0.4.0", "pin-project-lite", "rustix 0.38.37", "tracing", @@ -2229,6 +2356,35 @@ dependencies = [ "winapi", ] +[[package]] +name = "procfs" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "943ca7f9f29bab5844ecd8fdb3992c5969b6622bb9609b9502fef9b4310e3f1f" +dependencies = [ + "bitflags 1.3.2", + "byteorder", + "chrono", + "flate2", + "hex", + "lazy_static", + "rustix 0.36.17", +] + +[[package]] +name = "ptrace-inject" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e38d641527d8ff055f523e6c2b87cdce2df9b9d172d2b165e22ed012afc6e897" +dependencies = [ + "eyre", + "libloading 0.7.4", + "log", + "nix 0.26.4", + "pete", + "procfs", +] + [[package]] name = "quick-xml" version = "0.30.0" @@ -2358,6 +2514,20 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +[[package]] +name = "rustix" +version = "0.36.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "305efbd14fde4139eb501df5f136994bb520b033fa9fbdce287507dc23b8c7ed" +dependencies = [ + "bitflags 1.3.2", + "errno", + "io-lifetimes", + "libc", + "linux-raw-sys 0.1.4", + "windows-sys 0.45.0", +] + [[package]] name = "rustix" version = "0.38.37" @@ -2528,7 +2698,7 @@ checksum = "16257adbfaef1ee58b1363bdc0664c9b8e1e30aed86049635fb5f147d065a9c0" dependencies = [ "log", "termcolor", - "time 0.3.37", + "time", ] [[package]] @@ -2665,17 +2835,6 @@ dependencies = [ "syn", ] -[[package]] -name = "time" -version = "0.1.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b797afad3f312d1c66a56d11d0316f916356d11bd158fbc6ca6389ff6bf805a" -dependencies = [ - "libc", - "wasi 0.10.0+wasi-snapshot-preview1", - "winapi", -] - [[package]] name = "time" version = "0.3.37" @@ -2749,15 +2908,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" -[[package]] -name = "toml" -version = "0.4.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "758664fc71a3a69038656bee8b6be6477d2a6c315a6b81f7081f591bffa4111f" -dependencies = [ - "serde", -] - [[package]] name = "toml_datetime" version = "0.6.8" @@ -2824,7 +2974,7 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89daebc3e6fd160ac4aa9fc8b3bf71e1f74fbf92367ae71fb83a037e8bf164b9" dependencies = [ - "memoffset", + "memoffset 0.9.1", "tempfile", "winapi", ] @@ -2883,12 +3033,6 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "wasi" -version = "0.10.0+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a143597ca7c7793eff794def352d41792a93c481eb1042423ff7ff72ba2c31f" - [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" @@ -3185,6 +3329,12 @@ dependencies = [ "syn", ] +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + [[package]] name = "windows-result" version = "0.2.0" @@ -3213,6 +3363,15 @@ dependencies = [ "windows-targets 0.42.2", ] +[[package]] +name = "windows-sys" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +dependencies = [ + "windows-targets 0.48.5", +] + [[package]] name = "windows-sys" version = "0.52.0" @@ -3490,7 +3649,7 @@ dependencies = [ "as-raw-xcb-connection", "gethostname", "libc", - "libloading", + "libloading 0.8.6", "once_cell", "rustix 0.38.37", "x11rb-protocol", @@ -3565,7 +3724,7 @@ dependencies = [ "futures-sink", "futures-util", "hex", - "nix", + "nix 0.29.0", "ordered-stream", "rand", "serde", diff --git a/agent_loader/Cargo.toml b/agent_loader/Cargo.toml index a99ef63..bfdf399 100644 --- a/agent_loader/Cargo.toml +++ b/agent_loader/Cargo.toml @@ -17,5 +17,5 @@ ctor = "0.2.8" log = "0.4.25" simplelog = "0.12.2" libloading = "0.8.0" -jvmti = "0.5.0" - +jni = "0.21" +libc = "0.2" \ No newline at end of file diff --git a/agent_loader/src/lib.rs b/agent_loader/src/lib.rs index 3772b23..4973495 100644 --- a/agent_loader/src/lib.rs +++ b/agent_loader/src/lib.rs @@ -3,8 +3,6 @@ extern crate log; extern crate simplelog; use ctor::*; -use jvmti::agent::Agent; -use jvmti::native::jvmti_native::{jsize, JNI_GetCreatedJavaVMs, JavaVM}; use libloading::{Library, Symbol}; use log::{error, info, LevelFilter}; use simplelog::{Config, WriteLogger}; @@ -16,10 +14,13 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Mutex, OnceLock}; use std::time::Duration; use std::{path, thread}; +use jni::JavaVM; +use jni::sys::{jsize, JNI_GetCreatedJavaVMs, JNI_OK}; // Global variable to keep track of the loaded library static CLIENT_LIBRARY: OnceLock>> = OnceLock::new(); static RUNNING: AtomicBool = AtomicBool::new(true); +static JVM_MONITOR: OnceLock> = OnceLock::new(); // Function called when the agent is loaded #[no_mangle] @@ -38,39 +39,16 @@ fn agent_onload() { // Initialize the global variable for the library CLIENT_LIBRARY.get_or_init(|| Mutex::new(None)); - unsafe { - if let Err(e) = register_agent() { - error!("Failed to register agent: {}", e); - return; - } - } + // Setup signal handlers for clean shutdown + setup_signal_handlers(); + + // Start monitoring the JVM + start_jvm_monitor(); // Start the socket server for commands start_command_server(); } -unsafe fn register_agent() -> Result<(), &'static str> { - let mut java_vm: *mut JavaVM = std::ptr::null_mut(); - let mut count: jsize = 0; - - if JNI_GetCreatedJavaVMs(&mut java_vm, 1, &mut count) != 0 || count == 0 { - return Err("Failed to get Java VMs"); - } - - let mut agent = Agent::new(java_vm); - - agent.on_vm_death(Some(on_vm_death)); - - agent.update(); - - Ok(()) -} - -fn on_vm_death() { - info!("VM is shutting down"); - agent_onunload(); -} - // Function called when the agent is unloaded #[no_mangle] #[dtor] @@ -87,6 +65,122 @@ fn agent_onunload() { } } +// Setup signal handlers to detect process termination +fn setup_signal_handlers() { + use std::sync::atomic::AtomicBool; + + static SIGNAL_HANDLER_INSTALLED: AtomicBool = AtomicBool::new(false); + + if SIGNAL_HANDLER_INSTALLED.swap(true, Ordering::SeqCst) { + return; // Already installed + } + + #[cfg(unix)] + { + extern "C" fn handle_signal(_: libc::c_int) { + info!("Received termination signal - cleaning up"); + agent_onunload(); + std::process::exit(0); + } + + unsafe { + libc::signal(libc::SIGTERM, handle_signal as libc::sighandler_t); + libc::signal(libc::SIGINT, handle_signal as libc::sighandler_t); + } + + info!("Signal handlers installed"); + } +} + +// Monitor the JVM status with multiple detection methods +fn start_jvm_monitor() { + let handle = thread::spawn(|| { + info!("JVM monitor thread started"); + + // Wait for JVM to be available + let jvm = loop { + if !RUNNING.load(Ordering::SeqCst) { + return; + } + + match get_jvm() { + Some(vm) => break vm, + None => { + thread::sleep(Duration::from_millis(500)); + } + } + }; + + info!("JVM detected, monitoring started"); + + // Monitor JVM health with multiple checks + let mut consecutive_failures = 0; + let max_failures = 3; + + while RUNNING.load(Ordering::SeqCst) { + thread::sleep(Duration::from_millis(500)); + + // Method 1: Try to attach to the JVM + let attach_ok = jvm.attach_current_thread_as_daemon().is_ok(); + + // Method 2: Check if we can access Java classes + let classes_ok = if attach_ok { + if let Ok(mut env) = jvm.attach_current_thread_as_daemon() { + env.find_class("java/lang/System").is_ok() + } else { + false + } + } else { + false + }; + + // Method 3: Check if the JVM pointer is still valid + let jvm_valid = { + let jvm_ptr = jvm.get_java_vm_pointer(); + !jvm_ptr.is_null() + }; + + if !attach_ok || !classes_ok || !jvm_valid { + consecutive_failures += 1; + info!( + "JVM health check failed ({}/{}): attach={}, classes={}, valid={}", + consecutive_failures, max_failures, attach_ok, classes_ok, jvm_valid + ); + + if consecutive_failures >= max_failures { + info!("JVM appears to be shutting down or dead"); + on_vm_death(); + break; + } + } else { + consecutive_failures = 0; + } + } + + info!("JVM monitor thread stopped"); + }); + + JVM_MONITOR.set(handle).ok(); +} + +fn get_jvm() -> Option { + unsafe { + let mut java_vm: *mut jni::sys::JavaVM = std::ptr::null_mut(); + let mut count: jsize = 0; + + if JNI_GetCreatedJavaVMs(&mut java_vm, 1, &mut count) != JNI_OK || count == 0 { + return None; + } + + JavaVM::from_raw(java_vm).ok() + } +} + +fn on_vm_death() { + info!("VM death detected - initiating cleanup"); + agent_onunload(); +} + // Function to load the client library fn load_client_library(lib_path: &str) -> Result<(), Box> { let client_path = PathBuf::from(lib_path); @@ -183,7 +277,8 @@ fn reload_client_library(lib_path: &str) -> Result<(), Box Result<(), Error> { // First time: load the agent_loader let loader_path = PathBuf::from(format!("{}.so", AGENT_NAME)); let lib_path = PathBuf::from(format!("{}.so", LIBRARY_NAME)); - // Check if agent_loader is already loaded - if !find_library(pid, "agent_loader") { + if !find_library(pid, format!("{}.so", AGENT_NAME).as_str()) { info!("Loading Agent Loader"); - // Load agent_loader via JVMTI - match Command::new("jcmd") - .arg(pid.to_string()) - .arg("JVMTI.agent_load") - .arg(format!("{:?}", path::absolute(&loader_path)?)) - .output() - { - Ok(output) if output.status.success() => { - info!("Agent Loader loaded via jcmd: {:?}", loader_path); + let proc = match Process::get(pid) { + Ok(p) => p, + Err(e) => { + error!("Failed to get Process for pid {}: {:?}", pid, e); + return Err(Error::new(std::io::ErrorKind::Other, format!("Process::get failed: {:?}", e))); } - Ok(output) => { - error!( - "jcmd failed (stderr): {}", - String::from_utf8_lossy(&output.stderr) - ); + }; + + match Injector::attach(proc) { + Ok(mut injector) => { + match injector.inject(&loader_path) { + Ok(_) => { + info!("Successfully injected library: {}", loader_path.to_string_lossy()); + } + Err(e) => { + error!("Injection failed: {:?}", e); + return Err(Error::new(std::io::ErrorKind::Other, e.to_string())); + } + } } Err(e) => { - error!("Unable to execute jcmd: {:?}", e); + error!("Failed to attach to pid {}: {:?}", pid, e); + return Err(Error::new(std::io::ErrorKind::Other, e.to_string())); } } @@ -131,13 +136,14 @@ pub fn find_pid() -> Option { fn find_library(pid: u32, lib_name: &str) -> bool { let maps = get_process_maps(pid as i32).ok(); if maps.is_none() { + error!("Failed to get process maps"); return false; } let maps = maps.unwrap(); for map in maps { if let Some(path) = map.filename() { - if path.ends_with(format!("{}.so", lib_name)) { + if path.ends_with(lib_name) { // Library loaded return true; }