We have a contract to let the library unload at process exit

MolotovCherry · Nov 14, 2024 · 27e48f1 · 27e48f1
1 parent a417dda
commit 27e48f1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 61 deletions.
diff --git a/crates/loader/src/lib.rs b/crates/loader/src/lib.rs
@@ -34,7 +34,7 @@ use windows::{
 use client::{TrySend as _, CLIENT};
 use loader::load_plugins;
 use logging::setup_logging;
-use utils::{FreeSelfLibrary, HInstance, Plugin};
+use utils::Plugin;
 
 declare_plugin! {
     "Loader",
@@ -43,11 +43,10 @@ declare_plugin! {
 }
 
 static LOADED_PLUGINS: LazyLock<Mutex<Vec<Plugin>>> = LazyLock::new(Mutex::default);
-static MODULE: OnceLock<HInstance> = OnceLock::new();
 
 #[unsafe(no_mangle)]
 extern "stdcall-unwind" fn DllMain(
-    module: HINSTANCE,
+    _module: HINSTANCE,
     fdw_reason: u32,
     _lpv_reserved: *const c_void,
 ) -> bool {
@@ -59,7 +58,6 @@ extern "stdcall-unwind" fn DllMain(
             // > Consider calling DisableThreadLibraryCalls when receiving DLL_PROCESS_ATTACH, unless your DLL is
             // > linked with static C run-time library (CRT).
             // _ = unsafe { DisableThreadLibraryCalls(module) };
-            _ = MODULE.set(HInstance(module));
 
             if !is_yabg3nml() {
                 unsupported_operation();
@@ -96,14 +94,6 @@ unsafe extern "system-unwind" fn Init(data: &ThreadData) -> u32 {
         assert!(size_of::<&ThreadData>() == size_of::<usize>());
     }
 
-    let module = *MODULE.get().unwrap();
-    // incref self library to ensure it cannot be unloaded (and plugins unloaded) until all the plugins Init is done
-    let _guard = FreeSelfLibrary::new(module.0);
-    if let Err(e) = _guard {
-        error!("can't create freeselflibrary: {e}");
-        return 0;
-    }
-
     if !is_yabg3nml() {
         unsupported_operation();
         return 0;
@@ -118,6 +108,7 @@ unsafe extern "system-unwind" fn Init(data: &ThreadData) -> u32 {
 
         setup_logging(&data.log).context("failed to setup logging")?;
 
+        // blocking call which waits for all plugins to finish DllMain/Init
         load_plugins()?;
 
         Ok::<_, Error>(())

diff --git a/crates/loader/src/utils.rs b/crates/loader/src/utils.rs
@@ -1,13 +1,4 @@
-use eyre::Result;
-use windows::{
-    core::PCSTR,
-    Win32::{
-        Foundation::{FreeLibrary, HINSTANCE, HMODULE},
-        System::LibraryLoader::{
-            FreeLibraryAndExitThread, GetModuleHandleExA, GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
-        },
-    },
-};
+use windows::Win32::Foundation::{FreeLibrary, HMODULE};
 
 /// Container for a loaded plugin. Frees itself on drop
 #[derive(Default)]
@@ -20,42 +11,3 @@ impl Drop for Plugin {
         _ = unsafe { FreeLibrary(self.0) };
     }
 }
-
-/// Wrapper for HINSTANCE being thread safe
-#[derive(Copy, Clone, Default)]
-pub struct HInstance(pub HINSTANCE);
-unsafe impl Send for HInstance {}
-unsafe impl Sync for HInstance {}
-
-/// This increfs self dll, decrefs on drop
-/// This ensures that while we hold the guard, our library cannot be freed
-///
-/// Warning, instantly shuts down thread after exit to avoid last decref
-/// causing code to disappear. This means things like handle.join() on a
-/// thread will break rust. NO CODE after the drop of this is executed
-///
-/// https://devblogs.microsoft.com/oldnewthing/20131105-00/?p=2733
-pub struct FreeSelfLibrary(HMODULE);
-unsafe impl Send for FreeSelfLibrary {}
-unsafe impl Sync for FreeSelfLibrary {}
-
-impl FreeSelfLibrary {
-    pub fn new(inst: HINSTANCE) -> Result<Self> {
-        let mut hmod = HMODULE::default();
-        unsafe {
-            GetModuleHandleExA(
-                GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
-                PCSTR(inst.0.cast()),
-                &mut hmod,
-            )?;
-        }
-
-        Ok(Self(hmod))
-    }
-}
-
-impl Drop for FreeSelfLibrary {
-    fn drop(&mut self) {
-        unsafe { FreeLibraryAndExitThread(self.0, 0) };
-    }
-}