From 54456b7ea678434a79fcb8811a9c520f25da3a27 Mon Sep 17 00:00:00 2001 From: James Liu Date: Sat, 20 Apr 2024 19:49:42 -0700 Subject: [PATCH] Make SystemParam::new_archetype and QueryState::new_archetype unsafe functions (#13044) # Objective Fix #2128. Both `Query::new_archetype` and `SystemParam::new_archetype` do not check if the `Archetype` comes from the same World the state is initialized from. This could result in unsoundness via invalid accesses if called incorrectly. ## Solution Make them `unsafe` functions and lift the invariant to the caller. This also caught one instance of us not validating the World in `SystemState::update_archetypes_unsafe_world_cell`'s implementation. --- ## Changelog Changed: `QueryState::new_archetype` is now an unsafe function. Changed: `SystemParam::new_archetype` is now an unsafe function. ## Migration Guide `QueryState::new_archetype` and `SystemParam::new_archetype` are now an unsafe functions that must be sure that the provided `Archetype` is from the same `World` that the state was initialized from. Callers may need to add additional assertions or propagate the safety invariant upwards through the callstack to ensure safety. --- crates/bevy_ecs/macros/src/lib.rs | 10 +++--- crates/bevy_ecs/src/query/state.rs | 36 ++++++++++++++----- crates/bevy_ecs/src/system/function_system.rs | 8 +++-- crates/bevy_ecs/src/system/system_param.rs | 35 ++++++++++++------ crates/bevy_gizmos/src/gizmos.rs | 9 +++-- 5 files changed, 72 insertions(+), 26 deletions(-) diff --git a/crates/bevy_ecs/macros/src/lib.rs b/crates/bevy_ecs/macros/src/lib.rs index 267969a55e..75a43e8e48 100644 --- a/crates/bevy_ecs/macros/src/lib.rs +++ b/crates/bevy_ecs/macros/src/lib.rs @@ -239,8 +239,9 @@ pub fn impl_param_set(_input: TokenStream) -> TokenStream { (#(#param,)*) } - fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) { - <(#(#param,)*) as SystemParam>::new_archetype(state, archetype, system_meta); + unsafe fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) { + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`. + unsafe { <(#(#param,)*) as SystemParam>::new_archetype(state, archetype, system_meta); } } fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) { @@ -425,8 +426,9 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream { } } - fn new_archetype(state: &mut Self::State, archetype: &#path::archetype::Archetype, system_meta: &mut #path::system::SystemMeta) { - <#fields_alias::<'_, '_, #punctuated_generic_idents> as #path::system::SystemParam>::new_archetype(&mut state.state, archetype, system_meta) + unsafe fn new_archetype(state: &mut Self::State, archetype: &#path::archetype::Archetype, system_meta: &mut #path::system::SystemMeta) { + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`. + unsafe { <#fields_alias::<'_, '_, #punctuated_generic_idents> as #path::system::SystemParam>::new_archetype(&mut state.state, archetype, system_meta) } } fn apply(state: &mut Self::State, system_meta: &#path::system::SystemMeta, world: &mut #path::world::World) { diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 47618fb936..9cf1abaab0 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -161,8 +161,12 @@ impl QueryState { ) -> Self { let mut state = Self::new_uninitialized(world); for archetype in world.archetypes.iter() { - if state.new_archetype_internal(archetype) { - state.update_archetype_component_access(archetype, access); + // SAFETY: The state was just initialized from the `world` above, and the archetypes being added + // come directly from the same world. + unsafe { + if state.new_archetype_internal(archetype) { + state.update_archetype_component_access(archetype, access); + } } } state.archetype_generation = world.archetypes.generation(); @@ -342,7 +346,11 @@ impl QueryState { std::mem::replace(&mut self.archetype_generation, archetypes.generation()); for archetype in &archetypes[old_generation..] { - self.new_archetype_internal(archetype); + // SAFETY: The validate_world call ensures that the world is the same the QueryState + // was initialized from. + unsafe { + self.new_archetype_internal(archetype); + } } } @@ -371,13 +379,19 @@ impl QueryState { /// (if applicable, i.e. if the archetype has any intersecting [`ComponentId`] with the current [`QueryState`]). /// /// The passed in `access` will be updated with any new accesses introduced by the new archetype. - pub fn new_archetype( + /// + /// # Safety + /// `archetype` must be from the `World` this state was initialized from. + pub unsafe fn new_archetype( &mut self, archetype: &Archetype, access: &mut Access, ) { - if self.new_archetype_internal(archetype) { - self.update_archetype_component_access(archetype, access); + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from. + let matches = unsafe { self.new_archetype_internal(archetype) }; + if matches { + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from. + unsafe { self.update_archetype_component_access(archetype, access) }; } } @@ -386,7 +400,10 @@ impl QueryState { /// /// Returns `true` if the given `archetype` matches the query. Otherwise, returns `false`. /// If there is no match, then there is no need to update the query's [`FilteredAccess`]. - fn new_archetype_internal(&mut self, archetype: &Archetype) -> bool { + /// + /// # Safety + /// `archetype` must be from the `World` this state was initialized from. + unsafe fn new_archetype_internal(&mut self, archetype: &Archetype) -> bool { if D::matches_component_set(&self.fetch_state, &|id| archetype.contains(id)) && F::matches_component_set(&self.filter_state, &|id| archetype.contains(id)) && self.matches_component_set(&|id| archetype.contains(id)) @@ -431,7 +448,10 @@ impl QueryState { /// For the given `archetype`, adds any component accessed used by this query's underlying [`FilteredAccess`] to `access`. /// /// The passed in `access` will be updated with any new accesses introduced by the new archetype. - pub fn update_archetype_component_access( + /// + /// # Safety + /// `archetype` must be from the `World` this state was initialized from. + pub unsafe fn update_archetype_component_access( &mut self, archetype: &Archetype, access: &mut Access, diff --git a/crates/bevy_ecs/src/system/function_system.rs b/crates/bevy_ecs/src/system/function_system.rs index 1ad187904a..ac96075017 100644 --- a/crates/bevy_ecs/src/system/function_system.rs +++ b/crates/bevy_ecs/src/system/function_system.rs @@ -284,12 +284,15 @@ impl SystemState { /// This method only accesses world metadata. #[inline] pub fn update_archetypes_unsafe_world_cell(&mut self, world: UnsafeWorldCell) { + assert_eq!(self.world_id, world.id(), "Encountered a mismatched World. A System cannot be used with Worlds other than the one it was initialized with."); + let archetypes = world.archetypes(); let old_generation = std::mem::replace(&mut self.archetype_generation, archetypes.generation()); for archetype in &archetypes[old_generation..] { - Param::new_archetype(&mut self.param_state, archetype, &mut self.meta); + // SAFETY: The assertion above ensures that the param_state was initialized from `world`. + unsafe { Param::new_archetype(&mut self.param_state, archetype, &mut self.meta) }; } } @@ -527,7 +530,8 @@ where for archetype in &archetypes[old_generation..] { let param_state = self.param_state.as_mut().unwrap(); - F::Param::new_archetype(param_state, archetype, &mut self.system_meta); + // SAFETY: The assertion above ensures that the param_state was initialized from `world`. + unsafe { F::Param::new_archetype(param_state, archetype, &mut self.system_meta) }; } } diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs index 5d63b310a5..d85f50ae25 100644 --- a/crates/bevy_ecs/src/system/system_param.rs +++ b/crates/bevy_ecs/src/system/system_param.rs @@ -137,12 +137,16 @@ pub unsafe trait SystemParam: Sized { /// and creates a new instance of this param's [`State`](Self::State). fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State; - /// For the specified [`Archetype`], registers the components accessed by this [`SystemParam`] (if applicable). + /// For the specified [`Archetype`], registers the components accessed by this [`SystemParam`] (if applicable).a + /// + /// # Safety + /// `archetype` must be from the [`World`] used to initialize `state` in `init_state`. #[inline] - fn new_archetype( - _state: &mut Self::State, - _archetype: &Archetype, - _system_meta: &mut SystemMeta, + #[allow(unused_variables)] + unsafe fn new_archetype( + state: &mut Self::State, + archetype: &Archetype, + system_meta: &mut SystemMeta, ) { } @@ -208,7 +212,11 @@ unsafe impl SystemParam for Qu state } - fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) { + unsafe fn new_archetype( + state: &mut Self::State, + archetype: &Archetype, + system_meta: &mut SystemMeta, + ) { state.new_archetype(archetype, &mut system_meta.archetype_component_access); } @@ -1343,8 +1351,10 @@ macro_rules! impl_system_param_tuple { } #[inline] - fn new_archetype(($($param,)*): &mut Self::State, _archetype: &Archetype, _system_meta: &mut SystemMeta) { - $($param::new_archetype($param, _archetype, _system_meta);)* + #[allow(unused_unsafe)] + unsafe fn new_archetype(($($param,)*): &mut Self::State, _archetype: &Archetype, _system_meta: &mut SystemMeta) { + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`. + unsafe { $($param::new_archetype($param, _archetype, _system_meta);)* } } #[inline] @@ -1487,8 +1497,13 @@ unsafe impl SystemParam for StaticSystemParam<'_, '_, P::init_state(world, system_meta) } - fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) { - P::new_archetype(state, archetype, system_meta); + unsafe fn new_archetype( + state: &mut Self::State, + archetype: &Archetype, + system_meta: &mut SystemMeta, + ) { + // SAFETY: The caller guarantees that the provided `archetype` matches the World used to initialize `state`. + unsafe { P::new_archetype(state, archetype, system_meta) }; } fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) { diff --git a/crates/bevy_gizmos/src/gizmos.rs b/crates/bevy_gizmos/src/gizmos.rs index f161792bca..bf0cc8a9c2 100644 --- a/crates/bevy_gizmos/src/gizmos.rs +++ b/crates/bevy_gizmos/src/gizmos.rs @@ -55,21 +55,26 @@ pub struct GizmosFetchState { unsafe impl SystemParam for Gizmos<'_, '_, T> { type State = GizmosFetchState; type Item<'w, 's> = Gizmos<'w, 's, T>; + fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State { GizmosFetchState { state: GizmosState::::init_state(world, system_meta), } } - fn new_archetype( + + unsafe fn new_archetype( state: &mut Self::State, archetype: &bevy_ecs::archetype::Archetype, system_meta: &mut SystemMeta, ) { - GizmosState::::new_archetype(&mut state.state, archetype, system_meta); + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`. + unsafe { GizmosState::::new_archetype(&mut state.state, archetype, system_meta) }; } + fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) { GizmosState::::apply(&mut state.state, system_meta, world); } + unsafe fn get_param<'w, 's>( state: &'s mut Self::State, system_meta: &SystemMeta,