Skip to content

Prevent TOCTOU bugs in ComponentsQueuedRegistrator #20016

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 27 additions & 27 deletions crates/bevy_ecs/src/component.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1232,58 +1232,58 @@ impl<'w> ComponentsQueuedRegistrator<'w> {
Self { components, ids }
}

/// Queues this function to run as a component registrator.
/// Queues this function to run as a component registrator if the given
/// type is not already queued as a component.
///
/// # Safety
///
/// The [`TypeId`] must not already be registered or queued as a component.
unsafe fn force_register_arbitrary_component(
/// The [`TypeId`] must not already be registered as a component.
unsafe fn register_arbitrary_component(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why remove force from the names? They're still queuing without first checking that the components are registered, which is a kind of "forcing".

The names queue_register_component and register_arbitrary_component don't do a great job of communicating the difference, although I suppose force_register_arbitrary_component wasn't much better.

Hmm, I wonder if it would make sense to make a future PR to inline these into their callers so that we don't need two names.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see your point, but at the same time they are no longer really forcing the registration. They are more like register_if_not_enqueued or something like that, but it doesn't sound so good as a name.

&self,
type_id: TypeId,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
) -> ComponentId {
let id = self.ids.next();
self.components
.queued
.write()
.unwrap_or_else(PoisonError::into_inner)
.components
.insert(
type_id,
.entry(type_id)
.or_insert_with(|| {
// SAFETY: The id was just generated.
unsafe { QueuedRegistration::new(id, descriptor, func) },
);
id
unsafe { QueuedRegistration::new(self.ids.next(), descriptor, func) }
})
.id
}

/// Queues this function to run as a resource registrator.
/// Queues this function to run as a resource registratorif the given
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/// Queues this function to run as a resource registratorif the given
/// Queues this function to run as a resource registrator if the given

/// type is not already queued as a resource.
///
/// # Safety
///
/// The [`TypeId`] must not already be registered or queued as a resource.
unsafe fn force_register_arbitrary_resource(
/// The [`TypeId`] must not already be registered as a resource.
unsafe fn register_arbitrary_resource(
&self,
type_id: TypeId,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
) -> ComponentId {
let id = self.ids.next();
self.components
.queued
.write()
.unwrap_or_else(PoisonError::into_inner)
.resources
.insert(
type_id,
.entry(type_id)
.or_insert_with(|| {
// SAFETY: The id was just generated.
unsafe { QueuedRegistration::new(id, descriptor, func) },
);
id
unsafe { QueuedRegistration::new(self.ids.next(), descriptor, func) }
})
.id
}

/// Queues this function to run as a dynamic registrator.
fn force_register_arbitrary_dynamic(
fn register_arbitrary_dynamic(
&self,
descriptor: ComponentDescriptor,
func: impl FnOnce(&mut ComponentsRegistrator, ComponentId, ComponentDescriptor) + 'static,
Expand Down Expand Up @@ -1314,9 +1314,9 @@ impl<'w> ComponentsQueuedRegistrator<'w> {
#[inline]
pub fn queue_register_component<T: Component>(&self) -> ComponentId {
self.component_id::<T>().unwrap_or_else(|| {
// SAFETY: We just checked that this type was not in the queue.
// SAFETY: We just checked that this type was not already registered.
unsafe {
self.force_register_arbitrary_component(
self.register_arbitrary_component(
TypeId::of::<T>(),
ComponentDescriptor::new::<T>(),
|registrator, id, _descriptor| {
Expand Down Expand Up @@ -1344,7 +1344,7 @@ impl<'w> ComponentsQueuedRegistrator<'w> {
&self,
descriptor: ComponentDescriptor,
) -> ComponentId {
self.force_register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
self.register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
// SAFETY: Id uniqueness handled by caller.
unsafe {
registrator.register_component_inner(id, descriptor);
Expand All @@ -1366,9 +1366,9 @@ impl<'w> ComponentsQueuedRegistrator<'w> {
pub fn queue_register_resource<T: Resource>(&self) -> ComponentId {
let type_id = TypeId::of::<T>();
self.get_resource_id(type_id).unwrap_or_else(|| {
// SAFETY: We just checked that this type was not in the queue.
// SAFETY: We just checked that this type was not already registered.
unsafe {
self.force_register_arbitrary_resource(
self.register_arbitrary_resource(
type_id,
ComponentDescriptor::new_resource::<T>(),
move |registrator, id, descriptor| {
Expand Down Expand Up @@ -1398,9 +1398,9 @@ impl<'w> ComponentsQueuedRegistrator<'w> {
pub fn queue_register_non_send<T: Any>(&self) -> ComponentId {
let type_id = TypeId::of::<T>();
self.get_resource_id(type_id).unwrap_or_else(|| {
// SAFETY: We just checked that this type was not in the queue.
// SAFETY: We just checked that this type was not already registered.
unsafe {
self.force_register_arbitrary_resource(
self.register_arbitrary_resource(
type_id,
ComponentDescriptor::new_non_send::<T>(StorageType::default()),
move |registrator, id, descriptor| {
Expand Down Expand Up @@ -1429,7 +1429,7 @@ impl<'w> ComponentsQueuedRegistrator<'w> {
&self,
descriptor: ComponentDescriptor,
) -> ComponentId {
self.force_register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
self.register_arbitrary_dynamic(descriptor, |registrator, id, descriptor| {
// SAFETY: Id uniqueness handled by caller.
unsafe {
registrator.register_component_inner(id, descriptor);
Expand Down
13 changes: 13 additions & 0 deletions crates/bevy_ecs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2776,4 +2776,17 @@ mod tests {

fn custom_clone(_source: &SourceComponent, _ctx: &mut ComponentCloneCtx) {}
}

#[test]
fn queue_register_component_toctou() {
for _ in 0..1000 {
let w = World::new();

std::thread::scope(|s| {
let c1 = s.spawn(|| w.components_queue().queue_register_component::<A>());
let c2 = s.spawn(|| w.components_queue().queue_register_component::<A>());
assert_eq!(c1.join().unwrap(), c2.join().unwrap());
});
}
}
}