Fix unsoundness in QueryIter::sort_by (#17826)

# Objective

`QueryIter::sort_by()` is unsound. It passes the lens items with the
full `'w` lifetime, and a malicious user could smuggle them out of the
closure where they could alias with the query results.

## Solution

Make the sort closures generic in the lifetime parameter of the lens
item. This ensures the lens items cannot outlive the call to the
closure.

## Testing

Added a compile-fail test that demonstrates the unsound pattern.

## Migration Guide

The `sort` family of methods on `QueryIter` unsoundly gave access
`L::Item<'w>` with the full `'w` lifetime. It has been shortened to
`L::Item<'w>` so that items cannot escape the comparer. If you get
lifetime errors using these methods, you will need to make the comparer
generic in the new lifetime. Often this can be done by replacing named
`'w` with `'_`, or by replacing the use of a function item with a
closure.

```rust
// Before: Now fails with "error: implementation of `FnMut` is not general enough"
query.iter().sort_by::<&C>(Ord::cmp);
// After: Wrap in a closure
query.iter().sort_by::<&C>(|l, r| Ord::cmp(l, r));

query.iter().sort_by::<&C>(comparer);
// Before: Uses specific `'w` lifetime from some outer scope
// now fails with "error: implementation of `FnMut` is not general enough"
fn comparer(left: &&'w C, right: &&'w C) -> Ordering { /* ... */ }
// After: Accepts any lifetime using inferred lifetime parameter
fn comparer(left: &&C, right: &&C) -> Ordering { /* ... */ }
This commit is contained in:
Chris Russell 2025-02-26 15:36:37 -05:00 committed by GitHub
parent 23cf8ffe6c
commit 94e6fa168f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 59 additions and 24 deletions

View File

@ -0,0 +1,19 @@
use bevy_ecs::prelude::*;
use std::cmp::Ordering;
#[derive(Component)]
struct A(usize);
fn system(mut query: Query<&mut A>) {
let iter = query.iter_mut();
let mut stored: Option<&A> = None;
let mut sorted = iter.sort_by::<&A>(|left, _right| {
// Try to smuggle the lens item out of the closure.
stored = Some(left);
//~^ E0521
Ordering::Equal
});
let r: &A = stored.unwrap();
let m: &mut A = &mut sorted.next().unwrap();
assert!(std::ptr::eq(r, m));
}

View File

@ -0,0 +1,10 @@
error[E0521]: borrowed data escapes outside of closure
--> tests/ui\system_query_iter_sort_lifetime_safety.rs:12:9
|
9 | let mut stored: Option<&A> = None;
| ---------- `stored` declared here, outside of the closure body
10 | let mut sorted = iter.sort_by::<&A>(|left, _right| {
| ---- `left` is a reference that is only valid in the closure body
11 | // Try to smuggle the lens item out of the closure.
12 | stored = Some(left);
| ^^^^^^^^^^^^^^^^^^^ `left` escapes the closure body here

View File

@ -487,7 +487,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// # schedule.add_systems((system_1, system_2, system_3));
/// # schedule.run(&mut world);
/// ```
pub fn sort<L: ReadOnlyQueryData<Item<'w>: Ord> + 'w>(
pub fn sort<L: ReadOnlyQueryData + 'w>(
self,
) -> QuerySortedIter<
'w,
@ -495,7 +495,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
>
where
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
}
@ -541,7 +544,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// # schedule.add_systems((system_1));
/// # schedule.run(&mut world);
/// ```
pub fn sort_unstable<L: ReadOnlyQueryData<Item<'w>: Ord> + 'w>(
pub fn sort_unstable<L: ReadOnlyQueryData + 'w>(
self,
) -> QuerySortedIter<
'w,
@ -549,7 +552,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
>
where
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
}
@ -604,7 +610,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// ```
pub fn sort_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedIter<
'w,
's,
@ -636,7 +642,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedIter<
'w,
's,
@ -688,7 +694,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// #[derive(Component)]
/// struct AvailableMarker;
///
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
/// enum Rarity {
/// Common,
/// Rare,
@ -716,7 +722,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
/// (
/// entity_ref.contains::<AvailableMarker>(),
/// entity_ref.get::<Rarity>()
/// entity_ref.get::<Rarity>().copied()
/// )
/// })
/// .rev()
@ -728,7 +734,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// ```
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedIter<
'w,
's,
@ -761,7 +767,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedIter<
'w,
's,
@ -796,7 +802,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedIter<
'w,
's,
@ -826,7 +832,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
fn sort_impl<L: ReadOnlyQueryData + 'w>(
self,
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
f: impl FnOnce(&mut Vec<(L::Item<'_>, NeutralOrd<Entity>)>),
) -> QuerySortedIter<
'w,
's,
@ -1334,7 +1340,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
>
where
L::Item<'w>: Ord,
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
}
@ -1392,7 +1398,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
>
where
L::Item<'w>: Ord,
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
}
@ -1449,7 +1455,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// ```
pub fn sort_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedManyIter<
'w,
's,
@ -1480,7 +1486,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedManyIter<
'w,
's,
@ -1532,7 +1538,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// #[derive(Component)]
/// struct AvailableMarker;
///
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
/// enum Rarity {
/// Common,
/// Rare,
@ -1562,7 +1568,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
/// (
/// entity_ref.contains::<AvailableMarker>(),
/// entity_ref.get::<Rarity>()
// entity_ref.get::<Rarity>().copied()
/// )
/// })
/// .rev()
@ -1574,7 +1580,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// ```
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedManyIter<
'w,
's,
@ -1606,7 +1612,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedManyIter<
'w,
's,
@ -1640,7 +1646,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedManyIter<
'w,
's,
@ -1669,7 +1675,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
fn sort_impl<L: ReadOnlyQueryData + 'w>(
self,
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
f: impl FnOnce(&mut Vec<(L::Item<'_>, NeutralOrd<Entity>)>),
) -> QuerySortedManyIter<
'w,
's,
@ -2902,13 +2908,13 @@ mod tests {
{
let mut query = query_state
.iter_many_mut(&mut world, [id, id])
.sort_by::<&C>(Ord::cmp);
.sort_by::<&C>(|l, r| Ord::cmp(l, r));
while query.fetch_next().is_some() {}
}
{
let mut query = query_state
.iter_many_mut(&mut world, [id, id])
.sort_unstable_by::<&C>(Ord::cmp);
.sort_unstable_by::<&C>(|l, r| Ord::cmp(l, r));
while query.fetch_next().is_some() {}
}
{