From e25f4f2af6e5e522f8b6ad0ca0b84f607a73cae2 Mon Sep 17 00:00:00 2001 From: HampusM Date: Sun, 16 Jun 2024 17:59:49 +0200 Subject: fix(ecs): prevent archetype creation from causing oob memory accesses --- ecs/src/component/storage.rs | 98 ++++++++++++++++++++++++-------------------- ecs/src/query.rs | 8 ++-- 2 files changed, 57 insertions(+), 49 deletions(-) diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs index bd53da0..f174ea9 100644 --- a/ecs/src/component/storage.rs +++ b/ecs/src/component/storage.rs @@ -1,7 +1,7 @@ use std::any::type_name; use std::collections::{HashMap, HashSet}; use std::hash::{DefaultHasher, Hash, Hasher}; -use std::ptr::NonNull; +use std::slice::Iter as SliceIter; use crate::component::{Component, Id as ComponentId, IsOptional as ComponentIsOptional}; use crate::lock::Lock; @@ -12,16 +12,39 @@ use crate::EntityComponent; pub struct ComponentStorage { archetypes: Vec, - archetype_lookup: HashMap>>, + archetype_lookup: HashMap>, pending_archetype_lookup_entries: Vec>, } +#[derive(Debug)] +pub struct ArchetypeRefIter<'component_storage> +{ + inner: SliceIter<'component_storage, usize>, + archetypes: &'component_storage Vec, +} + +impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage> +{ + type Item = &'component_storage Archetype; + + fn next(&mut self) -> Option + { + let archetype_index = *self.inner.next()?; + + Some( + self.archetypes + .get(archetype_index) + .expect("Archetype index in archetype lookup entry was not found"), + ) + } +} + impl ComponentStorage { pub fn find_entities( &self, component_ids: &[(ComponentId, ComponentIsOptional)], - ) -> Option<&[&Archetype]> + ) -> Option> { let ids = component_ids .iter() @@ -35,12 +58,9 @@ impl ComponentStorage self.archetype_lookup .get(&ArchetypeComponentsHash::new(ids)) - .map(|archetypes| - // SAFETY: All NonNulls are references to items of the - // archetypes field and the items won't be dropped until the whole - // struct is dropped - unsafe { - nonnull_slice_to_ref_slice(archetypes.as_slice()) + .map(|archetypes_indices| ArchetypeRefIter { + inner: archetypes_indices.iter(), + archetypes: &self.archetypes, }) } @@ -57,7 +77,7 @@ impl ComponentStorage .join(", ") ); - let archetypes = self + let archetypes_indices = self .archetype_lookup .entry(ArchetypeComponentsHash::new( components @@ -68,18 +88,18 @@ impl ComponentStorage .or_insert_with(|| { self.archetypes.push(Archetype::default()); - vec![NonNull::from(self.archetypes.last().unwrap())] + vec![self.archetypes.len() - 1] }); - // SAFETY: All NonNulls are references to items of the - // archetypes field and the items won't be dropped until the whole - // struct is dropped - let archetype = unsafe { - archetypes - .first_mut() - .expect("Archetype has disappeared") - .as_mut() - }; + let archetype = self + .archetypes + .get_mut( + archetypes_indices + .first() + .copied() + .expect("No archetype index found"), + ) + .expect("Archetype is gone"); archetype .component_ids @@ -111,22 +131,26 @@ impl ComponentStorage .map(|component_id| *component_id) .collect(); - let matching_archetypes = self.archetypes.iter().filter_map(|archetype| { - if archetype.component_ids.is_superset(&components_set) { - return Some(NonNull::from(archetype)); - } + let matching_archetype_indices = self + .archetypes + .iter() + .enumerate() + .filter_map(|(index, archetype)| { + if archetype.component_ids.is_superset(&components_set) { + return Some(index); + } - None - }); + None + }); - let lookup_archetypes = self + let archetype_indices = self .archetype_lookup .entry(ArchetypeComponentsHash::new( pending_entry.into_iter().copied().clone(), )) .or_default(); - lookup_archetypes.extend(matching_archetypes); + archetype_indices.extend(matching_archetype_indices); } } } @@ -168,21 +192,10 @@ impl ArchetypeComponentsHash } } -/// Casts a `&[NonNull]` to a `&[&Item]`. -/// -/// # Safety -/// All items in the slice must be initialized, properly aligned and follow Rust's -/// aliasing rules. -const unsafe fn nonnull_slice_to_ref_slice(slice: &[NonNull]) -> &[&Item] -{ - unsafe { &*(std::ptr::from_ref(slice) as *const [&Item]) } -} - #[cfg(test)] mod tests { use std::collections::HashSet; - use std::ptr::addr_of; use ecs_macros::Component; @@ -257,14 +270,11 @@ mod tests ])) .expect("Expected entry in archetype lookup map"); - let archetype_from_lookup = lookup + let first_archetype_index = lookup .first() .expect("Expected archetype lookup to contain a archetype reference"); - assert_eq!( - archetype_from_lookup.as_ptr() as usize, - addr_of!(*archetype) as usize - ); + assert_eq!(*first_archetype_index, 0); } #[test] diff --git a/ecs/src/query.rs b/ecs/src/query.rs index a2edc4d..dcd0b0e 100644 --- a/ecs/src/query.rs +++ b/ecs/src/query.rs @@ -2,10 +2,9 @@ use std::any::{type_name, Any}; use std::collections::HashSet; use std::iter::{Flatten, Map}; use std::marker::PhantomData; -use std::slice::Iter as SliceIter; use std::sync::{Arc, Weak}; -use crate::component::storage::Archetype; +use crate::component::storage::{Archetype, ArchetypeRefIter}; use crate::component::{ Id as ComponentId, IsOptional as ComponentIsOptional, @@ -46,7 +45,6 @@ where .component_storage .find_entities(&Comps::ids()) .unwrap_or_else(|| panic!("Could not find {:?}", type_name::())) - .iter() .map((|archetype| archetype.components.as_slice()) as ComponentIterMapFn) .flatten(), comps_pd: PhantomData, @@ -218,11 +216,11 @@ where } } -type ComponentIterMapFn = for<'a> fn(&'a &'a Archetype) -> &'a [Vec]; +type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> &'a [Vec]; pub struct ComponentIter<'world, Comps> { - entities: Flatten, ComponentIterMapFn>>, + entities: Flatten, ComponentIterMapFn>>, comps_pd: PhantomData, } -- cgit v1.2.3-18-g5258