fix: fix soundess issue in bones_ecs related to UntypedResource acc…

…ess. (#185)

Previously it was possible to change the underlying `SchemaBox` for an
`UntypedResource` through it's cell, which would allow de-syncing the
schema ID of the box, which was relied on for the typed resources store.
This would be unsound.

To fix it we prevent accessing the `SchemaBox` in an `UntypedResource`
directly and only give out atomic wrappers around `SchemaRef` and
`SchemaRefMut`.

This required switching from `atomic_refcell` to `atomicell`, which had
the ability to do unsafe maps on the atomic borrows. Both libraries are
fundamentally similar in idea and implementation, and I like the way
that `atomicell` allows us to create our own borrowed types, so I think
it's a benefit either way.

framework_crates/bones_ecs/Cargo.toml

@@ -23,12 +23,12 @@ keysize32 = []
 bones_utils  = { version = "0.3", path = "../bones_utils" }
 bones_schema = { version = "0.3", path = "../bones_schema" }
 
-anyhow         = "1.0"
-atomic_refcell = "0.1"
-bitset-core    = "0.1"
-thiserror      = "1.0"
-glam           = { version = "0.24", optional = true }
-paste          = { version = "1.0", optional = true }
+anyhow      = "1.0"
+atomicell   = "0.1"
+bitset-core = "0.1"
+thiserror   = "1.0"
+glam        = { version = "0.24", optional = true }
+paste       = { version = "1.0", optional = true }
 
 [dev-dependencies]
 glam = "0.24"

framework_crates/bones_ecs/src/components.rs

@@ -12,15 +12,15 @@ pub use iterator::*;
 pub use typed::*;
 pub use untyped::*;
 
-type AtomicComponentStore<T> = Arc<AtomicRefCell<ComponentStore<T>>>;
+type AtomicComponentStore<T> = Arc<AtomicCell<ComponentStore<T>>>;
 
 /// A collection of [`ComponentStore<T>`].
 ///
 /// [`ComponentStores`] is used to in [`World`] to store all component types that have been
 /// initialized for that world.
 #[derive(Default)]
 pub struct ComponentStores {
-    pub(crate) components: HashMap<SchemaId, Arc<AtomicRefCell<UntypedComponentStore>>>,
+    pub(crate) components: HashMap<SchemaId, Arc<AtomicCell<UntypedComponentStore>>>,
 }
 
 // SOUND: all of the functions for ComponentStores requires that the types stored implement Sync +
@@ -46,9 +46,9 @@ impl ComponentStores {
     /// Initialize component storage for type `T`.
     pub fn init<T: HasSchema>(&mut self) {
         let schema = T::schema();
-        self.components.entry(schema.id()).or_insert_with(|| {
-            Arc::new(AtomicRefCell::new(UntypedComponentStore::for_type::<T>()))
-        });
+        self.components
+            .entry(schema.id())
+            .or_insert_with(|| Arc::new(AtomicCell::new(UntypedComponentStore::for_type::<T>())));
     }
 
     /// Get the components of a certain type
@@ -59,16 +59,16 @@ impl ComponentStores {
         // `UntypedComponentStore`.
         unsafe {
             Ok(std::mem::transmute::<
-                Arc<AtomicRefCell<UntypedComponentStore>>,
-                Arc<AtomicRefCell<ComponentStore<T>>>,
+                Arc<AtomicCell<UntypedComponentStore>>,
+                Arc<AtomicCell<ComponentStore<T>>>,
             >(untyped))
         }
     }
 
     /// Borrow a component store.
     /// # Errors
     /// Errors if the component store has not been initialized yet.
-    pub fn get<T: HasSchema>(&self) -> Result<AtomicRef<ComponentStore<T>>, EcsError> {
+    pub fn get<T: HasSchema>(&self) -> Result<Ref<ComponentStore<T>>, EcsError> {
         let id = T::schema().id();
         let atomicref = self
             .components
@@ -77,7 +77,7 @@ impl ComponentStores {
             .borrow();
 
         // SOUND: ComponentStore<T> is repr(transparent) over UntypedComponent store.
-        let atomicref = AtomicRef::map(atomicref, |x| unsafe {
+        let atomicref = Ref::map(atomicref, |x| unsafe {
             std::mem::transmute::<&UntypedComponentStore, &ComponentStore<T>>(x)
         });
 
@@ -87,7 +87,7 @@ impl ComponentStores {
     /// Borrow a component store.
     /// # Errors
     /// Errors if the component store has not been initialized yet.
-    pub fn get_mut<T: HasSchema>(&self) -> Result<AtomicRefMut<ComponentStore<T>>, EcsError> {
+    pub fn get_mut<T: HasSchema>(&self) -> Result<RefMut<ComponentStore<T>>, EcsError> {
         let id = T::schema().id();
         let atomicref = self
             .components
@@ -96,7 +96,7 @@ impl ComponentStores {
             .borrow_mut();
 
         // SOUND: ComponentStore<T> is repr(transparent) over UntypedComponent store.
-        let atomicref = AtomicRefMut::map(atomicref, |x| unsafe {
+        let atomicref = RefMut::map(atomicref, |x| unsafe {
             std::mem::transmute::<&mut UntypedComponentStore, &mut ComponentStore<T>>(x)
         });
 
@@ -107,7 +107,7 @@ impl ComponentStores {
     pub fn get_cell_by_schema_id(
         &self,
         id: SchemaId,
-    ) -> Result<Arc<AtomicRefCell<UntypedComponentStore>>, EcsError> {
+    ) -> Result<Arc<AtomicCell<UntypedComponentStore>>, EcsError> {
         self.components
             .get(&id)
             .cloned()

framework_crates/bones_ecs/src/lib.rs

@@ -9,7 +9,7 @@ pub mod atomic {
     //! Atomic Refcells are from the [`atomic_refcell`] crate.
     //!
     //! [`atomic_refcell`]: https://docs.rs/atomic_refcell
-    pub use atomic_refcell::*;
+    pub use atomicell::*;
 }
 pub mod bitset;
 pub mod components;
@@ -30,7 +30,7 @@ pub use world::{FromWorld, World};
 /// The prelude.
 pub mod prelude {
     pub use {
-        atomic_refcell::*, bitset_core::BitSet, bones_schema::prelude::*, bones_utils::prelude::*,
+        atomicell::*, bitset_core::BitSet, bones_schema::prelude::*, bones_utils::prelude::*,
     };
 
     pub use crate::{

framework_crates/bones_ecs/src/resources.rs

@@ -2,6 +2,8 @@
 
 use std::{fmt::Debug, marker::PhantomData, sync::Arc};
 
+use atomicell::borrow::{AtomicBorrow, AtomicBorrowMut};
+
 use crate::prelude::*;
 
 /// Storage for un-typed resources.
@@ -16,10 +18,8 @@ pub struct UntypedResources {
 }
 
 /// An untyped resource that may be inserted into [`UntypedResources`].
-#[derive(Deref, DerefMut)]
 pub struct UntypedAtomicResource {
-    #[deref]
-    cell: Arc<AtomicRefCell<SchemaBox>>,
+    cell: Arc<AtomicCell<SchemaBox>>,
     schema: &'static Schema,
 }
 
@@ -28,15 +28,15 @@ impl UntypedAtomicResource {
     pub fn new(resource: SchemaBox) -> Self {
         Self {
             schema: resource.schema(),
-            cell: Arc::new(AtomicRefCell::new(resource)),
+            cell: Arc::new(AtomicCell::new(resource)),
         }
     }
 
     /// Create a new [`UntypedAtomicResource`] for the given schema, initially populated with the default
     /// value for the schema.
     pub fn from_schema(schema: &'static Schema) -> Self {
         Self {
-            cell: Arc::new(AtomicRefCell::new(SchemaBox::default(schema))),
+            cell: Arc::new(AtomicCell::new(SchemaBox::default(schema))),
             schema,
         }
     }
@@ -49,16 +49,64 @@ impl UntypedAtomicResource {
         }
     }
 
+    /// Borrow the resource.
+    pub fn borrow(&self) -> AtomicSchemaRef {
+        // SOUND: we keep the borrow along with the reference.
+        let (reference, borrow) = unsafe { Ref::into_split(self.cell.borrow()) };
+        let schema_ref = reference.as_ref();
+        AtomicSchemaRef { schema_ref, borrow }
+    }
+
+    /// Mutably borrow the resource.
+    pub fn borrow_mut(&self) -> AtomicSchemaRefMut {
+        // SOUND: we keep the borrow along with the reference.
+        let (reference, borrow) = unsafe { RefMut::into_split(self.cell.borrow_mut()) };
+        let schema_ref = reference.as_mut();
+        AtomicSchemaRefMut { schema_ref, borrow }
+    }
+
     /// Get the schema of the resource.
     pub fn schema(&self) -> &'static Schema {
         self.schema
     }
 }
 
+/// An atomic borrow of a [`SchemaRef`].
+#[derive(Deref)]
+pub struct AtomicSchemaRef<'a> {
+    #[deref]
+    schema_ref: SchemaRef<'a>,
+    borrow: AtomicBorrow<'a>,
+}
+
+impl<'a> AtomicSchemaRef<'a> {
+    /// # Safety
+    /// You must know that T represents the data in the [`SchemaRef`].
+    pub unsafe fn deref<T>(self) -> Ref<'a, T> {
+        Ref::with_borrow(self.schema_ref.deref(), self.borrow)
+    }
+}
+
+/// An atomic borrow of a [`SchemaRefMut`].
+#[derive(Deref, DerefMut)]
+pub struct AtomicSchemaRefMut<'a> {
+    #[deref]
+    schema_ref: SchemaRefMut<'a, 'a>,
+    borrow: AtomicBorrowMut<'a>,
+}
+
+impl<'a> AtomicSchemaRefMut<'a> {
+    /// # Safety
+    /// You must know that T represents the data in the [`SchemaRefMut`].
+    pub unsafe fn deref_mut<T>(self) -> RefMut<'a, T> {
+        RefMut::with_borrow(self.schema_ref.deref_mut(), self.borrow)
+    }
+}
+
 impl Clone for UntypedAtomicResource {
     fn clone(&self) -> Self {
         Self {
-            cell: Arc::new(AtomicRefCell::new(self.cell.borrow().clone())),
+            cell: Arc::new(AtomicCell::new((*self.cell.borrow()).clone())),
             schema: self.schema,
         }
     }
@@ -82,7 +130,7 @@ impl UntypedResources {
         &mut self,
         resource: UntypedAtomicResource,
     ) -> Option<UntypedAtomicResource> {
-        let id = resource.cell.borrow().schema().id();
+        let id = resource.schema().id();
         self.resources.insert(id, resource)
     }
 
@@ -92,17 +140,13 @@ impl UntypedResources {
     }
 
     /// Get a reference to an untyped resource.
-    pub fn get(&self, schema_id: SchemaId) -> Option<AtomicRef<SchemaBox>> {
-        self.resources
-            .get(&schema_id)
-            .map(|x| AtomicRefCell::borrow(&x.cell))
+    pub fn get(&self, schema_id: SchemaId) -> Option<AtomicSchemaRef> {
+        self.resources.get(&schema_id).map(|x| x.borrow())
     }
 
     /// Get a mutable reference to an untyped resource.
-    pub fn get_mut(&mut self, schema_id: SchemaId) -> Option<AtomicRefMut<SchemaBox>> {
-        self.resources
-            .get(&schema_id)
-            .map(|x| AtomicRefCell::borrow_mut(&x.cell))
+    pub fn get_mut(&mut self, schema_id: SchemaId) -> Option<AtomicSchemaRefMut> {
+        self.resources.get(&schema_id).map(|x| x.borrow_mut())
     }
 
     /// Remove a resource.
@@ -138,24 +182,18 @@ impl Resources {
     }
 
     /// Borrow a resource.
-    pub fn get<T: HasSchema>(&self) -> Option<AtomicRef<T>> {
-        self.untyped.get(T::schema().id()).map(|sbox| {
-            AtomicRef::map(sbox, |sbox| {
-                // SOUND: schema matches as checked by retreiving from the untyped store by the schema
-                // ID.
-                unsafe { sbox.as_ref().deref() }
-            })
+    pub fn get<T: HasSchema>(&self) -> Option<Ref<T>> {
+        self.untyped.get(T::schema().id()).map(|x| {
+            // SOUND: untyped resources returns data matching the schema of T.
+            unsafe { x.deref() }
         })
     }
 
     /// Mutably borrow a resource.
-    pub fn get_mut<T: HasSchema>(&mut self) -> Option<AtomicRefMut<T>> {
-        self.untyped.get_mut(T::schema().id()).map(|sbox| {
-            AtomicRefMut::map(sbox, |sbox| {
-                // SOUND: schema matches as checked by retreiving from the untyped store by the
-                // schema ID.
-                unsafe { sbox.as_mut().deref_mut() }
-            })
+    pub fn get_mut<T: HasSchema>(&mut self) -> Option<RefMut<T>> {
+        self.untyped.get_mut(T::schema().id()).map(|x| {
+            // SOUND: untyped resources returns data matching the schema of T.
+            unsafe { x.deref_mut() }
         })
     }
 
@@ -249,19 +287,19 @@ impl<T: HasSchema> AtomicResource<T> {
     /// Lock the resource for reading.
     ///
     /// This returns a read guard, very similar to an [`RwLock`][std::sync::RwLock].
-    pub fn borrow(&self) -> AtomicRef<T> {
-        let borrow = AtomicRefCell::borrow(&self.untyped);
+    pub fn borrow(&self) -> Ref<T> {
+        let borrow = self.untyped.borrow();
         // SOUND: We know that the data pointer is valid for type T.
-        AtomicRef::map(borrow, |data| unsafe { data.as_ref().deref() })
+        unsafe { borrow.deref() }
     }
 
     /// Lock the resource for read-writing.
     ///
     /// This returns a write guard, very similar to an [`RwLock`][std::sync::RwLock].
-    pub fn borrow_mut(&self) -> AtomicRefMut<T> {
-        let borrow = AtomicRefCell::borrow_mut(&self.untyped);
+    pub fn borrow_mut(&self) -> RefMut<T> {
+        let borrow = self.untyped.borrow_mut();
         // SOUND: We know that the data pointer is valid for type T.
-        AtomicRefMut::map(borrow, |data| unsafe { data.as_mut().deref_mut() })
+        unsafe { borrow.deref_mut() }
     }
 }
 
@@ -298,4 +336,28 @@ mod test {
         assert_eq!(resources.get::<B>().unwrap().0, 2);
         assert_eq!(resources.get::<A>().unwrap().0, "world");
     }
+
+    #[test]
+    fn untyped_resources_no_rewrite() {
+        let mut store = UntypedResources::default();
+
+        // Create two datas with different types.
+        let data1 = SchemaBox::new(0usize);
+        let data1_schema = data1.schema();
+        let mut data2 = SchemaBox::new(String::from("bad_data"));
+
+        {
+            // Try to "cheat" and overwrite the the data1 resource with data of a different type.
+            store.insert(data1);
+            let data1_cell = store.get_cell(data1_schema.id()).unwrap();
+            let mut data1_borrow = data1_cell.borrow_mut();
+            // This can't actually write to data1, it just changes the reference we have.
+            *data1_borrow = data2.as_mut();
+        }
+
+        // Validate that data1 is unchanged
+        let data1_cell = store.get_cell(data1_schema.id()).unwrap();
+        let data1_borrow = data1_cell.borrow();
+        assert_eq!(data1_borrow.cast::<usize>(), &0);
+    }
 }

framework_crates/bones_ecs/src/stage.rs

@@ -301,7 +301,7 @@ impl CommandQueue {
 ///
 /// This is a shortcut for [`ResMut<CommandQueue>`].
 #[derive(Deref, DerefMut)]
-pub struct Commands<'a>(AtomicRefMut<'a, CommandQueue>);
+pub struct Commands<'a>(RefMut<'a, CommandQueue>);
 
 impl<'a> SystemParam for Commands<'a> {
     type State = AtomicResource<CommandQueue>;