Fix unsoundness in EntityMut::world_scope (#7387)

# Objective

Found while working on #7385.

The struct `EntityMut` has the safety invariant that it's cached `EntityLocation` must always accurately specify where the entity is stored. Thus, any time its location might be invalidated (such as by calling `EntityMut::world_mut` and moving archetypes), the cached location *must* be updated by calling `EntityMut::update_location`.

The method `world_scope` encapsulates this pattern in safe API by requiring world mutations to be done in a closure, after which `update_location` will automatically be called. However, this method has a soundness hole: if a panic occurs within the closure, then `update_location` will never get called. If the panic is caught in an outer scope, then the `EntityMut` will be left with an outdated location, which is undefined behavior.

An example of this can be seen in the unit test `entity_mut_world_scope_panic`, which has been added to this PR as a regression test. Without the other changes in this PR, that test will invoke undefined behavior in safe code.

## Solution

Call `EntityMut::update_location()` from within a `Drop` impl, which ensures that it will get executed even if `EntityMut::world_scope` unwinds.
This commit is contained in:
JoJoJet 2023-01-29 00:10:45 +00:00
parent bfafa781c1
commit 209f6f8e83

View File

@ -564,14 +564,29 @@ impl<'w> EntityMut<'w> {
/// # assert_eq!(new_r.0, 1); /// # assert_eq!(new_r.0, 1);
/// ``` /// ```
pub fn world_scope<U>(&mut self, f: impl FnOnce(&mut World) -> U) -> U { pub fn world_scope<U>(&mut self, f: impl FnOnce(&mut World) -> U) -> U {
let val = f(self.world); struct Guard<'w, 'a> {
self.update_location(); entity_mut: &'a mut EntityMut<'w>,
val }
impl Drop for Guard<'_, '_> {
#[inline]
fn drop(&mut self) {
self.entity_mut.update_location();
}
}
// When `guard` is dropped at the end of this scope,
// it will update the cached `EntityLocation` for this instance.
// This will run even in case the closure `f` unwinds.
let guard = Guard { entity_mut: self };
f(guard.entity_mut.world)
} }
/// Updates the internal entity location to match the current location in the internal /// Updates the internal entity location to match the current location in the internal
/// [`World`]. This is only needed if the user called [`EntityMut::world`], which enables the /// [`World`].
/// location to change. ///
/// This is *only* required when using the unsafe function [`EntityMut::world_mut`],
/// which enables the location to change.
pub fn update_location(&mut self) { pub fn update_location(&mut self) {
self.location = self.world.entities().get(self.entity).unwrap(); self.location = self.world.entities().get(self.entity).unwrap();
} }
@ -856,6 +871,8 @@ pub(crate) unsafe fn take_component<'a>(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::panic::AssertUnwindSafe;
use crate as bevy_ecs; use crate as bevy_ecs;
use crate::component::ComponentId; use crate::component::ComponentId;
use crate::prelude::*; // for the `#[derive(Component)]` use crate::prelude::*; // for the `#[derive(Component)]`
@ -947,4 +964,28 @@ mod tests {
assert!(entity.get_by_id(invalid_component_id).is_none()); assert!(entity.get_by_id(invalid_component_id).is_none());
assert!(entity.get_mut_by_id(invalid_component_id).is_none()); assert!(entity.get_mut_by_id(invalid_component_id).is_none());
} }
// regression test for https://github.com/bevyengine/bevy/pull/7387
#[test]
fn entity_mut_world_scope_panic() {
let mut world = World::new();
let mut entity = world.spawn_empty();
let old_location = entity.location();
let id = entity.id();
let res = std::panic::catch_unwind(AssertUnwindSafe(|| {
entity.world_scope(|w| {
// Change the entity's `EntityLocation`, 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());
// Ensure that the location has been properly updated.
assert!(entity.location() != old_location);
}
} }