Unify lazy atomics in entropy backends

src/backends/efi_rng.rs

@@ -2,8 +2,7 @@
 use crate::Error;
 use core::{
     mem::MaybeUninit,
-    ptr::{self, NonNull, null_mut},
-    sync::atomic::{AtomicPtr, Ordering::Relaxed},
+    ptr::{self, NonNull},
 };
 use r_efi::{
     efi::{BootServices, Handle},
@@ -17,8 +16,6 @@ pub use crate::util::{inner_u32, inner_u64};
 #[cfg(not(target_os = "uefi"))]
 compile_error!("`efi_rng` backend can be enabled only for UEFI targets!");
 
-static RNG_PROTOCOL: AtomicPtr<rng::Protocol> = AtomicPtr::new(null_mut());
-
 #[cold]
 #[inline(never)]
 fn init() -> Result<NonNull<rng::Protocol>, Error> {
@@ -36,7 +33,7 @@ fn init() -> Result<NonNull<rng::Protocol>, Error> {
         ((*boot_services.as_ptr()).locate_handle)(
             r_efi::efi::BY_PROTOCOL,
             &mut guid,
-            null_mut(),
+            ptr::null_mut(),
             &mut buf_size,
             handles.as_mut_ptr(),
         )
@@ -88,18 +85,19 @@ fn init() -> Result<NonNull<rng::Protocol>, Error> {
             continue;
         }
 
-        RNG_PROTOCOL.store(protocol.as_ptr(), Relaxed);
         return Ok(protocol);
     }
     Err(Error::NO_RNG_HANDLE)
 }
 
 #[inline]
 pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
-    let protocol = match NonNull::new(RNG_PROTOCOL.load(Relaxed)) {
-        Some(p) => p,
-        None => init()?,
-    };
+    #[path = "../utils/lazy_ptr.rs"]
+    mod lazy;
+
+    static RNG_PROTOCOL: lazy::LazyPtr<rng::Protocol> = lazy::LazyPtr::new();
+
+    let protocol = RNG_PROTOCOL.try_unsync_init(init)?;
 
     let mut alg_guid = rng::ALGORITHM_RAW;
     let ret = unsafe {

src/backends/linux_android_with_fallback.rs

@@ -4,8 +4,7 @@ use crate::Error;
 use core::{
     ffi::c_void,
     mem::{MaybeUninit, transmute},
-    ptr::NonNull,
-    sync::atomic::{AtomicPtr, Ordering},
+    ptr::{self, NonNull},
 };
 use use_file::utils;
 
@@ -17,18 +16,12 @@ type GetRandomFn = unsafe extern "C" fn(*mut c_void, libc::size_t, libc::c_uint)
 /// or not supported by kernel.
 const NOT_AVAILABLE: NonNull<c_void> = unsafe { NonNull::new_unchecked(usize::MAX as *mut c_void) };
 
-static GETRANDOM_FN: AtomicPtr<c_void> = AtomicPtr::new(core::ptr::null_mut());
-
 #[cold]
 #[inline(never)]
 fn init() -> NonNull<c_void> {
     // Use static linking to `libc::getrandom` on MUSL targets and `dlsym` everywhere else
     #[cfg(not(target_env = "musl"))]
-    let raw_ptr = {
-        static NAME: &[u8] = b"getrandom\0";
-        let name_ptr = NAME.as_ptr().cast::<libc::c_char>();
-        unsafe { libc::dlsym(libc::RTLD_DEFAULT, name_ptr) }
-    };
+    let raw_ptr = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"getrandom".as_ptr()) };
     #[cfg(target_env = "musl")]
     let raw_ptr = {
         let fptr: GetRandomFn = libc::getrandom;
@@ -37,10 +30,9 @@ fn init() -> NonNull<c_void> {
 
     let res_ptr = match NonNull::new(raw_ptr) {
         Some(fptr) => {
-            let getrandom_fn = unsafe { transmute::<NonNull<c_void>, GetRandomFn>(fptr) };
-            let dangling_ptr = NonNull::dangling().as_ptr();
+            let getrandom_fn = unsafe { transmute::<*mut c_void, GetRandomFn>(fptr.as_ptr()) };
             // Check that `getrandom` syscall is supported by kernel
-            let res = unsafe { getrandom_fn(dangling_ptr, 0, 0) };
+            let res = unsafe { getrandom_fn(ptr::dangling_mut(), 0, 0) };
             if cfg!(getrandom_test_linux_fallback) {
                 NOT_AVAILABLE
             } else if res.is_negative() {
@@ -65,7 +57,6 @@ fn init() -> NonNull<c_void> {
         panic!("Fallback is triggered with enabled `getrandom_test_linux_without_fallback`")
     }
 
-    GETRANDOM_FN.store(res_ptr.as_ptr(), Ordering::Release);
     res_ptr
 }
 
@@ -77,23 +68,17 @@ fn use_file_fallback(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
 
 #[inline]
 pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
-    // Despite being only a single atomic variable, we still cannot always use
-    // Ordering::Relaxed, as we need to make sure a successful call to `init`
-    // is "ordered before" any data read through the returned pointer (which
-    // occurs when the function is called). Our implementation mirrors that of
-    // the one in libstd, meaning that the use of non-Relaxed operations is
-    // probably unnecessary.
-    let raw_ptr = GETRANDOM_FN.load(Ordering::Acquire);
-    let fptr = match NonNull::new(raw_ptr) {
-        Some(p) => p,
-        None => init(),
-    };
+    #[path = "../utils/lazy_ptr.rs"]
+    mod lazy;
+
+    static GETRANDOM_FN: lazy::LazyPtr<c_void> = lazy::LazyPtr::new();
+    let fptr = GETRANDOM_FN.unsync_init(init);
 
     if fptr == NOT_AVAILABLE {
         use_file_fallback(dest)
     } else {
         // note: `transmute` is currently the only way to convert a pointer into a function reference
-        let getrandom_fn = unsafe { transmute::<NonNull<c_void>, GetRandomFn>(fptr) };
+        let getrandom_fn = unsafe { transmute::<*mut c_void, GetRandomFn>(fptr.as_ptr()) };
         utils::sys_fill_exact(dest, |buf| unsafe {
             getrandom_fn(buf.as_mut_ptr().cast(), buf.len(), 0)
         })

src/backends/netbsd.rs

@@ -8,8 +8,7 @@ use core::{
     cmp,
     ffi::c_void,
     mem::{self, MaybeUninit},
-    ptr,
-    sync::atomic::{AtomicPtr, Ordering},
+    ptr::{self, NonNull},
 };
 
 pub use crate::util::{inner_u32, inner_u64};
@@ -42,36 +41,28 @@ unsafe extern "C" fn polyfill_using_kern_arand(
 
 type GetRandomFn = unsafe extern "C" fn(*mut c_void, libc::size_t, libc::c_uint) -> libc::ssize_t;
 
-static GETRANDOM: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());
-
 #[cold]
 #[inline(never)]
-fn init() -> *mut c_void {
-    static NAME: &[u8] = b"getrandom\0";
-    let name_ptr = NAME.as_ptr().cast::<libc::c_char>();
-    let mut ptr = unsafe { libc::dlsym(libc::RTLD_DEFAULT, name_ptr) };
-    if ptr.is_null() || cfg!(getrandom_test_netbsd_fallback) {
-        // Verify `polyfill_using_kern_arand` has the right signature.
-        const POLYFILL: GetRandomFn = polyfill_using_kern_arand;
-        ptr = POLYFILL as *mut c_void;
+fn init() -> NonNull<c_void> {
+    let ptr = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"getrandom".as_ptr()) };
+    if !cfg!(getrandom_test_netbsd_fallback) {
+        if let Some(ptr) = NonNull::new(ptr) {
+            return ptr;
+        }
     }
-    GETRANDOM.store(ptr, Ordering::Release);
-    ptr
+    const POLYFILL: GetRandomFn = polyfill_using_kern_arand;
+    unsafe { NonNull::new_unchecked(POLYFILL as *mut c_void) }
 }
 
 #[inline]
 pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
-    // Despite being only a single atomic variable, we still cannot always use
-    // Ordering::Relaxed, as we need to make sure a successful call to `init`
-    // is "ordered before" any data read through the returned pointer (which
-    // occurs when the function is called). Our implementation mirrors that of
-    // the one in libstd, meaning that the use of non-Relaxed operations is
-    // probably unnecessary.
-    let mut fptr = GETRANDOM.load(Ordering::Acquire);
-    if fptr.is_null() {
-        fptr = init();
-    }
-    let fptr = unsafe { mem::transmute::<*mut c_void, GetRandomFn>(fptr) };
+    #[path = "../utils/lazy_ptr.rs"]
+    mod lazy;
+
+    static GETRANDOM_FN: lazy::LazyPtr<c_void> = lazy::LazyPtr::new();
+
+    let fptr = GETRANDOM_FN.unsync_init(init);
+    let fptr = unsafe { mem::transmute::<*mut c_void, GetRandomFn>(fptr.as_ptr()) };
     utils::sys_fill_exact(dest, |buf| unsafe {
         fptr(buf.as_mut_ptr().cast::<c_void>(), buf.len(), 0)
     })

src/backends/rdrand.rs

@@ -2,9 +2,6 @@
 use crate::{Error, util::slice_as_uninit};
 use core::mem::{MaybeUninit, size_of};
 
-#[path = "../utils/lazy.rs"]
-mod lazy;
-
 #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
 compile_error!("`rdrand` backend can be enabled only for x86 and x86-64 targets!");
 
@@ -20,8 +17,6 @@ cfg_if! {
     }
 }
 
-static RDRAND_GOOD: lazy::LazyBool = lazy::LazyBool::new();
-
 // Recommendation from "Intel® Digital Random Number Generator (DRNG) Software
 // Implementation Guide" - Section 5.2.1 and "Intel® 64 and IA-32 Architectures
 // Software Developer’s Manual" - Volume 1 - Section 7.3.17.1.
@@ -72,7 +67,9 @@ fn self_test() -> bool {
     fails <= 2
 }
 
-fn is_rdrand_good() -> bool {
+#[cold]
+#[inline(never)]
+fn init() -> bool {
     #[cfg(not(target_feature = "rdrand"))]
     {
         // SAFETY: All Rust x86 targets are new enough to have CPUID, and we
@@ -115,6 +112,15 @@ fn is_rdrand_good() -> bool {
     unsafe { self_test() }
 }
 
+fn is_rdrand_good() -> bool {
+    #[path = "../utils/lazy_bool.rs"]
+    mod lazy;
+
+    static RDRAND_GOOD: lazy::LazyBool = lazy::LazyBool::new();
+
+    RDRAND_GOOD.unsync_init(init)
+}
+
 #[target_feature(enable = "rdrand")]
 fn rdrand_exact(dest: &mut [MaybeUninit<u8>]) -> Option<()> {
     // We use chunks_exact_mut instead of chunks_mut as it allows almost all
@@ -162,7 +168,7 @@ fn rdrand_u64() -> Option<u64> {
 
 #[inline]
 pub fn inner_u32() -> Result<u32, Error> {
-    if !RDRAND_GOOD.unsync_init(is_rdrand_good) {
+    if !is_rdrand_good() {
         return Err(Error::NO_RDRAND);
     }
     // SAFETY: After this point, we know rdrand is supported.
@@ -171,7 +177,7 @@ pub fn inner_u32() -> Result<u32, Error> {
 
 #[inline]
 pub fn inner_u64() -> Result<u64, Error> {
-    if !RDRAND_GOOD.unsync_init(is_rdrand_good) {
+    if !is_rdrand_good() {
         return Err(Error::NO_RDRAND);
     }
     // SAFETY: After this point, we know rdrand is supported.
@@ -180,7 +186,7 @@ pub fn inner_u64() -> Result<u64, Error> {
 
 #[inline]
 pub fn fill_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
-    if !RDRAND_GOOD.unsync_init(is_rdrand_good) {
+    if !is_rdrand_good() {
         return Err(Error::NO_RDRAND);
     }
     // SAFETY: After this point, we know rdrand is supported.

src/backends/rndr.rs

@@ -69,8 +69,9 @@ fn is_rndr_available() -> bool {
 
 #[cfg(not(target_feature = "rand"))]
 fn is_rndr_available() -> bool {
-    #[path = "../utils/lazy.rs"]
+    #[path = "../utils/lazy_bool.rs"]
     mod lazy;
+
     static RNDR_GOOD: lazy::LazyBool = lazy::LazyBool::new();
 
     cfg_if::cfg_if! {

src/utils/lazy_bool.rs

@@ -0,0 +1,37 @@
+//! Lazily caches a `bool` in an `AtomicU8`.
+//!
+//! Initialization is intentionally unsynchronized: concurrent callers may race
+//! and run `init` more than once. Once a value is produced, it is cached and
+//! reused by subsequent calls.
+//!
+//! Uses `Ordering::Relaxed` because this helper only publishes the cached
+//! value itself.
+
+use core::sync::atomic::{AtomicU8, Ordering};
+
+pub(crate) struct LazyBool(AtomicU8);
+
+impl LazyBool {
+    const UNINIT: u8 = 2;
+
+    pub const fn new() -> Self {
+        Self(AtomicU8::new(Self::UNINIT))
+    }
+
+    #[cold]
+    fn do_init(&self, init: impl FnOnce() -> bool) -> bool {
+        let val = u8::from(init());
+        self.0.store(val, Ordering::Relaxed);
+        val != 0
+    }
+
+    #[inline]
+    pub fn unsync_init(&self, init: impl FnOnce() -> bool) -> bool {
+        let val = self.0.load(Ordering::Relaxed);
+        if val != Self::UNINIT {
+            val != 0
+        } else {
+            self.do_init(init)
+        }
+    }
+}