fixed bugs, memory leaks, etc

injector/crates/loader/src/bridge.rs

@@ -8,6 +8,8 @@ static JVMTI: AtomicPtr<*const *const c_void> = AtomicPtr::new(std::ptr::null_mu
 /// jclass whose bytes the ClassFileLoadHook should capture during the current RetransformClasses call (0 = none).
 static CAPTURE_TARGET: AtomicUsize = AtomicUsize::new(0);
 static CAPTURED: Mutex<Option<Vec<u8>>> = Mutex::new(None);
+/// Serialises classBytes calls: CAPTURE_TARGET holds a thread-local jclass reference, so only one thread may capture at a time.
+static CAPTURE_LOCK: Mutex<()> = Mutex::new(());
 
 /// Acquires the JVMTI capabilities and enables the class file load hook. Called once from the worker thread.
 pub unsafe fn init_jvmti(vm: VmPtr) -> Result<(), String> {
@@ -56,6 +58,7 @@ pub unsafe extern "system" fn native_class_bytes(env: EnvPtr, _cls: JClass, targ
     if jvmti.is_null() {
         return std::ptr::null_mut();
     }
+    let _serial = CAPTURE_LOCK.lock();
     if let Ok(mut slot) = CAPTURED.lock() {
         *slot = None;
     }
@@ -86,6 +89,9 @@ pub unsafe extern "system" fn native_redefine(env: EnvPtr, _cls: JClass, target:
     }
     let jni = Jni::new(env);
     let len = jni.get_array_length(bytes);
+    if len < 0 {
+        return -1;
+    }
     let mut buf = vec![0u8; len as usize];
     jni.get_byte_array_region(bytes, 0, len, buf.as_mut_ptr());
     let def = ClassDefinition { klass: target, class_byte_count: len, class_bytes: buf.as_ptr() };
@@ -104,3 +110,30 @@ pub unsafe extern "system" fn native_log(env: EnvPtr, _cls: JClass, message: JOb
     let text = Jni::new(env).get_string_utf(message);
     crate::log!("[java] {text}");
 }
+
+extern "system" {
+    fn FreeLibraryAndExitThread(module: *mut c_void, exit_code: u32) -> !;
+}
+
+/// void release(): stops the JVMTI hook and unloads this DLL after a short delay (self destruct).
+/// Java must not call any native of NativeBridge afterwards.
+pub unsafe extern "system" fn native_release(env: EnvPtr, cls: JClass) {
+    let jvmti = JVMTI.swap(std::ptr::null_mut(), Ordering::SeqCst);
+    if !jvmti.is_null() {
+        if let Err(e) = disable_class_file_load_hook(jvmti) {
+            crate::log!("release: {e}");
+        }
+    }
+    // Late calls from Java then throw UnsatisfiedLinkError instead of jumping into the unmapped DLL.
+    let rc = Jni::new(env).unregister_natives(cls);
+    if rc != JNI_OK {
+        crate::log!("release: UnregisterNatives failed ({rc})");
+    }
+    crate::log!("[java] released, unloading the loader");
+    let module = crate::MODULE as usize;
+    // The thread below runs code from this DLL until FreeLibraryAndExitThread; the delay lets in-flight native calls return.
+    let _ = std::thread::Builder::new().name("prestige-unload".into()).spawn(move || {
+        std::thread::sleep(std::time::Duration::from_millis(2000));
+        FreeLibraryAndExitThread(module as *mut c_void, 0);
+    });
+}

injector/crates/loader/src/jni.rs

@@ -123,6 +123,9 @@ impl Jni {
     pub unsafe fn register_natives(&self, c: JClass, methods: &[NativeMethod]) -> JInt {
         f!(self, 215, unsafe extern "system" fn(EnvPtr, JClass, *const NativeMethod, JInt) -> JInt)(self.env, c, methods.as_ptr(), methods.len() as JInt)
     }
+    pub unsafe fn unregister_natives(&self, c: JClass) -> JInt {
+        f!(self, 216, unsafe extern "system" fn(EnvPtr, JClass) -> JInt)(self.env, c)
+    }
     pub unsafe fn exception_check(&self) -> bool {
         f!(self, 228, unsafe extern "system" fn(EnvPtr) -> JBool)(self.env) != 0
     }
@@ -168,6 +171,12 @@ pub unsafe fn attach_current_thread(vm: VmPtr) -> Result<EnvPtr, String> {
     Ok(env)
 }
 
+/// Detaches the current thread (it must not hold Java frames); a thread attached from native code has to detach before it ends.
+pub unsafe fn detach_current_thread(vm: VmPtr) {
+    let detach = std::mem::transmute::<*const c_void, unsafe extern "system" fn(VmPtr) -> JInt>(*(*vm).add(5));
+    detach(vm);
+}
+
 pub unsafe fn get_jvmti(vm: VmPtr) -> Result<JvmtiPtr, String> {
     let get_env = std::mem::transmute::<*const c_void, unsafe extern "system" fn(VmPtr, *mut *mut c_void, JInt) -> JInt>(*(*vm).add(6));
     let mut out: *mut c_void = std::ptr::null_mut();
@@ -253,3 +262,19 @@ pub unsafe fn redefine_class(j: JvmtiPtr, def: &ClassDefinition) -> JInt {
     let f = std::mem::transmute::<*const c_void, unsafe extern "system" fn(JvmtiPtr, JInt, *const ClassDefinition) -> JInt>(*(*j).add(86));
     f(j, 1, def)
 }
+
+/// Disables the ClassFileLoadHook event and clears the callback table, so the JVM no longer calls into this DLL.
+pub unsafe fn disable_class_file_load_hook(j: JvmtiPtr) -> Result<(), String> {
+    let set_mode = std::mem::transmute::<*const c_void, unsafe extern "C" fn(JvmtiPtr, JInt, JInt, JObject, ...) -> JInt>(*(*j).add(1));
+    let rc = set_mode(j, 0, EVENT_CLASS_FILE_LOAD_HOOK, std::ptr::null_mut());
+    if rc != 0 {
+        return Err(format!("SetEventNotificationMode(disable) failed ({rc})"));
+    }
+    let callbacks = [0usize; 40];
+    let set_callbacks = std::mem::transmute::<*const c_void, unsafe extern "system" fn(JvmtiPtr, *const usize, JInt) -> JInt>(*(*j).add(121));
+    let rc = set_callbacks(j, callbacks.as_ptr(), std::mem::size_of_val(&callbacks) as JInt);
+    if rc != 0 {
+        return Err(format!("SetEventCallbacks(clear) failed ({rc})"));
+    }
+    Ok(())
+}

injector/crates/loader/src/worker.rs

@@ -1,5 +1,5 @@
 //! The loader worker: JVM discovery, payload installation into Fabric's Knot class loader, native registration, bootstrap.
-use crate::bridge::{init_jvmti, native_class_bytes, native_log, native_redefine};
+use crate::bridge::{init_jvmti, native_class_bytes, native_log, native_redefine, native_release};
 use crate::jni::*;
 use crate::log;
 use std::ffi::{c_char, c_void};
@@ -28,6 +28,13 @@ pub fn run() {
 unsafe fn install() -> Result<(), String> {
     let vm = find_vm()?;
     let env = attach_current_thread(vm)?;
+    struct Detach(VmPtr);
+    impl Drop for Detach {
+        fn drop(&mut self) {
+            unsafe { detach_current_thread(self.0) }
+        }
+    }
+    let _detach = Detach(vm);
     let jni = Jni::new(env);
     log!("attached to the JVM");
 
@@ -160,6 +167,11 @@ unsafe fn register_natives(jni: &Jni, loader: JObject) -> Result<(), String> {
             signature: c"(Ljava/lang/Class;[B)I".as_ptr(),
             fn_ptr: native_redefine as *const () as usize as *mut c_void,
         },
+        NativeMethod {
+            name: c"release".as_ptr(),
+            signature: c"()V".as_ptr(),
+            fn_ptr: native_release as *const () as usize as *mut c_void,
+        },
         NativeMethod {
             name: c"log".as_ptr(),
             signature: c"(Ljava/lang/String;)V".as_ptr(),