zeroize: replace atomic_fence with optimization_barrier (#1252)

Author: newpavlov

zeroize/src/aarch64.rs

@@ -1,6 +1,6 @@
 //! [`Zeroize`] impls for ARM64 SIMD registers.
 
-use crate::{Zeroize, atomic_fence, volatile_write};
+use crate::{Zeroize, optimization_barrier, volatile_write};
 
 use core::arch::aarch64::*;
 
@@ -11,7 +11,7 @@ macro_rules! impl_zeroize_for_simd_register {
                 #[inline]
                 fn zeroize(&mut self) {
                     volatile_write(self, unsafe { core::mem::zeroed() });
-                    atomic_fence();
+                    optimization_barrier(self);
                 }
             }
         )+

zeroize/src/barrier.rs

@@ -0,0 +1,91 @@
+/// Observe the referenced data and prevent the compiler from removing previous writes to it.
+///
+/// This function acts like [`core::hint::black_box`] but takes a reference and
+/// does not return the passed value.
+///
+/// It's implemented using the [`core::arch::asm!`] macro on target arches where `asm!` is stable,
+/// i.e. `aarch64`, `arm`, `arm64ec`, `loongarch64`, `riscv32`, `riscv64`, `s390x`, `x86`, and
+/// `x86_64`.
+///
+/// On all other targets it's implemented using [`core::hint::black_box`] and custom `black_box`
+/// implemented using `#[inline(never)]` and `read_volatile`.
+///
+/// # Examples
+/// ```ignore
+/// use core::num::NonZeroU32;
+/// use zeroize::{ZeroizeOnDrop, zeroize_flat_type};
+///
+/// struct DataToZeroize {
+///     buf: [u8; 32],
+///     pos: NonZeroU32,
+/// }
+///
+/// struct SomeMoreFlatData(u64);
+///
+/// impl Drop for DataToZeroize {
+///     fn drop(&mut self) {
+///         self.buf = [0u8; 32];
+///         self.pos = NonZeroU32::new(32).unwrap();
+///         zeroize::optimization_barrier(self);
+///     }
+/// }
+///
+/// impl zeroize::ZeroizeOnDrop for DataToZeroize {}
+///
+/// let mut data = DataToZeroize {
+///     buf: [3u8; 32],
+///     pos: NonZeroU32::new(32).unwrap(),
+/// };
+///
+/// // data gets zeroized when dropped
+/// ```
+pub(crate) fn optimization_barrier<T: ?Sized>(val: &T) {
+    #[cfg(all(
+        not(miri),
+        any(
+            target_arch = "aarch64",
+            target_arch = "arm",
+            target_arch = "arm64ec",
+            target_arch = "loongarch64",
+            target_arch = "riscv32",
+            target_arch = "riscv64",
+            target_arch = "s390x",
+            target_arch = "x86",
+            target_arch = "x86_64",
+        )
+    ))]
+    unsafe {
+        core::arch::asm!(
+            "# {}",
+            in(reg) val as *const T as *const (),
+            options(readonly, preserves_flags, nostack),
+        );
+    }
+    #[cfg(not(all(
+        not(miri),
+        any(
+            target_arch = "aarch64",
+            target_arch = "arm",
+            target_arch = "arm64ec",
+            target_arch = "loongarch64",
+            target_arch = "riscv32",
+            target_arch = "riscv64",
+            target_arch = "s390x",
+            target_arch = "x86",
+            target_arch = "x86_64",
+        )
+    )))]
+    {
+        /// Custom version of `core::hint::black_box` implemented using
+        /// `#[inline(never)]` and `read_volatile`.
+        #[inline(never)]
+        fn custom_black_box(p: *const u8) {
+            let _ = unsafe { core::ptr::read_volatile(p) };
+        }
+
+        core::hint::black_box(val);
+        if size_of_val(val) > 0 {
+            custom_black_box(val as *const T as *const u8);
+        }
+    }
+}

zeroize/src/lib.rs

@@ -250,6 +250,9 @@ mod aarch64;
 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
 mod x86;
 
+mod barrier;
+use barrier::optimization_barrier;
+
 use core::{
     marker::{PhantomData, PhantomPinned},
     mem::{MaybeUninit, size_of},
@@ -259,7 +262,6 @@ use core::{
     },
     ops, ptr,
     slice::IterMut,
-    sync::atomic,
 };
 
 #[cfg(feature = "alloc")]
@@ -300,7 +302,7 @@ where
 {
     fn zeroize(&mut self) {
         volatile_write(self, Z::default());
-        atomic_fence();
+        optimization_barrier(self);
     }
 }
 
@@ -334,7 +336,7 @@ macro_rules! impl_zeroize_for_non_zero {
                         None => unreachable!(),
                     };
                     volatile_write(self, ONE);
-                    atomic_fence();
+                    optimization_barrier(self);
                 }
             }
         )+
@@ -425,7 +427,7 @@ where
         // done so by take().
         unsafe { ptr::write_volatile(self, None) }
 
-        atomic_fence();
+        optimization_barrier(self);
     }
 }
 
@@ -441,7 +443,7 @@ impl<Z> Zeroize for MaybeUninit<Z> {
         // Safety:
         // `MaybeUninit` is valid for any byte pattern, including zeros.
         unsafe { ptr::write_volatile(self, MaybeUninit::zeroed()) }
-        atomic_fence();
+        optimization_barrier(self);
     }
 }
 
@@ -467,7 +469,7 @@ impl<Z> Zeroize for [MaybeUninit<Z>] {
         // and 0 is a valid value for `MaybeUninit<Z>`
         // The memory of the slice should not wrap around the address space.
         unsafe { volatile_set(ptr, MaybeUninit::zeroed(), size) }
-        atomic_fence();
+        optimization_barrier(self);
     }
 }
 
@@ -493,7 +495,7 @@ where
         // `self.len()` is also not larger than an `isize`, because of the assertion above.
         // The memory of the slice should not wrap around the address space.
         unsafe { volatile_set(self.as_mut_ptr(), Z::default(), self.len()) };
-        atomic_fence();
+        optimization_barrier(self);
     }
 }
 
@@ -753,14 +755,6 @@ where
     }
 }
 
-/// Use fences to prevent accesses from being reordered before this
-/// point, which should hopefully help ensure that all accessors
-/// see zeroes after this point.
-#[inline(always)]
-fn atomic_fence() {
-    atomic::compiler_fence(atomic::Ordering::SeqCst);
-}
-
 /// Perform a volatile write to the destination
 #[inline(always)]
 fn volatile_write<T: Copy + Sized>(dst: &mut T, src: T) {
@@ -851,7 +845,7 @@ pub unsafe fn zeroize_flat_type<F: Sized>(data: *mut F) {
     unsafe {
         volatile_set(data as *mut u8, 0, size);
     }
-    atomic_fence()
+    optimization_barrier(&data);
 }
 
 /// Internal module used as support for `AssertZeroizeOnDrop`.

zeroize/src/x86.rs

@@ -1,6 +1,6 @@
 //! [`Zeroize`] impls for x86 SIMD registers
 
-use crate::{Zeroize, atomic_fence, volatile_write};
+use crate::{Zeroize, optimization_barrier, volatile_write};
 
 #[cfg(target_arch = "x86")]
 use core::arch::x86::*;
@@ -15,7 +15,7 @@ macro_rules! impl_zeroize_for_simd_register {
                 #[inline]
                 fn zeroize(&mut self) {
                     volatile_write(self, unsafe { core::mem::zeroed() });
-                    atomic_fence();
+                    optimization_barrier(self);
                 }
             }
         )*