Skip to content

Commit 209f6f8

Browse files
committed
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.
1 parent bfafa78 commit 209f6f8

File tree

1 file changed

+46
-5
lines changed

1 file changed

+46
-5
lines changed

crates/bevy_ecs/src/world/entity_ref.rs

+46-5
Original file line numberDiff line numberDiff line change
@@ -564,14 +564,29 @@ impl<'w> EntityMut<'w> {
564564
/// # assert_eq!(new_r.0, 1);
565565
/// ```
566566
pub fn world_scope<U>(&mut self, f: impl FnOnce(&mut World) -> U) -> U {
567-
let val = f(self.world);
568-
self.update_location();
569-
val
567+
struct Guard<'w, 'a> {
568+
entity_mut: &'a mut EntityMut<'w>,
569+
}
570+
571+
impl Drop for Guard<'_, '_> {
572+
#[inline]
573+
fn drop(&mut self) {
574+
self.entity_mut.update_location();
575+
}
576+
}
577+
578+
// When `guard` is dropped at the end of this scope,
579+
// it will update the cached `EntityLocation` for this instance.
580+
// This will run even in case the closure `f` unwinds.
581+
let guard = Guard { entity_mut: self };
582+
f(guard.entity_mut.world)
570583
}
571584

572585
/// Updates the internal entity location to match the current location in the internal
573-
/// [`World`]. This is only needed if the user called [`EntityMut::world`], which enables the
574-
/// location to change.
586+
/// [`World`].
587+
///
588+
/// This is *only* required when using the unsafe function [`EntityMut::world_mut`],
589+
/// which enables the location to change.
575590
pub fn update_location(&mut self) {
576591
self.location = self.world.entities().get(self.entity).unwrap();
577592
}
@@ -856,6 +871,8 @@ pub(crate) unsafe fn take_component<'a>(
856871

857872
#[cfg(test)]
858873
mod tests {
874+
use std::panic::AssertUnwindSafe;
875+
859876
use crate as bevy_ecs;
860877
use crate::component::ComponentId;
861878
use crate::prelude::*; // for the `#[derive(Component)]`
@@ -947,4 +964,28 @@ mod tests {
947964
assert!(entity.get_by_id(invalid_component_id).is_none());
948965
assert!(entity.get_mut_by_id(invalid_component_id).is_none());
949966
}
967+
968+
// regression test for https://github.com/bevyengine/bevy/pull/7387
969+
#[test]
970+
fn entity_mut_world_scope_panic() {
971+
let mut world = World::new();
972+
973+
let mut entity = world.spawn_empty();
974+
let old_location = entity.location();
975+
let id = entity.id();
976+
let res = std::panic::catch_unwind(AssertUnwindSafe(|| {
977+
entity.world_scope(|w| {
978+
// Change the entity's `EntityLocation`, which invalidates the original `EntityMut`.
979+
// This will get updated at the end of the scope.
980+
w.entity_mut(id).insert(TestComponent(0));
981+
982+
// Ensure that the entity location still gets updated even in case of a panic.
983+
panic!("this should get caught by the outer scope")
984+
});
985+
}));
986+
assert!(res.is_err());
987+
988+
// Ensure that the location has been properly updated.
989+
assert!(entity.location() != old_location);
990+
}
950991
}

0 commit comments

Comments
 (0)