Open a BCrypt algorithm handle

Author: ChrisDenton

library/std/src/sys/windows/c.rs

@@ -66,6 +66,7 @@ pub type LPSYSTEM_INFO = *mut SYSTEM_INFO;
 pub type LPWSABUF = *mut WSABUF;
 pub type LPWSAOVERLAPPED = *mut c_void;
 pub type LPWSAOVERLAPPED_COMPLETION_ROUTINE = *mut c_void;
+pub type BCRYPT_ALG_HANDLE = LPVOID;
 
 pub type PCONDITION_VARIABLE = *mut CONDITION_VARIABLE;
 pub type PLARGE_INTEGER = *mut c_longlong;
@@ -278,14 +279,16 @@ pub const STATUS_INVALID_PARAMETER: NTSTATUS = 0xc000000d_u32 as _;
 pub const STATUS_PENDING: NTSTATUS = 0x103 as _;
 pub const STATUS_END_OF_FILE: NTSTATUS = 0xC0000011_u32 as _;
 pub const STATUS_NOT_IMPLEMENTED: NTSTATUS = 0xC0000002_u32 as _;
+pub const STATUS_NOT_SUPPORTED: NTSTATUS = 0xC00000BB_u32 as _;
 
 // Equivalent to the `NT_SUCCESS` C preprocessor macro.
 // See: https://docs.microsoft.com/en-us/windows-hardware/drivers/kernel/using-ntstatus-values
 pub fn nt_success(status: NTSTATUS) -> bool {
     status >= 0
 }
 
-pub const BCRYPT_RNG_ALG_HANDLE: usize = 0x81;
+// "RNG\0"
+pub const BCRYPT_RNG_ALGORITHM: &[u16] = &[b'R' as u16, b'N' as u16, b'G' as u16, 0];
 
 #[repr(C)]
 pub struct UNICODE_STRING {
@@ -1229,11 +1232,18 @@ extern "system" {
     // >= Vista / Server 2008
     // https://docs.microsoft.com/en-us/windows/win32/api/bcrypt/nf-bcrypt-bcryptgenrandom
     pub fn BCryptGenRandom(
-        hAlgorithm: LPVOID,
+        hAlgorithm: BCRYPT_ALG_HANDLE,
         pBuffer: *mut u8,
         cbBuffer: ULONG,
         dwFlags: ULONG,
     ) -> NTSTATUS;
+    pub fn BCryptOpenAlgorithmProvider(
+        phalgorithm: *mut BCRYPT_ALG_HANDLE,
+        pszAlgId: LPCWSTR,
+        pszimplementation: LPCWSTR,
+        dwflags: ULONG,
+    ) -> NTSTATUS;
+    pub fn BCryptCloseAlgorithmProvider(hAlgorithm: BCRYPT_ALG_HANDLE, dwFlags: ULONG) -> NTSTATUS;
 }
 
 // Functions that aren't available on every version of Windows that we support,

library/std/src/sys/windows/rand.rs

@@ -22,7 +22,6 @@
 //! [`RtlGenRandom`]: https://docs.microsoft.com/en-us/windows/win32/api/ntsecapi/nf-ntsecapi-rtlgenrandom
 //! [`BCryptGenRandom`]: https://docs.microsoft.com/en-us/windows/win32/api/bcrypt/nf-bcrypt-bcryptgenrandom
 //! [Pseudo-handle]: https://docs.microsoft.com/en-us/windows/win32/seccng/cng-algorithm-pseudo-handles
-use crate::io;
 use crate::mem;
 use crate::ptr;
 use crate::sys::c;
@@ -34,35 +33,82 @@ use crate::sys::c;
 /// [`HashMap`]: crate::collections::HashMap
 /// [`RandomState`]: crate::collections::hash_map::RandomState
 pub fn hashmap_random_keys() -> (u64, u64) {
-    let mut v = (0, 0);
-    let ret = unsafe {
-        let size = mem::size_of_val(&v).try_into().unwrap();
-        c::BCryptGenRandom(
-            // BCRYPT_RNG_ALG_HANDLE is only supported in Windows 10+.
-            // So for Windows 8.1 and Windows 7 we'll need a fallback when this fails.
-            ptr::invalid_mut(c::BCRYPT_RNG_ALG_HANDLE),
-            ptr::addr_of_mut!(v).cast(),
-            size,
-            0,
-        )
-    };
-    if ret != 0 { fallback_rng() } else { v }
+    Rng::open().and_then(|rng| rng.gen_random_keys()).unwrap_or_else(fallback_rng)
+}
+
+struct Rng(c::BCRYPT_ALG_HANDLE);
+impl Rng {
+    // Open a handle to the RNG algorithm.
+    fn open() -> Result<Self, c::NTSTATUS> {
+        use crate::sync::atomic::AtomicPtr;
+        use crate::sync::atomic::Ordering::{Acquire, Release};
+        const ERROR_VALUE: c::LPVOID = ptr::invalid_mut(usize::MAX);
+
+        // An atomic is used so we don't need to reopen the handle every time.
+        static HANDLE: AtomicPtr<crate::ffi::c_void> = AtomicPtr::new(ptr::null_mut());
+
+        let mut handle = HANDLE.load(Acquire);
+        // We use a sentinel value to designate an error occurred last time.
+        if handle == ERROR_VALUE {
+            Err(c::STATUS_NOT_SUPPORTED)
+        } else if handle.is_null() {
+            let status = unsafe {
+                c::BCryptOpenAlgorithmProvider(
+                    &mut handle,
+                    c::BCRYPT_RNG_ALGORITHM.as_ptr(),
+                    ptr::null(),
+                    0,
+                )
+            };
+            if c::nt_success(status) {
+                // If another thread opens a handle first then use that handle instead.
+                let result = HANDLE.compare_exchange(ptr::null_mut(), handle, Release, Acquire);
+                if let Err(previous_handle) = result {
+                    // Close our handle and return the previous one.
+                    unsafe { c::BCryptCloseAlgorithmProvider(handle, 0) };
+                    handle = previous_handle;
+                }
+                Ok(Self(handle))
+            } else {
+                HANDLE.store(ERROR_VALUE, Release);
+                Err(status)
+            }
+        } else {
+            Ok(Self(handle))
+        }
+    }
+
+    fn gen_random_keys(self) -> Result<(u64, u64), c::NTSTATUS> {
+        let mut v = (0, 0);
+        let status = unsafe {
+            let size = mem::size_of_val(&v).try_into().unwrap();
+            c::BCryptGenRandom(self.0, ptr::addr_of_mut!(v).cast(), size, 0)
+        };
+        if c::nt_success(status) { Ok(v) } else { Err(status) }
+    }
 }
 
 /// Generate random numbers using the fallback RNG function (RtlGenRandom)
 #[cfg(not(target_vendor = "uwp"))]
 #[inline(never)]
-fn fallback_rng() -> (u64, u64) {
+fn fallback_rng(rng_status: c::NTSTATUS) -> (u64, u64) {
     let mut v = (0, 0);
     let ret =
         unsafe { c::RtlGenRandom(&mut v as *mut _ as *mut u8, mem::size_of_val(&v) as c::ULONG) };
 
-    if ret != 0 { v } else { panic!("fallback RNG broken: {}", io::Error::last_os_error()) }
+    if ret != 0 {
+        v
+    } else {
+        panic!(
+            "RNG broken: {rng_status:#x}, fallback RNG broken: {}",
+            crate::io::Error::last_os_error()
+        )
+    }
 }
 
 /// We can't use RtlGenRandom with UWP, so there is no fallback
 #[cfg(target_vendor = "uwp")]
 #[inline(never)]
-fn fallback_rng() -> (u64, u64) {
-    panic!("fallback RNG broken: RtlGenRandom() not supported on UWP");
+fn fallback_rng(rng_status: c::NTSTATUS) -> (u64, u64) {
+    panic!("RNG broken: {rng_status:#x} fallback RNG broken: RtlGenRandom() not supported on UWP");
 }