[Merged by Bors] - Fix unsoundness in EntityMut::world_scope

crates/bevy_ecs/src/world/entity_ref.rs

@@ -564,9 +564,22 @@ impl<'w> EntityMut<'w> {
  /// # assert_eq!(new_r.0, 1);
  /// ```
  pub fn world_scope<U>(&mut self, f: impl FnOnce(&mut World) -> U) -> U {
- let val = f(self.world);
- self.update_location();
- val
+ // Erase the lifetime of self, so we can get around the borrow checker.
+ let this_ptr = self as *mut Self;
+
+ // Update the stored `EntityLocation` for this instance.
+ // This will get called at the end of this scope, even if the closure `f` unwinds.
+ let _cleanup_guard = bevy_utils::OnDrop::new(move || {
+ // SAFETY:
+ // - When this closure gets called at the end of the outer scope,
+ // the reference to `self` will be inactive, which ensures that this will not alias.
+ // - The pointer must otherwise be safe to dereference, since it was obtained
+ // from a mutable reference.
+ let this = unsafe { &mut *this_ptr };
+ this.update_location();
+ });
+
+ f(self.world)
  }
 
  /// Updates the internal entity location to match the current location in the internal
@@ -856,6 +869,8 @@ pub(crate) unsafe fn take_component<'a>(
 
 #[cfg(test)]
 mod tests {
+ use std::panic::AssertUnwindSafe;
+
  use crate as bevy_ecs;
  use crate::component::ComponentId;
  use crate::prelude::*; // for the `#[derive(Component)]`
@@ -947,4 +962,26 @@ mod tests {
  assert!(entity.get_by_id(invalid_component_id).is_none());
  assert!(entity.get_mut_by_id(invalid_component_id).is_none());
  }
+
+ #[test]
+ fn entity_mut_world_scope_panic() {
+ let mut world = World::new();
+
+ let mut entity = world.spawn_empty();
+ let id = entity.id();
+ let res = std::panic::catch_unwind(AssertUnwindSafe(|| {
+ entity.world_scope(|w| {
+ // Change the entity's `EnityLocation`, which invalidates the original `EntityMut`.
+ // This will get updated at the end of the scope.
+ w.entity_mut(id).insert(TestComponent(0));
+
+ // Ensure that the entity location still gets updated even in case of a panic.
+ panic!("this should get caught by the outer scope")
+ });
+ }));
+ assert!(res.is_err());
+
+ // If the `EntityLocation` has not been updated, this will cause UB.
+ entity.insert(TestComponent(0));
+ }
 }