[Merged by Bors] - Improve soundness of CommandQueue

crates/bevy_ecs/src/system/commands/command_queue.rs

@@ -1,21 +1,23 @@
+use std::mem::{ManuallyDrop, MaybeUninit};
+
 use super::Command;
 use crate::world::World;
 
 struct CommandMeta {
     offset: usize,
-    func: unsafe fn(value: *mut u8, world: &mut World),
+    func: unsafe fn(value: *mut MaybeUninit<u8>, world: &mut World),
 }
 
 /// A queue of [`Command`]s
 //
-// NOTE: [`CommandQueue`] is implemented via a `Vec<u8>` over a `Vec<Box<dyn Command>>`
+// NOTE: [`CommandQueue`] is implemented via a `Vec<MaybeUninit<u8>>` over a `Vec<Box<dyn Command>>`
 // as an optimization. Since commands are used frequently in systems as a way to spawn
 // entities/components/resources, and it's not currently possible to parallelize these
 // due to mutable [`World`] access, maximizing performance for [`CommandQueue`] is
 // preferred to simplicity of implementation.
 #[derive(Default)]
 pub struct CommandQueue {
-    bytes: Vec<u8>,
+    bytes: Vec<MaybeUninit<u8>>,
     metas: Vec<CommandMeta>,
 }
 
@@ -35,7 +37,7 @@ impl CommandQueue {
         /// SAFE: This function is only every called when the `command` bytes is the associated
         /// [`Commands`] `T` type. Also this only reads the data via `read_unaligned` so unaligned
         /// accesses are safe.
-        unsafe fn write_command<T: Command>(command: *mut u8, world: &mut World) {
+        unsafe fn write_command<T: Command>(command: *mut MaybeUninit<u8>, world: &mut World) {
             let command = command.cast::<T>().read_unaligned();
             command.write(world);
         }
@@ -48,25 +50,30 @@ impl CommandQueue {
             func: write_command::<C>,
         });
 
+        // Use `ManuallyDrop` to forget `command` right away, avoiding
+        // any use of it after the `ptr::copy_nonoverlapping`.
+        let command = ManuallyDrop::new(command);
+
         if size > 0 {
             self.bytes.reserve(size);
 
             // SAFE: The internal `bytes` vector has enough storage for the
-            // command (see the call the `reserve` above), and the vector has
-            // its length set appropriately.
-            // Also `command` is forgotten at the end of this function so that
-            // when `apply` is called later, a double `drop` does not occur.
+            // command (see the call the `reserve` above), the vector has
+            // its length set appropriately and can contain any kind of bytes.
+            // In case we're writing a ZST and the `Vec` hasn't allocated yet
+            // then `as_mut_ptr` will be a dangling (non null) pointer, and
+            // thus valid for ZST writes.
+            // Also `command` is forgotten so that  when `apply` is called
+            // later, a double `drop` does not occur.
             unsafe {
                 std::ptr::copy_nonoverlapping(
-                    &command as *const C as *const u8,
+                    &*command as *const C as *const MaybeUninit<u8>,
                     self.bytes.as_mut_ptr().add(old_len),
                     size,
                 );
                 self.bytes.set_len(old_len + size);
             }
         }
-
-        std::mem::forget(command);
     }
 
     /// Execute the queued [`Command`]s in the world.
@@ -81,27 +88,12 @@ impl CommandQueue {
         // unnecessary allocations.
         unsafe { self.bytes.set_len(0) };
 
-        let byte_ptr = if self.bytes.as_mut_ptr().is_null() {
-            // SAFE: If the vector's buffer pointer is `null` this mean nothing has been pushed to its bytes.
-            // This means either that:
-            //
-            // 1) There are no commands so this pointer will never be read/written from/to.
-            //
-            // 2) There are only zero-sized commands pushed.
-            //    According to https://doc.rust-lang.org/std/ptr/index.html
-            //    "The canonical way to obtain a pointer that is valid for zero-sized accesses is NonNull::dangling"
-            //    therefore it is safe to call `read_unaligned` on a pointer produced from `NonNull::dangling` for
-            //    zero-sized commands.
-            unsafe { std::ptr::NonNull::dangling().as_mut() }
-        } else {
-            self.bytes.as_mut_ptr()
-        };
-
         for meta in self.metas.drain(..) {
             // SAFE: The implementation of `write_command` is safe for the according Command type.
+            // It's ok to read from `bytes.as_mut_ptr()` because we just wrote to it in `push`.
             // The bytes are safely cast to their original type, safely read, and then dropped.
             unsafe {
-                (meta.func)(byte_ptr.add(meta.offset), world);
+                (meta.func)(self.bytes.as_mut_ptr().add(meta.offset), world);
             }
         }
     }
@@ -234,4 +226,17 @@ mod test {
     fn test_command_is_send() {
         assert_is_send(SpawnCommand);
     }
+
+    struct CommandWithPadding(u8, u16);
+    impl Command for CommandWithPadding {
+        fn write(self, _: &mut World) {}
+    }
+
+    #[cfg(miri)]
+    #[test]
+    fn test_uninit_bytes() {
+        let mut queue = CommandQueue::default();
+        queue.push(CommandWithPadding(0, 0));
+        let _ = format!("{:?}", queue.bytes);
+    }
 }