Add more compat_fn machinery

seritools · seritools · commit 16aaea2baec5 · 2025-04-02T23:56:46.000+02:00
diff --git a/library/std/src/sys/pal/windows/c.rs b/library/std/src/sys/pal/windows/c.rs
@@ -5,7 +5,7 @@
 #![unstable(issue = "none", feature = "windows_c")]
 #![allow(clippy::style)]
 
-use core::ffi::{CStr, c_uint, c_ulong, c_ushort, c_void};
+use core::ffi::{c_uint, c_ulong, c_ushort, c_void};
 use core::{mem, ptr};
 
 mod windows_sys;
@@ -122,7 +122,7 @@ extern "system" {
 // Functions that aren't available on every version of Windows that we support,
 // but we still use them and just provide some form of a fallback implementation.
 compat_fn_with_fallback! {
-    pub static KERNEL32: &CStr = c"kernel32";
+    pub static KERNEL32: &CStr = c"kernel32" => { load: false, unicows: false };
 
     // >= Win10 1607
     // https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-setthreaddescription
@@ -189,7 +189,7 @@ compat_fn_optional! {
 
 #[cfg(any(target_vendor = "win7", target_vendor = "uwp"))]
 compat_fn_with_fallback! {
-    pub static NTDLL: &CStr = c"ntdll";
+    pub static NTDLL: &CStr = c"ntdll" => { load: true, unicows: false };
 
     #[cfg(target_vendor = "win7")]
     pub fn NtCreateKeyedEvent(
diff --git a/library/std/src/sys/pal/windows/c/bindings.txt b/library/std/src/sys/pal/windows/c/bindings.txt
@@ -2475,6 +2475,7 @@ Windows.Win32.System.LibraryLoader.GetModuleFileNameW
 Windows.Win32.System.LibraryLoader.GetModuleHandleA
 Windows.Win32.System.LibraryLoader.GetModuleHandleW
 Windows.Win32.System.LibraryLoader.GetProcAddress
+Windows.Win32.System.LibraryLoader.LoadLibraryA
 Windows.Win32.System.Performance.QueryPerformanceCounter
 Windows.Win32.System.Performance.QueryPerformanceFrequency
 Windows.Win32.System.Pipes.CreateNamedPipeW
diff --git a/library/std/src/sys/pal/windows/c/windows_sys.rs b/library/std/src/sys/pal/windows/c/windows_sys.rs
@@ -65,6 +65,7 @@ windows_targets::link!("kernel32.dll" "system" fn GetWindowsDirectoryW(lpbuffer
 windows_targets::link!("kernel32.dll" "system" fn InitOnceBeginInitialize(lpinitonce : *mut INIT_ONCE, dwflags : u32, fpending : *mut BOOL, lpcontext : *mut *mut core::ffi::c_void) -> BOOL);
 windows_targets::link!("kernel32.dll" "system" fn InitOnceComplete(lpinitonce : *mut INIT_ONCE, dwflags : u32, lpcontext : *const core::ffi::c_void) -> BOOL);
 windows_targets::link!("kernel32.dll" "system" fn InitializeProcThreadAttributeList(lpattributelist : LPPROC_THREAD_ATTRIBUTE_LIST, dwattributecount : u32, dwflags : u32, lpsize : *mut usize) -> BOOL);
+windows_targets::link!("kernel32.dll" "system" fn LoadLibraryA(lplibfilename : PCSTR) -> HMODULE);
 windows_targets::link!("kernel32.dll" "system" fn LocalFree(hmem : HLOCAL) -> HLOCAL);
 windows_targets::link!("kernel32.dll" "system" fn LockFileEx(hfile : HANDLE, dwflags : LOCK_FILE_FLAGS, dwreserved : u32, nnumberofbytestolocklow : u32, nnumberofbytestolockhigh : u32, lpoverlapped : *mut OVERLAPPED) -> BOOL);
 windows_targets::link!("kernel32.dll" "system" fn MoveFileExW(lpexistingfilename : PCWSTR, lpnewfilename : PCWSTR, dwflags : MOVE_FILE_FLAGS) -> BOOL);
diff --git a/library/std/src/sys/pal/windows/compat.rs b/library/std/src/sys/pal/windows/compat.rs
@@ -37,7 +37,7 @@ use crate::sys::c;
 // file an issue for discussion; currently we don't guarantee any functionality
 // before main.
 // See https://docs.microsoft.com/en-us/cpp/c-runtime-library/crt-initialization?view=msvc-170
-#[cfg(target_vendor = "win7")]
+#[cfg(any(target_vendor = "win7", target_vendor = "rust9x"))]
 #[used]
 #[link_section = ".CRT$XCT"]
 static INIT_TABLE_ENTRY: unsafe extern "C" fn() = init;
@@ -68,6 +68,21 @@ unsafe extern "C" fn init() {
     load_synch_functions();
 }
 
+#[cfg(target_vendor = "rust9x")]
+unsafe extern "C" fn init() {
+    // In an exe this code is executed before main() so is single threaded.
+    // In a DLL the system's loader lock will be held thereby synchronizing
+    // access. So the same best practices apply here as they do to running in DllMain:
+    // https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices
+    //
+    // DO NOT do anything interesting or complicated in this function! DO NOT call
+    // any Rust functions or CRT functions if those functions touch any global state,
+    // because this function runs during global initialization. For example, DO NOT
+    // do any dynamic allocation, don't call LoadLibrary, etc.
+
+    checks::init_rust9x_checks();
+}
+
 /// Helper macro for creating CStrs from literals and symbol names.
 macro_rules! ansi_str {
     (sym $ident:ident) => {{ crate::sys::compat::const_cstr_from_bytes(concat!(stringify!($ident), "\0").as_bytes()) }};
@@ -118,6 +133,15 @@ impl Module {
         }
     }
 
+    #[allow(dead_code)]
+    pub unsafe fn load(name: &CStr) -> Option<Self> {
+        // SAFETY: A CStr is always null terminated.
+        unsafe {
+            let module = c::LoadLibraryA(name.as_ptr().cast::<u8>());
+            NonNull::new(module).map(Self)
+        }
+    }
+
     // Try to get the address of a function.
     pub fn proc_address(self, name: &CStr) -> Option<NonNull<c_void>> {
         unsafe {
@@ -131,43 +155,54 @@ impl Module {
     }
 }
 
+pub static UNICOWS: &CStr = c"unicows";
+
 /// Load a function or use a fallback implementation if that fails.
 macro_rules! compat_fn_with_fallback {
-    (pub static $module:ident: &CStr = $name:expr; $(
-        $(#[$meta:meta])*
-        $vis:vis fn $symbol:ident($($argname:ident: $argtype:ty),*) -> $rettype:ty $fallback_body:block
-    )*) => (
-        pub static $module: &CStr = $name;
+    {
+        pub static $module:ident: &CStr = $name:expr => { load: $load:expr, unicows: $unicows:expr };
+        $(
+            $(#[$meta:meta])*
+            $vis:vis fn $symbol:ident($($argname:ident: $argtype:ty),* $(,)?) $(-> $rettype:ty)? $fallback_body:block
+        )*
+    } => {
     $(
         $(#[$meta])*
         pub mod $symbol {
             #[allow(unused_imports)]
             use super::*;
             use crate::mem;
-            use crate::ffi::CStr;
+            use crate::ffi::{CStr, c_void};
             use crate::sync::atomic::{AtomicPtr, Ordering};
-            use crate::sys::compat::Module;
 
-            type F = unsafe extern "system" fn($($argtype),*) -> $rettype;
+            type F = unsafe extern "system" fn($($argtype),*) $(-> $rettype)?;
 
             /// `PTR` contains a function pointer to one of three functions.
             /// It starts with the `load` function.
             /// When that is called it attempts to load the requested symbol.
             /// If it succeeds, `PTR` is set to the address of that symbol.
             /// If it fails, then `PTR` is set to `fallback`.
-            static PTR: AtomicPtr<c_void> = AtomicPtr::new(load as *mut _);
+            pub(in crate::sys) static PTR: AtomicPtr<c_void> = AtomicPtr::new(load as *mut _);
 
-            unsafe extern "system" fn load($($argname: $argtype),*) -> $rettype {
+            unsafe extern "system" fn load($($argname: $argtype),*) $(-> $rettype)? {
                 unsafe {
-                    let func = load_from_module(Module::new($module));
+                    let func = load_from_module();
                     func($($argname),*)
                 }
             }
 
-            fn load_from_module(module: Option<Module>) -> F {
+            fn load_from_module() -> F {
                 unsafe {
                     static SYMBOL_NAME: &CStr = ansi_str!(sym $symbol);
-                    if let Some(f) = module.and_then(|m| m.proc_address(SYMBOL_NAME)) {
+
+                    let f = crate::sys::compat::load_from_module(
+                        $name,
+                        SYMBOL_NAME,
+                        $load,
+                        $unicows
+                    );
+
+                    if let Some(f) = f {
                         PTR.store(f.as_ptr(), Ordering::Relaxed);
                         mem::transmute(f)
                     } else {
@@ -177,13 +212,27 @@ macro_rules! compat_fn_with_fallback {
                 }
             }
 
+            #[allow(dead_code)]
+            pub fn available() -> Option<F> {
+                let mut ptr = PTR.load(Ordering::Relaxed);
+                if ptr == load as *mut _ {
+                    ptr = load_from_module() as *mut _;
+                }
+
+                if ptr != fallback as *mut _ {
+                    Some(unsafe { mem::transmute(ptr) })
+                } else {
+                    None
+                }
+            }
+
             #[allow(unused_variables)]
-            unsafe extern "system" fn fallback($($argname: $argtype),*) -> $rettype {
+            unsafe extern "system" fn fallback($($argname: $argtype),*) $(-> $rettype)? {
                 $fallback_body
             }
 
             #[inline(always)]
-            pub unsafe fn call($($argname: $argtype),*) -> $rettype {
+            pub unsafe fn call($($argname: $argtype),*) $(-> $rettype)? {
                 unsafe {
                     let func: F = mem::transmute(PTR.load(Ordering::Relaxed));
                     func($($argname),*)
@@ -193,7 +242,8 @@ macro_rules! compat_fn_with_fallback {
         #[allow(unused)]
         $(#[$meta])*
         $vis use $symbol::call as $symbol;
-    )*)
+    )*
+    }
 }
 
 /// Optionally loaded functions.
@@ -203,7 +253,7 @@ macro_rules! compat_fn_with_fallback {
 macro_rules! compat_fn_optional {
     ($(
         $(#[$meta:meta])*
-        $vis:vis fn $symbol:ident($($argname:ident: $argtype:ty),*) $(-> $rettype:ty)?;
+        $vis:vis fn $symbol:ident($($argname:ident: $argtype:ty),* $(,)?) $(-> $rettype:ty)?;
     )+) => (
         $(
             pub mod $symbol {
@@ -219,8 +269,19 @@ macro_rules! compat_fn_optional {
                 type F = unsafe extern "system" fn($($argtype),*) $(-> $rettype)?;
 
                 #[inline(always)]
+                #[allow(dead_code)]
                 pub fn option() -> Option<F> {
-                    NonNull::new(PTR.load(Ordering::Relaxed)).map(|f| unsafe { mem::transmute(f) })
+                    unsafe {
+                        NonNull::new(PTR.load(Ordering::Relaxed)).map(|f| unsafe { mem::transmute(f) })
+                    }
+                }
+
+                #[inline(always)]
+                #[allow(dead_code)]
+                pub unsafe fn call($($argname: $argtype),*) $(-> $rettype)? {
+                    unsafe {
+                        (mem::transmute::<_, F>(PTR.load(Ordering::Relaxed)))($($argname),*)
+                    }
                 }
             }
             #[inline]
@@ -231,6 +292,28 @@ macro_rules! compat_fn_optional {
     )
 }
 
+pub(crate) fn load_from_module(
+    module_name: &CStr,
+    symbol_name: &CStr,
+    load: bool,
+    unicows: bool,
+) -> Option<NonNull<c_void>> {
+    let in_unicows = if unicows {
+        unsafe { Module::new(UNICOWS).and_then(|m| m.proc_address(symbol_name)) }
+    } else {
+        None
+    };
+
+    in_unicows.or_else(|| {
+        if load {
+            unsafe { Module::new(module_name) }
+        } else {
+            unsafe { Module::load(module_name) }
+        }
+        .and_then(|m| m.proc_address(symbol_name))
+    })
+}
+
 /// Load all needed functions from "api-ms-win-core-synch-l1-2-0".
 #[cfg(target_vendor = "win7")]
 pub(super) fn load_synch_functions() {
@@ -250,6 +333,4 @@ pub(super) fn load_synch_functions() {
         c::WakeByAddressSingle::PTR.store(wake_by_address_single.as_ptr(), Ordering::Relaxed);
         Some(())
     }
-
-    try_load();
 }
