Skip to content

Commit

Permalink
Improve soundness of CommandQueue (#4863)
Browse files Browse the repository at this point in the history
# 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<MaybeUninit<u8>>` instead of `Vec<u8>`;
- 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.
  • Loading branch information
SkiFire13 committed May 30, 2022
1 parent 4828998 commit e543941
Showing 1 changed file with 33 additions and 28 deletions.
61 changes: 33 additions & 28 deletions crates/bevy_ecs/src/system/commands/command_queue.rs
Original file line number Diff line number Diff line change
@@ -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>,
}

Expand All @@ -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);
}
Expand All @@ -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.
Expand All @@ -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);
}
}
}
Expand Down Expand Up @@ -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);
}
}

0 comments on commit e543941

Please sign in to comment.