From e58fffb59d6702eb666581b6ba796f8a515a4f12 Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Mon, 30 May 2022 22:45:09 +0000 Subject: [PATCH] Improve soundness of `CommandQueue` (#4863) # Objective This PR aims to improve the soundness of `CommandQueue`. In particular it aims to: - make it sound to store commands that contain padding or uninitialized bytes; - avoid uses of commands after moving them in the queue's buffer (`std::mem::forget` is technically a use of its argument); - remove useless checks: `self.bytes.as_mut_ptr().is_null()` is always `false` because even `Vec`s that haven't allocated use a dangling pointer. Moreover the same pointer was used to write the command, so it ought to be valid for reads if it was for writes. ## Solution - To soundly store padding or uninitialized bytes `CommandQueue` was changed to contain a `Vec>` instead of `Vec`; - To avoid uses of the command through `std::mem::forget`, `ManuallyDrop` was used. ## Other observations While writing this PR I noticed that `CommandQueue` doesn't seem to drop the commands that weren't applied. While this is a pretty niche case (you would have to be manually using `CommandQueue`/`std::mem::swap`ping one), I wonder if it should be documented anyway. --- .../src/system/commands/command_queue.rs | 61 ++++++++++--------- 1 file changed, 33 insertions(+), 28 deletions(-) diff --git a/crates/bevy_ecs/src/system/commands/command_queue.rs b/crates/bevy_ecs/src/system/commands/command_queue.rs index 64a3bd7914a6fe..d73b529253f66d 100644 --- a/crates/bevy_ecs/src/system/commands/command_queue.rs +++ b/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, world: &mut World), } /// A queue of [`Command`]s // -// NOTE: [`CommandQueue`] is implemented via a `Vec` over a `Vec>` +// NOTE: [`CommandQueue`] is implemented via a `Vec>` over a `Vec>` // 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, + bytes: Vec>, metas: Vec, } @@ -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(command: *mut u8, world: &mut World) { + unsafe fn write_command(command: *mut MaybeUninit, world: &mut World) { let command = command.cast::().read_unaligned(); command.write(world); } @@ -48,25 +50,30 @@ impl CommandQueue { func: write_command::, }); + // 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, 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); + } }