Fix unsound lifetimes in Query::join and Query::join_filtered (#17972)

# Objective

Fix unsound lifetimes in `Query::join` and `Query::join_filtered`.  

The joined query allowed access from either input query, but it only
took the `'world` lifetime from `self`, not from `other`. This meant
that after the borrow of `other` ended, the joined query would unsoundly
alias `other`.

## Solution

Change the lifetimes on `join` and `join_filtered` to require mutable
borrows of the *same* lifetime for the input queries. This ensures both
input queries are borrowed for the full lifetime of the joined query.

Change `join_inner` to take `other` by value instead of reference so
that the returned query is still usable without needing to borrow from a
local variable.

## Testing

Added a compile-fail test.
This commit is contained in:
Chris Russell 2025-03-10 17:30:34 -04:00 committed by GitHub
parent 387f1c593b
commit 6df711ce7f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 68 additions and 21 deletions

View File

@ -1,24 +1,46 @@
use bevy_ecs::prelude::*;
use bevy_ecs::system::SystemState;
use bevy_ecs::system::{QueryLens, SystemState};
#[derive(Component, Eq, PartialEq, Debug)]
struct Foo(u32);
#[derive(Component, Eq, PartialEq, Debug)]
struct Bar(u32);
fn main() {
let mut world = World::default();
let e = world.spawn(Foo(10_u32)).id();
let e = world.spawn((Foo(10_u32), Bar(10_u32))).id();
let mut system_state = SystemState::<Query<&mut Foo>>::new(&mut world);
let mut system_state = SystemState::<(Query<&mut Foo>, Query<&mut Bar>)>::new(&mut world);
{
let mut query = system_state.get_mut(&mut world);
let mut lens = query.as_query_lens();
let (mut foo_query, mut bar_query) = system_state.get_mut(&mut world);
dbg!("hi");
{
let mut lens = foo_query.as_query_lens();
let mut data: Mut<Foo> = lens.query().get_inner(e).unwrap();
let mut data2: Mut<Foo> = lens.query().get_inner(e).unwrap();
//~^ E0499
assert_eq!(&mut *data, &mut *data2); // oops UB
}
{
let mut join: QueryLens<(&mut Foo, &mut Bar)> = foo_query.join(&mut bar_query);
let mut query = join.query();
let (_, mut data) = query.single_mut().unwrap();
let mut data2 = bar_query.single_mut().unwrap();
//~^ E0499
assert_eq!(&mut *data, &mut *data2); // oops UB
}
{
let mut join: QueryLens<(&mut Foo, &mut Bar)> =
foo_query.join_inner(bar_query.reborrow());
let mut query = join.query();
let (_, mut data) = query.single_mut().unwrap();
let mut data2 = bar_query.single_mut().unwrap();
//~^ E0499
assert_eq!(&mut *data, &mut *data2); // oops UB
}
dbg!("bye");
}
}

View File

@ -1,14 +1,38 @@
error[E0499]: cannot borrow `lens` as mutable more than once at a time
--> tests/ui/query_lens_lifetime_safety.rs:18:39
--> tests/ui/query_lens_lifetime_safety.rs:21:39
|
17 | let mut data: Mut<Foo> = lens.query().get_inner(e).unwrap();
20 | let mut data: Mut<Foo> = lens.query().get_inner(e).unwrap();
| ---- first mutable borrow occurs here
18 | let mut data2: Mut<Foo> = lens.query().get_inner(e).unwrap();
21 | let mut data2: Mut<Foo> = lens.query().get_inner(e).unwrap();
| ^^^^ second mutable borrow occurs here
19 |
20 | assert_eq!(&mut *data, &mut *data2); // oops UB
22 |
23 | assert_eq!(&mut *data, &mut *data2); // oops UB
| ---- first borrow later used here
error: aborting due to 1 previous error
error[E0499]: cannot borrow `bar_query` as mutable more than once at a time
--> tests/ui/query_lens_lifetime_safety.rs:30:29
|
27 | let mut join: QueryLens<(&mut Foo, &mut Bar)> = foo_query.join(&mut bar_query);
| -------------- first mutable borrow occurs here
...
30 | let mut data2 = bar_query.single_mut().unwrap();
| ^^^^^^^^^ second mutable borrow occurs here
31 |
32 | assert_eq!(&mut *data, &mut *data2); // oops UB
| ---- first borrow later used here
error[E0499]: cannot borrow `bar_query` as mutable more than once at a time
--> tests/ui/query_lens_lifetime_safety.rs:40:29
|
37 | foo_query.join_inner(bar_query.reborrow());
| --------- first mutable borrow occurs here
...
40 | let mut data2 = bar_query.single_mut().unwrap();
| ^^^^^^^^^ second mutable borrow occurs here
41 |
42 | assert_eq!(&mut *data, &mut *data2); // oops UB
| ---- first borrow later used here
error: aborting due to 3 previous errors
For more information about this error, try `rustc --explain E0499`.

View File

@ -2234,10 +2234,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
///
/// Like `transmute_lens` the query terms can be changed with some restrictions.
/// See [`Self::transmute_lens`] for more details.
pub fn join<OtherD: QueryData, NewD: QueryData>(
&mut self,
other: &mut Query<OtherD>,
) -> QueryLens<'_, NewD> {
pub fn join<'a, OtherD: QueryData, NewD: QueryData>(
&'a mut self,
other: &'a mut Query<OtherD>,
) -> QueryLens<'a, NewD> {
self.join_filtered(other)
}
@ -2263,7 +2263,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
/// - [`join`](Self::join) to join using a mutable borrow of the [`Query`].
pub fn join_inner<OtherD: QueryData, NewD: QueryData>(
self,
other: &mut Query<OtherD>,
other: Query<'w, '_, OtherD>,
) -> QueryLens<'w, NewD> {
self.join_filtered_inner(other)
}
@ -2276,15 +2276,16 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
/// terms like `Added` and `Changed` will only be respected if they are in
/// the type signature.
pub fn join_filtered<
'a,
OtherD: QueryData,
OtherF: QueryFilter,
NewD: QueryData,
NewF: QueryFilter,
>(
&mut self,
other: &mut Query<OtherD, OtherF>,
) -> QueryLens<'_, NewD, NewF> {
self.reborrow().join_filtered_inner(other)
&'a mut self,
other: &'a mut Query<OtherD, OtherF>,
) -> QueryLens<'a, NewD, NewF> {
self.reborrow().join_filtered_inner(other.reborrow())
}
/// Equivalent to [`Self::join_inner`] but also includes a [`QueryFilter`] type.
@ -2306,7 +2307,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
NewF: QueryFilter,
>(
self,
other: &mut Query<OtherD, OtherF>,
other: Query<'w, '_, OtherD, OtherF>,
) -> QueryLens<'w, NewD, NewF> {
let state = self
.state