From dea05e9af557b46c4893860bb94c03ee3649f1ae Mon Sep 17 00:00:00 2001
From: Boxy <supbscripter@gmail.com>
Date: Fri, 30 Oct 2020 18:43:39 +0000
Subject: [PATCH] Remove unsound cast in thread local resources (#749)

* Remove unsound cast in thread local resources

* Make ResourceRef(Mut)::new impossible to cause unsoundness with
---
 crates/bevy_ecs/src/resource/resources.rs | 59 +++++++++++++++++------
 1 file changed, 43 insertions(+), 16 deletions(-)

diff --git a/crates/bevy_ecs/src/resource/resources.rs b/crates/bevy_ecs/src/resource/resources.rs
index e957ae674e..acec0d2f07 100644
--- a/crates/bevy_ecs/src/resource/resources.rs
+++ b/crates/bevy_ecs/src/resource/resources.rs
@@ -33,7 +33,7 @@ trait ResourceStorage: Downcast {}
 impl_downcast!(ResourceStorage);
 
 struct StoredResource<T: 'static> {
-    value: T,
+    value: std::cell::UnsafeCell<T>,
     atomic_borrow: AtomicBorrow,
 }
 
@@ -45,27 +45,24 @@ impl<T: 'static> VecResourceStorage<T> {
     fn get(&self, index: usize) -> Option<ResourceRef<'_, T>> {
         self.stored
             .get(index)
-            .map(|stored| ResourceRef::new(&stored.value, &stored.atomic_borrow))
+            .map(|stored| ResourceRef::new(stored))
     }
 
     fn get_mut(&self, index: usize) -> Option<ResourceRefMut<'_, T>> {
-        self.stored.get(index).map(|stored|
-            // SAFE: ResourceRefMut ensures that this borrow is unique
-            unsafe {
-                let value = &stored.value as *const T as *mut T;
-                ResourceRefMut::new(&mut *value, &stored.atomic_borrow)
-            })
+        self.stored
+            .get(index)
+            .map(|stored| ResourceRefMut::new(stored))
     }
 
     fn push(&mut self, resource: T) {
         self.stored.push(StoredResource {
             atomic_borrow: AtomicBorrow::new(),
-            value: resource,
+            value: std::cell::UnsafeCell::new(resource),
         })
     }
 
     fn set(&mut self, index: usize, resource: T) {
-        self.stored[index].value = resource;
+        self.stored[index].value = std::cell::UnsafeCell::new(resource);
     }
 
     fn is_empty(&self) -> bool {
@@ -414,9 +411,24 @@ pub struct ResourceRef<'a, T: 'static> {
 
 impl<'a, T: 'static> ResourceRef<'a, T> {
     /// Creates a new resource borrow
-    pub fn new(resource: &'a T, borrow: &'a AtomicBorrow) -> Self {
-        borrow.borrow();
-        Self { resource, borrow }
+    fn new(
+        StoredResource {
+            value,
+            atomic_borrow,
+        }: &'a StoredResource<T>,
+    ) -> Self {
+        if atomic_borrow.borrow() {
+            Self {
+                // Safe because we acquired the lock
+                resource: unsafe { &*value.get() },
+                borrow: atomic_borrow,
+            }
+        } else {
+            panic!(
+                "Failed to acquire shared lock on resource: {}",
+                std::any::type_name::<T>()
+            );
+        }
     }
 }
 
@@ -454,9 +466,24 @@ pub struct ResourceRefMut<'a, T: 'static> {
 
 impl<'a, T: 'static> ResourceRefMut<'a, T> {
     /// Creates a new entity component mutable borrow
-    pub fn new(resource: &'a mut T, borrow: &'a AtomicBorrow) -> Self {
-        borrow.borrow_mut();
-        Self { resource, borrow }
+    fn new(
+        StoredResource {
+            value,
+            atomic_borrow,
+        }: &'a StoredResource<T>,
+    ) -> Self {
+        if atomic_borrow.borrow_mut() {
+            Self {
+                // Safe because we acquired the lock
+                resource: unsafe { &mut *value.get() },
+                borrow: atomic_borrow,
+            }
+        } else {
+            panic!(
+                "Failed to acquire exclusive lock on resource: {}",
+                std::any::type_name::<T>()
+            );
+        }
     }
 }