guard the WorldPointer against UB (#83)

Author: jrobsonchase

bevy_mod_scripting_core/src/world.rs

@@ -1,3 +1,4 @@
+use std::ops::Deref;
 use std::sync::Arc;
 
 use bevy::prelude::World;
@@ -6,37 +7,153 @@ use parking_lot::{
 };
 
 /// Pointer to a bevy world, safely allows multiple access via RwLock
-/// # Safety
-/// This pointer does not prevent dangling pointers, i.e. you must ensure the world is not dropped while any world pointers still exist,
-/// the world must also not change, from the moment a world pointer is created it must always point to the same world.
+///
+/// If the original [WorldPointerGuard] that created this pointer is dropped,
+/// the `read` and `write` methods will panic, and the "try" variants will
+/// return `None`.
 #[derive(Debug, Clone)]
-pub struct WorldPointer(Arc<RwLock<*mut World>>);
+pub struct WorldPointer(Arc<RwLock<Option<*mut World>>>);
+
+/// Guarded pointer to a bevy world, can be used to `clone` additional
+/// [WorldPointer]s for safe access.
+///
+/// # Safety
+/// The original `&mut World` used to create this guard _must not_ be used after
+/// the guard is created in order for the cloned [WorldPointer]s to be safe.
+///
+/// On [Drop], it will "take back" access to the `&mut World`, preventing the
+/// `WorldPointer`s from invoking UB.
+#[derive(Debug)]
+pub struct WorldPointerGuard(WorldPointer);
+
+impl Deref for WorldPointerGuard {
+    type Target = WorldPointer;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
 
 unsafe impl Send for WorldPointer {}
 unsafe impl Sync for WorldPointer {}
 
-impl WorldPointer {
+impl WorldPointerGuard {
     /// Creates a new world pointer.
     /// # Safety
-    /// satisfies world constancy, since it's impossible to change the underlying pointer
-    /// However you must ensure that the world does not go out of scope while this pointer is live   
+    /// The original `&mut World` must not be used while this guard is in scope.
+    /// The [World] may only be accessed through this guard or one of its cloned
+    /// [WorldPointer]s.
     #[allow(clippy::arc_with_non_send_sync)]
     pub unsafe fn new(world: &mut World) -> Self {
-        WorldPointer(Arc::new(RwLock::new(world)))
+        WorldPointerGuard(WorldPointer(Arc::new(RwLock::new(Some(world)))))
+    }
+}
+
+impl Drop for WorldPointerGuard {
+    fn drop(&mut self) {
+        // Being explicit about the types here to make sure we're getting things
+        // correct.
+        let world_ptr: &WorldPointer = &self.0;
+        let _: Option<*mut World> = RwLock::write(&world_ptr.0).take();
     }
+}
 
+impl WorldPointer {
     /// Returns a read guard which can be used for immutable world access.
+    ///
+    /// Panics if the pointer is already locked or has gone out of scope.
     pub fn read(&self) -> MappedRwLockReadGuard<World> {
-        RwLockReadGuard::map(self.0.try_read().expect(""), |ptr: &*mut World| unsafe {
-            &**ptr
-        })
+        self.try_read().expect("concurrent read/write world access")
     }
 
     /// Returns a write guard which can be used for mutable world access.
+    ///
+    /// Panics if the pointer is already locked or has gone out of scope.
     pub fn write(&self) -> MappedRwLockWriteGuard<World> {
-        RwLockWriteGuard::map(
-            self.0.try_write().expect(""),
-            |ptr: &mut *mut World| unsafe { &mut **ptr },
-        )
+        self.try_write()
+            .expect("concurrent read/write world access")
+    }
+
+    /// Returns a read guard which can be used for immutable world access.
+    ///
+    /// Returns `None` if the pointer is already locked or has gone out of
+    /// scope.
+    pub fn try_read(&self) -> Option<MappedRwLockReadGuard<World>> {
+        self.try_read_inner(false)
+    }
+
+    /// Returns a write guard which can be used for mutable world access.
+    ///
+    /// Returns `None` if the pointer is already locked or has gone out of
+    /// scope.
+    pub fn try_write(&self) -> Option<MappedRwLockWriteGuard<World>> {
+        self.try_write_inner(false)
+    }
+
+    /// Returns a read guard which can be used for immutable world access.
+    ///
+    /// Panics if the pointer has gone out of scope. May block if another thread
+    /// holds the lock.
+    pub fn read_blocking(&self) -> MappedRwLockReadGuard<World> {
+        self.try_read_blocking()
+            .expect("the world pointer is out of scope")
+    }
+
+    /// Returns a write guard which can be used for mutable world access.
+    ///
+    /// Panics if the pointer has gone out of scope. May block if another thread
+    /// holds the lock.
+    pub fn write_blocking(&self) -> MappedRwLockWriteGuard<World> {
+        self.try_write_blocking()
+            .expect("the world pointer is out of scope")
+    }
+
+    /// Returns a read guard which can be used for immutable world access.
+    ///
+    /// Returns `None` if has gone out of scope. May block if another thread
+    /// holds the lock.
+    pub fn try_read_blocking(&self) -> Option<MappedRwLockReadGuard<World>> {
+        self.try_read_inner(true)
+    }
+
+    /// Returns a write guard which can be used for mutable world access.
+    ///
+    /// Returns `None` if has gone out of scope. May block if another thread
+    /// holds the lock.
+    pub fn try_write_blocking(&self) -> Option<MappedRwLockWriteGuard<World>> {
+        self.try_write_inner(true)
+    }
+
+    fn try_read_inner(&self, blocking: bool) -> Option<MappedRwLockReadGuard<World>> {
+        let guard = if blocking {
+            self.0.read()
+        } else {
+            self.0.try_read()?
+        };
+        // Check if the inner pointer is there so we can invert the `Option`.
+        if guard.is_none() {
+            return None;
+        }
+
+        Some(RwLockReadGuard::map(
+            guard,
+            |ptr: &Option<*mut World>| unsafe { &*ptr.unwrap() },
+        ))
+    }
+
+    fn try_write_inner(&self, blocking: bool) -> Option<MappedRwLockWriteGuard<World>> {
+        let guard = if blocking {
+            self.0.write()
+        } else {
+            self.0.try_write()?
+        };
+        // Check if the inner pointer is there so we can invert the `Option`.
+        if guard.is_none() {
+            return None;
+        }
+
+        Some(RwLockWriteGuard::map(
+            guard,
+            |ptr: &mut Option<*mut World>| unsafe { &mut *ptr.unwrap() },
+        ))
     }
 }

languages/bevy_mod_scripting_lua/src/lib.rs

@@ -3,7 +3,7 @@ use crate::{
     docs::LuaDocFragment,
 };
 use bevy::{ecs::schedule::ScheduleLabel, prelude::*};
-use bevy_mod_scripting_core::{prelude::*, systems::*, world::WorldPointer};
+use bevy_mod_scripting_core::{prelude::*, systems::*, world::WorldPointerGuard};
 
 use std::fmt;
 use std::marker::PhantomData;
@@ -146,12 +146,12 @@ impl<A: LuaArg> ScriptHost for LuaScriptHost<A> {
     ) {
         // safety:
         // - we have &mut World access
-        // - we do not use world_ptr after using the world reference which it's derived from
-        let world_ptr = unsafe { WorldPointer::new(world) };
+        // - we do not use the original reference again anywhere in this function
+        let world = unsafe { WorldPointerGuard::new(world) };
 
         ctxs.for_each(|(script_data, ctx)| {
             providers
-                .setup_runtime_all(world_ptr.clone(), &script_data, ctx)
+                .setup_runtime_all(world.clone(), &script_data, ctx)
                 .expect("Could not setup script runtime");
 
             let ctx = ctx.get_mut().expect("Poison error in context");
@@ -172,7 +172,7 @@ impl<A: LuaArg> ScriptHost for LuaScriptHost<A> {
                 };
 
                 if let Err(error) = f.call::<_, ()>(event.args.clone()) {
-                    let mut world = world_ptr.write();
+                    let mut world = world.write();
                     let mut state: CachedScriptState<Self> = world.remove_resource().unwrap();
 
                     let (_, mut error_wrt, _) = state.event_state.get_mut(&mut world);

languages/bevy_mod_scripting_rhai/src/lib.rs

@@ -3,7 +3,7 @@ use crate::{
     docs::RhaiDocFragment,
 };
 use bevy::{ecs::schedule::ScheduleLabel, prelude::*};
-use bevy_mod_scripting_core::{prelude::*, systems::*, world::WorldPointer};
+use bevy_mod_scripting_core::{prelude::*, systems::*, world::WorldPointerGuard};
 use rhai::*;
 use std::marker::PhantomData;
 
@@ -150,12 +150,12 @@ impl<A: FuncArgs + Send + Clone + Sync + 'static> ScriptHost for RhaiScriptHost<
     ) {
         // safety:
         // - we have &mut World access
-        // - we do not use world_ptr after we use the original reference again anywhere in this function
-        let world_ptr = unsafe { WorldPointer::new(world) };
+        // - we do not use the original reference again anywhere in this function
+        let world = unsafe { WorldPointerGuard::new(world) };
 
         ctxs.for_each(|(fd, ctx)| {
             providers
-                .setup_runtime_all(world_ptr.clone(), &fd, ctx)
+                .setup_runtime_all(world.clone(), &fd, ctx)
                 .expect("Failed to setup script runtime");
 
             for event in events.iter() {
@@ -172,7 +172,7 @@ impl<A: FuncArgs + Send + Clone + Sync + 'static> ScriptHost for RhaiScriptHost<
                 ) {
                     Ok(v) => v,
                     Err(e) => {
-                        let mut world = world_ptr.write();
+                        let mut world = world.write();
                         let mut state: CachedScriptState<Self> = world.remove_resource().unwrap();
 
                         match *e {