Skip to content

Commit

Permalink
Fix unsoundness in EntityMut::world_scope (#7387)
Browse files Browse the repository at this point in the history
# 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.
  • Loading branch information
JoJoJet committed Jan 29, 2023
1 parent bfafa78 commit 209f6f8
Showing 1 changed file with 46 additions and 5 deletions.
51 changes: 46 additions & 5 deletions crates/bevy_ecs/src/world/entity_ref.rs
Original file line number Diff line number Diff line change
Expand Up @@ -564,14 +564,29 @@ 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
struct Guard<'w, 'a> {
entity_mut: &'a mut EntityMut<'w>,
}

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
/// [`World`]. This is only needed if the user called [`EntityMut::world`], which enables the
/// location to change.
/// [`World`].
///
/// This is *only* required when using the unsafe function [`EntityMut::world_mut`],
/// which enables the location to change.
pub fn update_location(&mut self) {
self.location = self.world.entities().get(self.entity).unwrap();
}
Expand Down Expand Up @@ -856,6 +871,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)]`
Expand Down Expand Up @@ -947,4 +964,28 @@ mod tests {
assert!(entity.get_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);
}
}

0 comments on commit 209f6f8

Please sign in to comment.