From b14a82fb0900bceb97267aa8153e40395c49879e Mon Sep 17 00:00:00 2001
From: Erich Gubler <erichdongubler@gmail.com>
Date: Thu, 21 Nov 2024 09:34:44 -0500
Subject: [PATCH] fix(core): validate bounds of timestamp write indices in
 compute passes

---
 CHANGELOG.md                     |  1 +
 wgpu-core/src/command/compute.rs | 14 ++++++++++++--
 wgpu-core/src/command/mod.rs     |  2 ++
 wgpu-core/src/command/query.rs   |  2 +-
 wgpu-core/src/command/render.rs  | 12 ++++++++++--
 5 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 574fb603eb5..ed4d0721a32 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -133,6 +133,7 @@ By @ErichDonGubler in [#6456](https://github.com/gfx-rs/wgpu/pull/6456), [#6148]
 - Replace potentially unsound usage of `PreHashedMap` with `FastHashMap`. By @jamienicol in [#6541](https://github.com/gfx-rs/wgpu/pull/6541).
 - Add missing validation for timestamp writes in compute and render passes (by @ErichDonGubler in [#6578](https://github.com/gfx-rs/wgpu/pull/6578)):
   - Check the status of the `TIMESTAMP_QUERY` feature before other validation.
+  - Check that indices are in-bounds for the query set.
 
 #### Naga
 
diff --git a/wgpu-core/src/command/compute.rs b/wgpu-core/src/command/compute.rs
index 630294e8ab6..db9ba51d39b 100644
--- a/wgpu-core/src/command/compute.rs
+++ b/wgpu-core/src/command/compute.rs
@@ -313,8 +313,8 @@ impl Global {
         arc_desc.timestamp_writes = if let Some(tw) = desc.timestamp_writes {
             let &PassTimestampWrites {
                 query_set,
-                beginning_of_pass_write_index: _,
-                end_of_pass_write_index: _,
+                beginning_of_pass_write_index,
+                end_of_pass_write_index,
             } = tw;
 
             match cmd_buf
@@ -335,6 +335,16 @@ impl Global {
                 Err(e) => return make_err(e.into(), arc_desc),
             }
 
+            for idx in [beginning_of_pass_write_index, end_of_pass_write_index]
+                .into_iter()
+                .flatten()
+            {
+                match query_set.validate_query(SimplifiedQueryType::Timestamp, idx, None) {
+                    Ok(()) => (),
+                    Err(e) => return make_err(e.into(), arc_desc),
+                }
+            }
+
             Some(ArcPassTimestampWrites {
                 query_set,
                 beginning_of_pass_write_index: tw.beginning_of_pass_write_index,
diff --git a/wgpu-core/src/command/mod.rs b/wgpu-core/src/command/mod.rs
index 6a06d1f83df..6c7b26b7867 100644
--- a/wgpu-core/src/command/mod.rs
+++ b/wgpu-core/src/command/mod.rs
@@ -654,6 +654,8 @@ pub enum CommandEncoderError {
     InvalidResource(#[from] InvalidResourceError),
     #[error(transparent)]
     MissingFeatures(#[from] MissingFeatures),
+    #[error(transparent)]
+    TimestampWritesInvalid(#[from] QueryUseError),
 }
 
 impl Global {
diff --git a/wgpu-core/src/command/query.rs b/wgpu-core/src/command/query.rs
index d783721fb46..b2dff276eb7 100644
--- a/wgpu-core/src/command/query.rs
+++ b/wgpu-core/src/command/query.rs
@@ -160,7 +160,7 @@ pub enum ResolveError {
 }
 
 impl QuerySet {
-    fn validate_query(
+    pub(crate) fn validate_query(
         self: &Arc<Self>,
         query_type: SimplifiedQueryType,
         query_index: u32,
diff --git a/wgpu-core/src/command/render.rs b/wgpu-core/src/command/render.rs
index 49318ecc97e..f06249fc789 100644
--- a/wgpu-core/src/command/render.rs
+++ b/wgpu-core/src/command/render.rs
@@ -1,6 +1,7 @@
 use crate::binding_model::BindGroup;
 use crate::command::{
     validate_and_begin_occlusion_query, validate_and_begin_pipeline_statistics_query,
+    SimplifiedQueryType,
 };
 use crate::init_tracker::BufferInitTrackerAction;
 use crate::pipeline::RenderPipeline;
@@ -1394,8 +1395,8 @@ impl Global {
             arc_desc.timestamp_writes = if let Some(tw) = desc.timestamp_writes {
                 let &PassTimestampWrites {
                     query_set,
-                    beginning_of_pass_write_index: _,
-                    end_of_pass_write_index: _,
+                    beginning_of_pass_write_index,
+                    end_of_pass_write_index,
                 } = tw;
 
                 let query_set = query_sets.get(query_set).get()?;
@@ -1404,6 +1405,13 @@ impl Global {
 
                 query_set.same_device(device)?;
 
+                for idx in [beginning_of_pass_write_index, end_of_pass_write_index]
+                    .into_iter()
+                    .flatten()
+                {
+                    query_set.validate_query(SimplifiedQueryType::Timestamp, idx, None)?;
+                }
+
                 Some(ArcPassTimestampWrites {
                     query_set,
                     beginning_of_pass_write_index: tw.beginning_of_pass_write_index,