From c82a522713f7f85dfacb0da04880a646244dc158 Mon Sep 17 00:00:00 2001
From: Joseph Richey <joerichey@google.com>
Date: Mon, 13 Jun 2022 12:10:19 -0700
Subject: [PATCH] Use AtomicPtr instead of AtomicUsize for Weak (#263)

This allows Strict Provenance to work properly, fixing #262. It also
now matches what `libstd` does:
https://github.com/rust-lang/rust/blob/9f7e997c8bc3cacd2ab4eb75e63cb5fa9279c7b0/library/std/src/sys/unix/weak.rs#L85-L141

Also, while reading the `libstd` code, I noticed that they use an
`Acquire` fence and `Release` store as the returned pointer should
have "consume" semantics. As this doesn't yet exist in Rust, we
instead do exactly what `libstd` does, which means:
  - Relaxed Load
  - Release Store
  - Acquire fence when returning pointer

Signed-off-by: Joe Richey <joerichey@google.com>

Co-authored-by: Joe ST <joe@fbstj.net>
---
 src/util_libc.rs | 57 ++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 12 deletions(-)

diff --git a/src/util_libc.rs b/src/util_libc.rs
index 6df1cd7d..d057071a 100644
--- a/src/util_libc.rs
+++ b/src/util_libc.rs
@@ -6,8 +6,13 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 #![allow(dead_code)]
-use crate::{util::LazyUsize, Error};
-use core::{num::NonZeroU32, ptr::NonNull};
+use crate::Error;
+use core::{
+    num::NonZeroU32,
+    ptr::NonNull,
+    sync::atomic::{fence, AtomicPtr, Ordering},
+};
+use libc::c_void;
 
 cfg_if! {
     if #[cfg(any(target_os = "netbsd", target_os = "openbsd", target_os = "android"))] {
@@ -76,29 +81,57 @@ pub fn sys_fill_exact(
 
 // A "weak" binding to a C function that may or may not be present at runtime.
 // Used for supporting newer OS features while still building on older systems.
-// F must be a function pointer of type `unsafe extern "C" fn`. Based off of the
-// weak! macro in libstd.
+// Based off of the DlsymWeak struct in libstd:
+// https://github.com/rust-lang/rust/blob/1.61.0/library/std/src/sys/unix/weak.rs#L84
+// except that the caller must manually cast self.ptr() to a function pointer.
 pub struct Weak {
     name: &'static str,
-    addr: LazyUsize,
+    addr: AtomicPtr<c_void>,
 }
 
 impl Weak {
+    // A non-null pointer value which indicates we are uninitialized. This
+    // constant should ideally not be a valid address of a function pointer.
+    // However, if by chance libc::dlsym does return UNINIT, there will not
+    // be undefined behavior. libc::dlsym will just be called each time ptr()
+    // is called. This would be inefficient, but correct.
+    // TODO: Replace with core::ptr::invalid_mut(1) when that is stable.
+    const UNINIT: *mut c_void = 1 as *mut c_void;
+
     // Construct a binding to a C function with a given name. This function is
     // unsafe because `name` _must_ be null terminated.
     pub const unsafe fn new(name: &'static str) -> Self {
         Self {
             name,
-            addr: LazyUsize::new(),
+            addr: AtomicPtr::new(Self::UNINIT),
         }
     }
 
-    // Return a function pointer if present at runtime. Otherwise, return null.
-    pub fn ptr(&self) -> Option<NonNull<libc::c_void>> {
-        let addr = self.addr.unsync_init(|| unsafe {
-            libc::dlsym(libc::RTLD_DEFAULT, self.name.as_ptr() as *const _) as usize
-        });
-        NonNull::new(addr as *mut _)
+    // Return the address of a function if present at runtime. Otherwise,
+    // return None. Multiple callers can call ptr() concurrently. It will
+    // always return _some_ value returned by libc::dlsym. However, the
+    // dlsym function may be called multiple times.
+    pub fn ptr(&self) -> Option<NonNull<c_void>> {
+        // Despite having only a single atomic variable (self.addr), we still
+        // cannot always use Ordering::Relaxed, as we need to make sure a
+        // successful call to dlsym() 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.
+        match self.addr.load(Ordering::Relaxed) {
+            Self::UNINIT => {
+                let symbol = self.name.as_ptr() as *const _;
+                let addr = unsafe { libc::dlsym(libc::RTLD_DEFAULT, symbol) };
+                // Synchronizes with the Acquire fence below
+                self.addr.store(addr, Ordering::Release);
+                NonNull::new(addr)
+            }
+            addr => {
+                let func = NonNull::new(addr)?;
+                fence(Ordering::Acquire);
+                Some(func)
+            }
+        }
     }
 }