diff --git a/crates/bevy_ecs/src/component.rs b/crates/bevy_ecs/src/component.rs index 615c5903f8fa8..3168044861695 100644 --- a/crates/bevy_ecs/src/component.rs +++ b/crates/bevy_ecs/src/component.rs @@ -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( &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 + /// 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, @@ -1314,9 +1314,9 @@ impl<'w> ComponentsQueuedRegistrator<'w> { #[inline] pub fn queue_register_component(&self) -> ComponentId { self.component_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_component( + self.register_arbitrary_component( TypeId::of::(), ComponentDescriptor::new::(), |registrator, id, _descriptor| { @@ -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); @@ -1366,9 +1366,9 @@ impl<'w> ComponentsQueuedRegistrator<'w> { pub fn queue_register_resource(&self) -> ComponentId { let type_id = TypeId::of::(); 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::(), move |registrator, id, descriptor| { @@ -1398,9 +1398,9 @@ impl<'w> ComponentsQueuedRegistrator<'w> { pub fn queue_register_non_send(&self) -> ComponentId { let type_id = TypeId::of::(); 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::(StorageType::default()), move |registrator, id, descriptor| { @@ -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); diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 8a07cdc8e1b92..96fd542b61bc3 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -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::()); + let c2 = s.spawn(|| w.components_queue().queue_register_component::()); + assert_eq!(c1.join().unwrap(), c2.join().unwrap()); + }); + } + } }