From 58f7dac689514d64e2d2d7f34257515337c2e48c Mon Sep 17 00:00:00 2001
From: Joseph <21144246+JoJoJet@users.noreply.github.com>
Date: Sat, 2 Sep 2023 16:43:22 -0700
Subject: [PATCH] Fix unsoundness in `QueryState::is_empty` (#9463)

# Objective

`QueryState::is_empty` is unsound, as it does not validate the world. If
a mismatched world is passed in, then the query filter may cast a
component to an incorrect type, causing undefined behavior.

## Solution

Add world validation. To prevent a performance regression in `Query`
(whose world does not need to be validated), the unchecked function
`is_empty_unsafe_world_cell` has been added. This also allows us to
remove one of the last usages of the private function
`UnsafeWorldCell::unsafe_world`, which takes us a step towards being
able to remove that method entirely.
---
 crates/bevy_ecs/src/query/state.rs  | 36 +++++++++++++++++++++++++++--
 crates/bevy_ecs/src/system/query.rs | 14 ++++++-----
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs
index 33d66271e5cfb..c1214c55585db 100644
--- a/crates/bevy_ecs/src/query/state.rs
+++ b/crates/bevy_ecs/src/query/state.rs
@@ -127,12 +127,44 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
     }
 
     /// Checks if the query is empty for the given [`World`], where the last change and current tick are given.
+    ///
+    /// # Panics
+    ///
+    /// If `world` does not match the one used to call `QueryState::new` for this instance.
     #[inline]
     pub fn is_empty(&self, world: &World, last_run: Tick, this_run: Tick) -> bool {
-        // SAFETY: NopFetch does not access any members while &self ensures no one has exclusive access
+        self.validate_world(world.id());
+        // SAFETY:
+        // - We have read-only access to the entire world.
+        // - The world has been validated.
+        unsafe {
+            self.is_empty_unsafe_world_cell(
+                world.as_unsafe_world_cell_readonly(),
+                last_run,
+                this_run,
+            )
+        }
+    }
+
+    /// Checks if the query is empty for the given [`UnsafeWorldCell`].
+    ///
+    /// # Safety
+    ///
+    /// - `world` must have permission to read any components required by this instance's `F` [`WorldQuery`].
+    /// - `world` must match the one used to create this [`QueryState`].
+    #[inline]
+    pub(crate) unsafe fn is_empty_unsafe_world_cell(
+        &self,
+        world: UnsafeWorldCell,
+        last_run: Tick,
+        this_run: Tick,
+    ) -> bool {
+        // SAFETY:
+        // - The caller ensures that `world` has permission to access any data used by the filter.
+        // - The caller ensures that the world matches.
         unsafe {
             self.as_nop()
-                .iter_unchecked_manual(world.as_unsafe_world_cell_readonly(), last_run, this_run)
+                .iter_unchecked_manual(world, last_run, this_run)
                 .next()
                 .is_none()
         }
diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs
index 5716497667f43..9fbad21a291ce 100644
--- a/crates/bevy_ecs/src/system/query.rs
+++ b/crates/bevy_ecs/src/system/query.rs
@@ -1367,12 +1367,14 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
     /// ```
     #[inline]
     pub fn is_empty(&self) -> bool {
-        self.state.is_empty(
-            // SAFETY: `QueryState::is_empty` does not access world data.
-            unsafe { self.world.unsafe_world() },
-            self.last_run,
-            self.this_run,
-        )
+        // SAFETY:
+        // - `self.world` has permission to read any data required by the WorldQuery.
+        // - `&self` ensures that no one currently has write access.
+        // - `self.world` matches `self.state`.
+        unsafe {
+            self.state
+                .is_empty_unsafe_world_cell(self.world, self.last_run, self.this_run)
+        }
     }
 
     /// Returns `true` if the given [`Entity`] matches the query.