From 62c1812e72dc2367c8d537d8224f1d55eae2535f Mon Sep 17 00:00:00 2001
From: Chris Russell <8494645+chescock@users.noreply.github.com>
Date: Wed, 12 Feb 2025 17:41:02 -0500
Subject: [PATCH] Shorten the 'world lifetime returned from
 `QueryLens::query()`. (#17694)

# Objective

Fix unsoundness introduced by #15858. `QueryLens::query()` would hand
out a `Query` with the full `'w` lifetime, and the new `_inner` methods
would let the results outlive the `Query`. This could be used to create
aliasing mutable references, like

```rust
fn bad<'w>(mut lens: QueryLens<'w, EntityMut>, entity: Entity) {
    let one: EntityMut<'w> = lens.query().get_inner(entity).unwrap();
    let two: EntityMut<'w> = lens.query().get_inner(entity).unwrap();
    assert!(one.entity() == two.entity());
}
```

Fixes #17693

## Solution

Restrict the `'world` lifetime in the `Query` returned by
`QueryLens::query()` to `'_`, the lifetime of the borrow of the
`QueryLens`.

The model here is that `Query<'w, 's, D, F>` and `QueryLens<'w, D, F>`
have permission to access their components for the lifetime `'w`. So
going from `&'a mut QueryLens<'w>` to `Query<'w, 'a>` would borrow the
permission only for the `'a` lifetime, but incorrectly give it out for
the full `'w` lifetime.

To handle any cases where users were calling `get_inner()` or
`iter_inner()` on the `Query` and expecting the full `'w` lifetime, we
introduce a new `QueryLens::query_inner()` method. This is only valid
for `ReadOnlyQueryData`, so it may safely hand out a copy of the
permission for the full `'w` lifetime. Since `get_inner()` and
`iter_inner()` were only valid on `ReadOnlyQueryData` prior to #15858,
that should cover any uses that relied on the longer lifetime.

## Migration Guide

Users of `QueryLens::query()` who were calling `get_inner()` or
`iter_inner()` will need to replace the call with
`QueryLens::query_inner()`.
---
 .../tests/ui/query_lens_lifetime_safety.rs    | 24 +++++++++++++++++++
 .../ui/query_lens_lifetime_safety.stderr      |  5 ++++
 crates/bevy_ecs/src/system/query.rs           | 20 +++++++++++++---
 3 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.rs
 create mode 100644 crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.stderr

diff --git a/crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.rs b/crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.rs
new file mode 100644
index 0000000000..462e9dbf1f
--- /dev/null
+++ b/crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.rs
@@ -0,0 +1,24 @@
+use bevy_ecs::prelude::*;
+use bevy_ecs::system::SystemState;
+
+#[derive(Component, Eq, PartialEq, Debug)]
+struct Foo(u32);
+
+fn main() {
+    let mut world = World::default();
+    let e = world.spawn(Foo(10_u32)).id();
+
+    let mut system_state = SystemState::<Query<&mut Foo>>::new(&mut world);
+    {
+        let mut query = system_state.get_mut(&mut world);
+        let mut lens = query.as_query_lens();
+        dbg!("hi");
+        {
+            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
+        }
+        dbg!("bye");
+    }
+}
diff --git a/crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.stderr b/crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.stderr
new file mode 100644
index 0000000000..060b320295
--- /dev/null
+++ b/crates/bevy_ecs/compile_fail/tests/ui/query_lens_lifetime_safety.stderr
@@ -0,0 +1,5 @@
+error[E0499]: cannot borrow `lens` as mutable more than once at a time
+  --> tests/ui\query_lens_lifetime_safety.rs:18:39
+   |
+17 |             let mut data: Mut<Foo> = lens.query().get_inner(e).unwrap();
+   |                                      ---- first mutable borrow occurs here
diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs
index fe5662febc..63622a7020 100644
--- a/crates/bevy_ecs/src/system/query.rs
+++ b/crates/bevy_ecs/src/system/query.rs
@@ -2127,7 +2127,21 @@ pub struct QueryLens<'w, Q: QueryData, F: QueryFilter = ()> {
 
 impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> {
     /// Create a [`Query`] from the underlying [`QueryState`].
-    pub fn query(&mut self) -> Query<'w, '_, Q, F> {
+    pub fn query(&mut self) -> Query<'_, '_, Q, F> {
+        Query {
+            world: self.world,
+            state: &self.state,
+            last_run: self.last_run,
+            this_run: self.this_run,
+        }
+    }
+}
+
+impl<'w, Q: ReadOnlyQueryData, F: QueryFilter> QueryLens<'w, Q, F> {
+    /// Create a [`Query`] from the underlying [`QueryState`].
+    /// This returns results with the actual "inner" world lifetime,
+    /// so it may only be used with read-only queries to prevent mutable aliasing.
+    pub fn query_inner(&self) -> Query<'w, '_, Q, F> {
         Query {
             world: self.world,
             state: &self.state,
@@ -2138,9 +2152,9 @@ impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> {
 }
 
 impl<'w, 's, Q: QueryData, F: QueryFilter> From<&'s mut QueryLens<'w, Q, F>>
-    for Query<'w, 's, Q, F>
+    for Query<'s, 's, Q, F>
 {
-    fn from(value: &'s mut QueryLens<'w, Q, F>) -> Query<'w, 's, Q, F> {
+    fn from(value: &'s mut QueryLens<'w, Q, F>) -> Query<'s, 's, Q, F> {
         value.query()
     }
 }