fix(ecs): prevent archetype creation from causing oob memory accesses

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>,
-    archetype_lookup: HashMap<ArchetypeComponentsHash, Vec<NonNull<Archetype>>>,
+    archetype_lookup: HashMap<ArchetypeComponentsHash, Vec<usize>>,
     pending_archetype_lookup_entries: Vec<Vec<ComponentId>>,
 }
 
+#[derive(Debug)]
+pub struct ArchetypeRefIter<'component_storage>
+{
+    inner: SliceIter<'component_storage, usize>,
+    archetypes: &'component_storage Vec<Archetype>,
+}
+
+impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
+{
+    type Item = &'component_storage Archetype;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        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<ArchetypeRefIter<'_>>
     {
         let ids = component_ids
             .iter()
@@ -35,12 +58,9 @@ impl ComponentStorage
 
         self.archetype_lookup
             .get(&ArchetypeComponentsHash::new(ids))
-            .map(|archetypes|
-                // SAFETY: All NonNull<Archetype>s 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 NonNull<Archetype>s 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<Item>]` 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<Item>(slice: &[NonNull<Item>]) -> &[&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::<Comps>()))
-                .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<EntityComponent>];
+type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> &'a [Vec<EntityComponent>];
 
 pub struct ComponentIter<'world, Comps>
 {
-    entities: Flatten<Map<SliceIter<'world, &'world Archetype>, ComponentIterMapFn>>,
+    entities: Flatten<Map<ArchetypeRefIter<'world>, ComponentIterMapFn>>,
     comps_pd: PhantomData<Comps>,
 }