Skip to content

Commit 94e6fa1

Browse files
chescockchescock
authored
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 { /* ... */ }
1 parent 23cf8ff commit 94e6fa1

File tree

3 files changed

+59
-24
lines changed

3 files changed

+59
-24
lines changed

‎crates/bevy_ecs/compile_fail/tests/ui/system_query_iter_sort_lifetime_safety.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
use bevy_ecs::prelude::*;
2+
use std::cmp::Ordering;
3+
4+
#[derive(Component)]
5+
struct A(usize);
6+
7+
fn system(mut query: Query<&mut A>) {
8+
let iter = query.iter_mut();
9+
let mut stored: Option<&A> = None;
10+
let mut sorted = iter.sort_by::<&A>(|left, _right| {
11+
// Try to smuggle the lens item out of the closure.
12+
stored = Some(left);
13+
//~^ E0521
14+
Ordering::Equal
15+
});
16+
let r: &A = stored.unwrap();
17+
let m: &mut A = &mut sorted.next().unwrap();
18+
assert!(std::ptr::eq(r, m));
19+
}

‎crates/bevy_ecs/compile_fail/tests/ui/system_query_iter_sort_lifetime_safety.stderr

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

‎crates/bevy_ecs/src/query/iter.rs

Lines changed: 30 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -487,15 +487,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
487487
/// # schedule.add_systems((system_1, system_2, system_3));
488488
/// # schedule.run(&mut world);
489489
/// ```
490-
pub fn sort<L: ReadOnlyQueryData<Item<'w>: Ord> + 'w>(
490+
pub fn sort<L: ReadOnlyQueryData + 'w>(
491491
self,
492492
) -> QuerySortedIter<
493493
'w,
494494
's,
495495
D,
496496
F,
497497
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
498-
> {
498+
>
499+
where
500+
for<'lw> L::Item<'lw>: Ord,
501+
{
499502
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
500503
}
501504

@@ -541,15 +544,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
541544
/// # schedule.add_systems((system_1));
542545
/// # schedule.run(&mut world);
543546
/// ```
544-
pub fn sort_unstable<L: ReadOnlyQueryData<Item<'w>: Ord> + 'w>(
547+
pub fn sort_unstable<L: ReadOnlyQueryData + 'w>(
545548
self,
546549
) -> QuerySortedIter<
547550
'w,
548551
's,
549552
D,
550553
F,
551554
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
552-
> {
555+
>
556+
where
557+
for<'lw> L::Item<'lw>: Ord,
558+
{
553559
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
554560
}
555561

@@ -604,7 +610,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
604610
/// ```
605611
pub fn sort_by<L: ReadOnlyQueryData + 'w>(
606612
self,
607-
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
613+
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
608614
) -> QuerySortedIter<
609615
'w,
610616
's,
@@ -636,7 +642,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
636642
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
637643
pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
638644
self,
639-
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
645+
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
640646
) -> QuerySortedIter<
641647
'w,
642648
's,
@@ -688,7 +694,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
688694
/// #[derive(Component)]
689695
/// struct AvailableMarker;
690696
///
691-
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
697+
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
692698
/// enum Rarity {
693699
/// Common,
694700
/// Rare,
@@ -716,7 +722,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
716722
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
717723
/// (
718724
/// entity_ref.contains::<AvailableMarker>(),
719-
/// entity_ref.get::<Rarity>()
725+
/// entity_ref.get::<Rarity>().copied()
720726
/// )
721727
/// })
722728
/// .rev()
@@ -728,7 +734,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
728734
/// ```
729735
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
730736
self,
731-
mut f: impl FnMut(&L::Item<'w>) -> K,
737+
mut f: impl FnMut(&L::Item<'_>) -> K,
732738
) -> QuerySortedIter<
733739
'w,
734740
's,
@@ -761,7 +767,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
761767
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
762768
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
763769
self,
764-
mut f: impl FnMut(&L::Item<'w>) -> K,
770+
mut f: impl FnMut(&L::Item<'_>) -> K,
765771
) -> QuerySortedIter<
766772
'w,
767773
's,
@@ -796,7 +802,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
796802
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
797803
pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
798804
self,
799-
mut f: impl FnMut(&L::Item<'w>) -> K,
805+
mut f: impl FnMut(&L::Item<'_>) -> K,
800806
) -> QuerySortedIter<
801807
'w,
802808
's,
@@ -826,7 +832,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
826832
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
827833
fn sort_impl<L: ReadOnlyQueryData + 'w>(
828834
self,
829-
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
835+
f: impl FnOnce(&mut Vec<(L::Item<'_>, NeutralOrd<Entity>)>),
830836
) -> QuerySortedIter<
831837
'w,
832838
's,
@@ -1334,7 +1340,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
13341340
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
13351341
>
13361342
where
1337-
L::Item<'w>: Ord,
1343+
for<'lw> L::Item<'lw>: Ord,
13381344
{
13391345
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
13401346
}
@@ -1392,7 +1398,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
13921398
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
13931399
>
13941400
where
1395-
L::Item<'w>: Ord,
1401+
for<'lw> L::Item<'lw>: Ord,
13961402
{
13971403
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
13981404
}
@@ -1449,7 +1455,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
14491455
/// ```
14501456
pub fn sort_by<L: ReadOnlyQueryData + 'w>(
14511457
self,
1452-
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
1458+
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
14531459
) -> QuerySortedManyIter<
14541460
'w,
14551461
's,
@@ -1480,7 +1486,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
14801486
/// called on [`QueryManyIter`] before.
14811487
pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
14821488
self,
1483-
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
1489+
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
14841490
) -> QuerySortedManyIter<
14851491
'w,
14861492
's,
@@ -1532,7 +1538,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
15321538
/// #[derive(Component)]
15331539
/// struct AvailableMarker;
15341540
///
1535-
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
1541+
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
15361542
/// enum Rarity {
15371543
/// Common,
15381544
/// Rare,
@@ -1562,7 +1568,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
15621568
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
15631569
/// (
15641570
/// entity_ref.contains::<AvailableMarker>(),
1565-
/// entity_ref.get::<Rarity>()
1571+
// entity_ref.get::<Rarity>().copied()
15661572
/// )
15671573
/// })
15681574
/// .rev()
@@ -1574,7 +1580,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
15741580
/// ```
15751581
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
15761582
self,
1577-
mut f: impl FnMut(&L::Item<'w>) -> K,
1583+
mut f: impl FnMut(&L::Item<'_>) -> K,
15781584
) -> QuerySortedManyIter<
15791585
'w,
15801586
's,
@@ -1606,7 +1612,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
16061612
/// called on [`QueryManyIter`] before.
16071613
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
16081614
self,
1609-
mut f: impl FnMut(&L::Item<'w>) -> K,
1615+
mut f: impl FnMut(&L::Item<'_>) -> K,
16101616
) -> QuerySortedManyIter<
16111617
'w,
16121618
's,
@@ -1640,7 +1646,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
16401646
/// called on [`QueryManyIter`] before.
16411647
pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
16421648
self,
1643-
mut f: impl FnMut(&L::Item<'w>) -> K,
1649+
mut f: impl FnMut(&L::Item<'_>) -> K,
16441650
) -> QuerySortedManyIter<
16451651
'w,
16461652
's,
@@ -1669,7 +1675,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
16691675
/// called on [`QueryManyIter`] before.
16701676
fn sort_impl<L: ReadOnlyQueryData + 'w>(
16711677
self,
1672-
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
1678+
f: impl FnOnce(&mut Vec<(L::Item<'_>, NeutralOrd<Entity>)>),
16731679
) -> QuerySortedManyIter<
16741680
'w,
16751681
's,
@@ -2902,13 +2908,13 @@ mod tests {
29022908
{
29032909
let mut query = query_state
29042910
.iter_many_mut(&mut world, [id, id])
2905-
.sort_by::<&C>(Ord::cmp);
2911+
.sort_by::<&C>(|l, r| Ord::cmp(l, r));
29062912
while query.fetch_next().is_some() {}
29072913
}
29082914
{
29092915
let mut query = query_state
29102916
.iter_many_mut(&mut world, [id, id])
2911-
.sort_unstable_by::<&C>(Ord::cmp);
2917+
.sort_unstable_by::<&C>(|l, r| Ord::cmp(l, r));
29122918
while query.fetch_next().is_some() {}
29132919
}
29142920
{

0 commit comments

Comments
 (0)