From 0018beeada4336e8c14b520f540520de9dd30b2d Mon Sep 17 00:00:00 2001
From: James Liu <contact@jamessliu.com>
Date: Wed, 11 Jan 2023 23:12:20 +0000
Subject: [PATCH] Ensure Ptr/PtrMut/OwningPtr are aligned when casting in debug
 builds (#7117)

# Objective
Improve safety testing when using `bevy_ptr` types. This is a follow-up to #7113.

## Solution
Add a debug-only assertion that pointers are aligned when casting to a concrete type. This should very quickly catch any unsoundness from unaligned pointers, even without miri. However, this can have a large negative perf impact on debug builds.

---

## Changelog
Added: `Ptr::deref` will now panic in debug builds if the pointer is not aligned.
Added: `PtrMut::deref_mut` will now panic in debug builds if the pointer is not aligned.
Added: `OwningPtr::read` will now panic in debug builds if the pointer is not aligned.
Added: `OwningPtr::drop_as` will now panic in debug builds if the pointer is not aligned.
---
 crates/bevy_ecs/src/storage/blob_vec.rs |  7 ++--
 crates/bevy_ptr/src/lib.rs              | 49 ++++++++++++++++++++++---
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_ecs/src/storage/blob_vec.rs b/crates/bevy_ecs/src/storage/blob_vec.rs
index c0bfa3302572c2..7b9437f384218d 100644
--- a/crates/bevy_ecs/src/storage/blob_vec.rs
+++ b/crates/bevy_ecs/src/storage/blob_vec.rs
@@ -46,10 +46,11 @@ impl BlobVec {
         drop: Option<unsafe fn(OwningPtr<'_>)>,
         capacity: usize,
     ) -> BlobVec {
+        let align = NonZeroUsize::new(item_layout.align()).expect("alignment must be > 0");
+        let data = bevy_ptr::dangling_with_align(align);
         if item_layout.size() == 0 {
-            let align = NonZeroUsize::new(item_layout.align()).expect("alignment must be > 0");
             BlobVec {
-                data: bevy_ptr::dangling_with_align(align),
+                data,
                 capacity: usize::MAX,
                 len: 0,
                 item_layout,
@@ -57,7 +58,7 @@ impl BlobVec {
             }
         } else {
             let mut blob_vec = BlobVec {
-                data: NonNull::dangling(),
+                data,
                 capacity: 0,
                 len: 0,
                 item_layout,
diff --git a/crates/bevy_ptr/src/lib.rs b/crates/bevy_ptr/src/lib.rs
index 57a8f0153cf14c..d793d626064387 100644
--- a/crates/bevy_ptr/src/lib.rs
+++ b/crates/bevy_ptr/src/lib.rs
@@ -164,7 +164,7 @@ impl<'a, A: IsAligned> Ptr<'a, A> {
     ///   for the pointee type `T`.
     #[inline]
     pub unsafe fn deref<T>(self) -> &'a T {
-        &*self.as_ptr().cast()
+        &*self.as_ptr().cast::<T>().debug_ensure_aligned()
     }
 
     /// Gets the underlying pointer, erasing the associated lifetime.
@@ -218,7 +218,7 @@ impl<'a, A: IsAligned> PtrMut<'a, A> {
     ///   for the pointee type `T`.
     #[inline]
     pub unsafe fn deref_mut<T>(self) -> &'a mut T {
-        &mut *self.as_ptr().cast()
+        &mut *self.as_ptr().cast::<T>().debug_ensure_aligned()
     }
 
     /// Gets the underlying pointer, erasing the associated lifetime.
@@ -287,7 +287,7 @@ impl<'a, A: IsAligned> OwningPtr<'a, A> {
     ///   for the pointee type `T`.
     #[inline]
     pub unsafe fn read<T>(self) -> T {
-        self.as_ptr().cast::<T>().read()
+        self.as_ptr().cast::<T>().debug_ensure_aligned().read()
     }
 
     /// Consumes the [`OwningPtr`] to drop the underlying data of type `T`.
@@ -298,7 +298,10 @@ impl<'a, A: IsAligned> OwningPtr<'a, A> {
     ///   for the pointee type `T`.
     #[inline]
     pub unsafe fn drop_as<T>(self) {
-        self.as_ptr().cast::<T>().drop_in_place();
+        self.as_ptr()
+            .cast::<T>()
+            .debug_ensure_aligned()
+            .drop_in_place();
     }
 
     /// Gets the underlying pointer, erasing the associated lifetime.
@@ -364,9 +367,10 @@ impl<'a, T> Copy for ThinSlicePtr<'a, T> {}
 impl<'a, T> From<&'a [T]> for ThinSlicePtr<'a, T> {
     #[inline]
     fn from(slice: &'a [T]) -> Self {
+        let ptr = slice.as_ptr() as *mut T;
         Self {
             // SAFETY: a reference can never be null
-            ptr: unsafe { NonNull::new_unchecked(slice.as_ptr() as *mut T) },
+            ptr: unsafe { NonNull::new_unchecked(ptr.debug_ensure_aligned()) },
             #[cfg(debug_assertions)]
             len: slice.len(),
             _marker: PhantomData,
@@ -377,6 +381,7 @@ impl<'a, T> From<&'a [T]> for ThinSlicePtr<'a, T> {
 /// Creates a dangling pointer with specified alignment.
 /// See [`NonNull::dangling`].
 pub fn dangling_with_align(align: NonZeroUsize) -> NonNull<u8> {
+    debug_assert!(align.is_power_of_two(), "Alignment must be power of two.");
     // SAFETY: The pointer will not be null, since it was created
     // from the address of a `NonZeroUsize`.
     unsafe { NonNull::new_unchecked(align.get() as *mut u8) }
@@ -429,3 +434,37 @@ impl<'a, T> UnsafeCellDeref<'a, T> for &'a UnsafeCell<T> {
         self.get().read()
     }
 }
+
+trait DebugEnsureAligned {
+    fn debug_ensure_aligned(self) -> Self;
+}
+
+// Disable this for miri runs as it already checks if pointer to reference
+// casts are properly aligned.
+#[cfg(all(debug_assertions, not(miri)))]
+impl<T: Sized> DebugEnsureAligned for *mut T {
+    #[track_caller]
+    fn debug_ensure_aligned(self) -> Self {
+        let align = core::mem::align_of::<T>();
+        // Implemenation shamelessly borrowed from the currently unstable
+        // ptr.is_aligned_to.
+        //
+        // Replace once https://github.com/rust-lang/rust/issues/96284 is stable.
+        assert!(
+            self as usize & (align - 1) == 0,
+            "pointer is not aligned. Address {:p} does not have alignemnt {} for type {}",
+            self,
+            align,
+            core::any::type_name::<T>(),
+        );
+        self
+    }
+}
+
+#[cfg(any(not(debug_assertions), miri))]
+impl<T: Sized> DebugEnsureAligned for *mut T {
+    #[inline(always)]
+    fn debug_ensure_aligned(self) -> Self {
+        self
+    }
+}