From 8a0138d1025b35dbfca73db1326dac7cb122db95 Mon Sep 17 00:00:00 2001
From: Simonas Kazlauskas <git@kazlauskas.me>
Date: Sat, 22 Aug 2020 05:55:54 +0300
Subject: [PATCH] Add os::windows::Library::this

---
 src/error.rs          | 14 ++++++++++++--
 src/os/windows/mod.rs | 41 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index dcbe2b114..80bab7af6 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -35,6 +35,13 @@ pub enum Error {
     LoadLibraryW { source: WindowsError },
     /// The `LoadLibraryW` call failed and system did not report an error.
     LoadLibraryWUnknown,
+    /// The `GetModuleHandleExW` call failed.
+    GetModuleHandleExW {
+        /// The source error.
+        source: WindowsError
+    },
+    /// The `LoadLibraryW` call failed and system did not report an error.
+    GetModuleHandleExWUnknown,
     /// The `GetProcAddress` call failed.
     GetProcAddress { source: WindowsError },
     /// The `GetProcAddressUnknown` call failed and system did not report an error.
@@ -75,9 +82,12 @@ impl std::fmt::Display for Error {
             DlSymUnknown => write!(f, "dlsym failed, but system did not report the error"),
             DlClose { ref desc } => write!(f, "{}", desc.0.to_string_lossy()),
             DlCloseUnknown => write!(f, "dlclose failed, but system did not report the error"),
-            LoadLibraryW { .. } => write!(f, "LoadLibraryW failed"),
+            LoadLibraryW { .. } => write!(f, "LoadLibraryExW failed"),
             LoadLibraryWUnknown =>
-                write!(f, "LoadLibraryW failed, but system did not report the error"),
+                write!(f, "LoadLibraryExW failed, but system did not report the error"),
+            GetModuleHandleExW { .. } => write!(f, "GetModuleHandleExW failed"),
+            GetModuleHandleExWUnknown =>
+                write!(f, "GetModuleHandleExWUnknown failed, but system did not report the error"),
             GetProcAddress { .. } => write!(f, "GetProcAddress failed"),
             GetProcAddressUnknown =>
                 write!(f, "GetProcAddress failed, but system did not report the error"),
diff --git a/src/os/windows/mod.rs b/src/os/windows/mod.rs
index 3620b3e38..898037ad3 100644
--- a/src/os/windows/mod.rs
+++ b/src/os/windows/mod.rs
@@ -39,6 +39,25 @@ impl Library {
         Library::load_with_flags(filename, 0)
     }
 
+    /// Load the `Library` representing the original program executable.
+    ///
+    /// Note that this is not a direct equivalent to [`os::unix::Library::this`].
+    ///
+    /// [`os::unix::Library::this`]: crate::os::unix::Library::this
+    pub fn this() -> Result<Library, crate::Error> {
+        let handle = unsafe {
+            let handle: HMODULE = std::ptr::null_mut();
+            with_get_last_error(|source| crate::Error::GetModuleHandleExW { source }, || {
+                let result = libloaderapi::GetModuleHandleExW(0, std::ptr::null_mut(), &mut handle);
+                if result == 0 {
+                    None
+                } else {
+                    Some(Library(handle))
+                }
+            }).map_err(|e| e.unwrap_or(crate::Error::GetModuleHandleExWUnknown));
+        };
+    }
+
     /// Find and load a shared library (module).
     ///
     /// Locations where library is searched for is platform specific and can’t be adjusted
@@ -52,7 +71,9 @@ impl Library {
         let ret = with_get_last_error(|source| crate::Error::LoadLibraryW { source }, || {
             // Make sure no winapi calls as a result of drop happen inside this closure, because
             // otherwise that might change the return value of the GetLastError.
-            let handle = unsafe { libloaderapi::LoadLibraryExW(wide_filename.as_ptr(), std::ptr::null_mut(), flags) };
+            let handle = unsafe {
+                libloaderapi::LoadLibraryExW(wide_filename.as_ptr(), std::ptr::null_mut(), flags)
+            };
             if handle.is_null()  {
                 None
             } else {
@@ -329,4 +350,22 @@ mod tests {
             assert_eq!(errhandlingapi::GetLastError(), gle())
         }
     }
+
+    #[test]
+    fn library_this_get() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+        static VARIABLE = AtomicBool::new();
+
+        #[no_mangle]
+        extern "C" fn library_this_get_test_fn() {
+            VARIABLE.store(true, Ordering::SeqCst);
+        }
+        let lib = Library::this().expect("this library");
+        let test_fn: Symbol<unsafe extern "C" fn()> = unsafe {
+            lib.get(b"library_this_get_test_fn\0").expect("get symbol");
+        };
+        assert_eq!(VARIABLE.load(Ordering::SeqCst), false);
+        test_fn();
+        assert_eq!(VARIABLE.load(Ordering::SeqCst), true);
+    }
 }