From 4b326fb4caed0bcad85954083f87846adface8ad Mon Sep 17 00:00:00 2001
From: JoJoJet <21144246+JoJoJet@users.noreply.github.com>
Date: Mon, 16 Jan 2023 15:41:12 +0000
Subject: [PATCH] Improve safety for `BlobVec::replace_unchecked` (#7181)

# Objective

- The function `BlobVec::replace_unchecked` has informal use of safety comments.
- This function does strange things with `OwningPtr` in order to get around the borrow checker.

## Solution

- Put safety comments in front of each unsafe operation. Describe the specific invariants of each operation and how they apply here.
- Added a guard type `OnDrop`, which is used to simplify ownership transfer in case of a panic.

---

## Changelog

+ Added the guard type `bevy_utils::OnDrop`.
+ Added conversions from `Ptr`, `PtrMut`, and `OwningPtr` to `NonNull<u8>`.
---
 crates/bevy_ecs/src/storage/blob_vec.rs | 61 +++++++++++++++++--------
 crates/bevy_ptr/src/lib.rs              |  6 +++
 crates/bevy_utils/src/lib.rs            | 56 +++++++++++++++++++++++
 3 files changed, 103 insertions(+), 20 deletions(-)

diff --git a/crates/bevy_ecs/src/storage/blob_vec.rs b/crates/bevy_ecs/src/storage/blob_vec.rs
index 7b9437f384218..aa7b3718aba08 100644
--- a/crates/bevy_ecs/src/storage/blob_vec.rs
+++ b/crates/bevy_ecs/src/storage/blob_vec.rs
@@ -6,6 +6,7 @@ use std::{
 };
 
 use bevy_ptr::{OwningPtr, Ptr, PtrMut};
+use bevy_utils::OnDrop;
 
 /// A flat, type-erased data storage type
 ///
@@ -153,31 +154,51 @@ impl BlobVec {
     /// [`BlobVec`]'s `item_layout`
     pub unsafe fn replace_unchecked(&mut self, index: usize, value: OwningPtr<'_>) {
         debug_assert!(index < self.len());
-        // If `drop` panics, then when the collection is dropped during stack unwinding, the
-        // collection's `Drop` impl will call `drop` again for the old value (which is still stored
-        // in the collection), so we get a double drop. To prevent that, we set len to 0 until we're
-        // done.
-        let old_len = self.len;
-        let ptr = self.get_unchecked_mut(index).promote().as_ptr();
-        self.len = 0;
-        // Drop the old value, then write back, justifying the promotion
-        // If the drop impl for the old value panics then we run the drop impl for `value` too.
+
+        // Pointer to the value in the vector that will get replaced.
+        // SAFETY: The caller ensures that `index` fits in this vector.
+        let destination = NonNull::from(self.get_unchecked_mut(index));
+        let source = value.as_ptr();
+
         if let Some(drop) = self.drop {
-            struct OnDrop<F: FnMut()>(F);
-            impl<F: FnMut()> Drop for OnDrop<F> {
-                fn drop(&mut self) {
-                    (self.0)();
-                }
-            }
-            let value = value.as_ptr();
-            let on_unwind = OnDrop(|| (drop)(OwningPtr::new(NonNull::new_unchecked(value))));
+            // Temporarily set the length to zero, so that if `drop` panics the caller
+            // will not be left with a `BlobVec` containing a dropped element within
+            // its initialized range.
+            let old_len = self.len;
+            self.len = 0;
 
-            (drop)(OwningPtr::new(NonNull::new_unchecked(ptr)));
+            // Transfer ownership of the old value out of the vector, so it can be dropped.
+            // SAFETY:
+            // - `destination` was obtained from a `PtrMut` in this vector, which ensures it is non-null,
+            //   well-aligned for the underlying type, and has proper provenance.
+            // - The storage location will get overwritten with `value` later, which ensures
+            //   that the element will not get observed or double dropped later.
+            // - If a panic occurs, `self.len` will remain `0`, which ensures a double-drop
+            //   does not occur. Instead, all elements will be forgotten.
+            let old_value = OwningPtr::new(destination);
+
+            // This closure will run in case `drop()` panics,
+            // which ensures that `value` does not get forgotten.
+            let on_unwind = OnDrop::new(|| drop(value));
+
+            drop(old_value);
 
+            // If the above code does not panic, make sure that `value` doesn't get dropped.
             core::mem::forget(on_unwind);
+
+            // Make the vector's contents observable again, since panics are no longer possible.
+            self.len = old_len;
         }
-        std::ptr::copy_nonoverlapping::<u8>(value.as_ptr(), ptr, self.item_layout.size());
-        self.len = old_len;
+
+        // Copy the new value into the vector, overwriting the previous value.
+        // SAFETY:
+        // - `source` and `destination` were obtained from `OwningPtr`s, which ensures they are
+        //   valid for both reads and writes.
+        // - The value behind `source` will only be dropped if the above branch panics,
+        //   so it must still be initialized and it is safe to transfer ownership into the vector.
+        // - `source` and `destination` were obtained from different memory locations,
+        //   both of which we have exclusive access to, so they are guaranteed not to overlap.
+        std::ptr::copy_nonoverlapping::<u8>(source, destination.as_ptr(), self.item_layout.size());
     }
 
     /// Pushes a value to the [`BlobVec`].
diff --git a/crates/bevy_ptr/src/lib.rs b/crates/bevy_ptr/src/lib.rs
index d793d62606438..39d26fec23637 100644
--- a/crates/bevy_ptr/src/lib.rs
+++ b/crates/bevy_ptr/src/lib.rs
@@ -78,6 +78,12 @@ macro_rules! impl_ptr {
             }
         }
 
+        impl<'a, A: IsAligned> From<$ptr<'a, A>> for NonNull<u8> {
+            fn from(ptr: $ptr<'a, A>) -> Self {
+                ptr.0
+            }
+        }
+
         impl<A: IsAligned> $ptr<'_, A> {
             /// Calculates the offset from a pointer.
             /// As the pointer is type-erased, there is no size information available. The provided
diff --git a/crates/bevy_utils/src/lib.rs b/crates/bevy_utils/src/lib.rs
index 2a522c087900f..e54065eb5f6bf 100644
--- a/crates/bevy_utils/src/lib.rs
+++ b/crates/bevy_utils/src/lib.rs
@@ -34,6 +34,7 @@ use std::{
     future::Future,
     hash::{BuildHasher, Hash, Hasher},
     marker::PhantomData,
+    mem::ManuallyDrop,
     ops::Deref,
     pin::Pin,
 };
@@ -233,3 +234,58 @@ impl<K: Hash + Eq + PartialEq + Clone, V> PreHashMapExt<K, V> for PreHashMap<K,
         }
     }
 }
+
+/// A type which calls a function when dropped.
+/// This can be used to ensure that cleanup code is run even in case of a panic.
+///
+/// Note that this only works for panics that [unwind](https://doc.rust-lang.org/nomicon/unwinding.html)
+/// -- any code within `OnDrop` will be skipped if a panic does not unwind.
+/// In most cases, this will just work.
+///
+/// # Examples
+///
+/// ```
+/// # use bevy_utils::OnDrop;
+/// # fn test_panic(do_panic: bool, log: impl FnOnce(&str)) {
+/// // This will print a message when the variable `_catch` gets dropped,
+/// // even if a panic occurs before we reach the end of this scope.
+/// // This is similar to a `try ... catch` block in languages such as C++.
+/// let _catch = OnDrop::new(|| log("Oops, a panic occured and this function didn't complete!"));
+///
+/// // Some code that may panic...
+/// // ...
+/// # if do_panic { panic!() }
+///
+/// // Make sure the message only gets printed if a panic occurs.
+/// // If we remove this line, then the message will be printed regardless of whether a panic occurs
+/// // -- similar to a `try ... finally` block.
+/// std::mem::forget(_catch);
+/// # }
+/// #
+/// # test_panic(false, |_| unreachable!());
+/// # let mut did_log = false;
+/// # std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+/// #   test_panic(true, |_| did_log = true);
+/// # }));
+/// # assert!(did_log);
+/// ```
+pub struct OnDrop<F: FnOnce()> {
+    callback: ManuallyDrop<F>,
+}
+
+impl<F: FnOnce()> OnDrop<F> {
+    /// Returns an object that will invoke the specified callback when dropped.
+    pub fn new(callback: F) -> Self {
+        Self {
+            callback: ManuallyDrop::new(callback),
+        }
+    }
+}
+
+impl<F: FnOnce()> Drop for OnDrop<F> {
+    fn drop(&mut self) {
+        // SAFETY: We may move out of `self`, since this instance can never be observed after it's dropped.
+        let callback = unsafe { ManuallyDrop::take(&mut self.callback) };
+        callback();
+    }
+}