From 7e4eacd5d132fac7a96bf174809d8402ecd9849f Mon Sep 17 00:00:00 2001 From: Ellen Date: Fri, 4 Mar 2022 04:08:11 +0000 Subject: [PATCH] >:( --- crates/bevy_ecs/src/reflect.rs | 6 ++- crates/bevy_ecs/src/system/query.rs | 9 ++-- crates/bevy_ecs/src/world/entity_ref.rs | 70 +++++++++++++------------ crates/bevy_ecs/src/world/mod.rs | 16 +++--- 4 files changed, 56 insertions(+), 45 deletions(-) diff --git a/crates/bevy_ecs/src/reflect.rs b/crates/bevy_ecs/src/reflect.rs index 4881dd83f0edde..5fbbede5b306de 100644 --- a/crates/bevy_ecs/src/reflect.rs +++ b/crates/bevy_ecs/src/reflect.rs @@ -103,15 +103,17 @@ impl FromType for ReflectComponent { .entity_mut(destination_entity) .insert(destination_component); }, - reflect_component: |world, entity| { + reflect_component: |world, entity| unsafe { world .get_entity(entity)? - .get::() + // SAFE: lifetimes force correct usage of returned data + .get_unchecked::() .map(|c| c as &dyn Reflect) }, reflect_component_mut: |world, entity| unsafe { world .get_entity(entity)? + // SAFE: lifetimes force correct usage of returned data .get_unchecked_mut::(world.last_change_tick(), world.read_change_tick()) .map(|c| ReflectMut { value: c.value as &mut dyn Reflect, diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 8d65633de149da..7c1878f1e692f8 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -728,9 +728,12 @@ where .archetype_component_access .has_read(archetype_component) { - entity_ref - .get::() - .ok_or(QueryComponentError::MissingComponent) + // SAFE: lifetimes enforce correct usage of returned borrow + unsafe { + entity_ref + .get_unchecked::() + .ok_or(QueryComponentError::MissingComponent) + } } else { Err(QueryComponentError::MissingReadAccess) } diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index a9e9cd56b6edfa..67097d3353ccd6 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -62,17 +62,26 @@ impl<'w> EntityRef<'w> { } #[inline] - pub fn get(&self) -> Option<&'w T> { + pub fn get(&self) -> Option<&'_ T> { + // SAFE: lifetimes enforce correct usage of returned borrow + unsafe { self.get_unchecked::() } + } + + /// This allows aliased mutability use after free bugs. + /// # Safety + /// - Must not be a mutable reference to the component. + /// - Must not use the returned reference after the component is removed + #[inline] + pub unsafe fn get_unchecked(&self) -> Option<&'w T> { // SAFE: entity location is valid and returned component is of type T - unsafe { - get_component_with_type(self.world, TypeId::of::(), self.entity, self.location) - .map(|value| &*value.cast::()) - } + get_component_with_type(self.world, TypeId::of::(), self.entity, self.location) + .map(|value| &*value.cast::()) } + /// This allows aliased mutability and use after free bugs. /// # Safety - /// This allows aliased mutability. You must make sure this call does not result in multiple - /// mutable references to the same component + /// - Must not be any other references to the component. + /// - Must not use the returned reference after the component is removed. #[inline] pub unsafe fn get_unchecked_mut( &self, @@ -145,39 +154,32 @@ impl<'w> EntityMut<'w> { } #[inline] - pub fn get(&self) -> Option<&'w T> { - // SAFE: entity location is valid and returned component is of type T - unsafe { - get_component_with_type(self.world, TypeId::of::(), self.entity, self.location) - .map(|value| &*value.cast::()) - } + pub fn get(&self) -> Option<&'_ T> { + // SAFE: lifetimes enforce correct usage of returned borrow + unsafe { self.get_unchecked::() } } #[inline] - pub fn get_mut(&mut self) -> Option> { - // SAFE: world access is unique, entity location is valid, and returned component is of type - // T - unsafe { - get_component_and_ticks_with_type( - self.world, - TypeId::of::(), - self.entity, - self.location, - ) - .map(|(value, ticks)| Mut { - value: &mut *value.cast::(), - ticks: Ticks { - component_ticks: &mut *ticks, - last_change_tick: self.world.last_change_tick(), - change_tick: self.world.change_tick(), - }, - }) - } + pub fn get_mut(&mut self) -> Option> { + // SAFE: world access is unique, and lifetimes enforce correct usage of returned borrow + unsafe { self.get_unchecked_mut::() } + } + + /// This allows aliased mutability use after free bugs. + /// # Safety + /// - Must not be a mutable reference to the component. + /// - Must not use the returned reference after the component is removed + #[inline] + pub unsafe fn get_unchecked(&self) -> Option<&'w T> { + // SAFE: entity location is valid and returned component is of type T + get_component_with_type(self.world, TypeId::of::(), self.entity, self.location) + .map(|value| &*value.cast::()) } + /// This allows aliased mutability and use after free bugs. /// # Safety - /// This allows aliased mutability. You must make sure this call does not result in multiple - /// mutable references to the same component + /// - Must not be any other references to the component. + /// - Must not use the returned reference after the component is removed. #[inline] pub unsafe fn get_unchecked_mut(&self) -> Option> { get_component_and_ticks_with_type(self.world, TypeId::of::(), self.entity, self.location) diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 22add5b970365a..b35d024c65fb1e 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -199,7 +199,8 @@ impl World { /// .insert(Position { x: 0.0, y: 0.0 }) /// .id(); /// - /// let position = world.entity(entity).get::().unwrap(); + /// let entity_ref = world.entity(entity); + /// let position = entity_ref.get::().unwrap(); /// assert_eq!(position.x, 0.0); /// ``` #[inline] @@ -226,8 +227,8 @@ impl World { /// let entity = world.spawn() /// .insert(Position { x: 0.0, y: 0.0 }) /// .id(); - /// - /// let mut position = world.entity_mut(entity).get_mut::().unwrap(); + /// let mut entity_mut = world.entity_mut(entity); + /// let mut position = entity_mut.get_mut::().unwrap(); /// position.x = 1.0; /// ``` #[inline] @@ -339,7 +340,8 @@ impl World { /// .insert_bundle((Num(1), Label("hello"))) // add a bundle of components /// .id(); /// - /// let position = world.entity(entity).get::().unwrap(); + /// let entity_ref = world.entity(entity); + /// let position = entity_ref.get::().unwrap(); /// assert_eq!(position.x, 0.0); /// ``` pub fn spawn(&mut self) -> EntityMut { @@ -416,7 +418,8 @@ impl World { /// ``` #[inline] pub fn get(&self, entity: Entity) -> Option<&T> { - self.get_entity(entity)?.get() + // SAFE: lifetimes enforce correct usage of returned borrow + unsafe { self.get_entity(entity)?.get_unchecked::() } } /// Retrieves a mutable reference to the given `entity`'s [Component] of the given type. @@ -439,7 +442,8 @@ impl World { /// ``` #[inline] pub fn get_mut(&mut self, entity: Entity) -> Option> { - self.get_entity_mut(entity)?.get_mut() + // SAFE: lifetimes enforce correct usage of returned borrow + unsafe { self.get_entity_mut(entity)?.get_unchecked_mut::() } } /// Despawns the given `entity`, if it exists. This will also remove all of the entity's