fix(core): Use more checked arithmetic

CHANGELOG.md

@@ -67,6 +67,7 @@ Bottom level categories:
 #### General
 
 - `Features::CLIP_DISTANCE`, `naga::Capabilities::CLIP_DISTANCE`, and `naga::BuiltIn::ClipDistance` have been renamed to `CLIP_DISTANCES` and `ClipDistances` (viz., pluralized) as appropriate, to match the WebGPU spec. By @ErichDonGubler in [#9267](https://github.com/gfx-rs/wgpu/pull/9267).
+- Added new `InvalidWorkgroupSizeError`, which is now used by `DrawError::InvalidGroupSize` and `StageError::InvalidWorkgroupSize`. By @andyleiserson in [#9357](https://github.com/gfx-rs/wgpu/pull/9357).
 
 #### Validation
 

wgpu-core/src/binding_model.rs

@@ -5,7 +5,7 @@ use alloc::{
     sync::{Arc, Weak},
     vec::Vec,
 };
-use core::{fmt, mem::ManuallyDrop, ops::Range};
+use core::{fmt, mem::ManuallyDrop, num::Saturating, ops::Range};
 
 use arrayvec::ArrayVec;
 use thiserror::Error;
@@ -367,9 +367,9 @@ impl BindingTypeMaxCountErrorKind {
 
 #[derive(Debug, Default)]
 pub(crate) struct PerStageBindingTypeCounter {
-    vertex: u32,
-    fragment: u32,
-    compute: u32,
+    vertex: Saturating<u32>,
+    fragment: Saturating<u32>,
+    compute: Saturating<u32>,
 }
 
 impl PerStageBindingTypeCounter {
@@ -397,7 +397,7 @@ impl PerStageBindingTypeCounter {
         if max_value == self.compute {
             stage |= wgt::ShaderStages::COMPUTE
         }
-        (BindingZone::Stage(stage), max_value)
+        (BindingZone::Stage(stage), max_value.0)
     }
 
     pub(crate) fn merge(&mut self, other: &Self) {

wgpu-core/src/command/bundle.rs

@@ -119,7 +119,10 @@ use crate::{
     resource_log,
     snatch::SnatchGuard,
     track::RenderBundleScope,
-    validation::{check_color_attachment_count, validate_color_attachment_bytes_per_sample},
+    validation::{
+        check_color_attachment_count, check_workgroup_sizes,
+        validate_color_attachment_bytes_per_sample,
+    },
     Label, LabelHelpers,
 };
 
@@ -848,21 +851,20 @@ fn draw_mesh_tasks(
 ) -> Result<(), RenderBundleErrorInner> {
     state.is_ready(DrawCommandFamily::DrawMeshTasks)?;
 
-    let groups_size_limit = state.device.limits.max_task_mesh_workgroups_per_dimension;
-    let max_groups = state.device.limits.max_task_mesh_workgroup_total_count;
-    if group_count_x > groups_size_limit
-        || group_count_y > groups_size_limit
-        || group_count_z > groups_size_limit
-        || group_count_x * group_count_y * group_count_z > max_groups
-    {
-        return Err(RenderBundleErrorInner::Draw(DrawError::InvalidGroupSize {
-            current: [group_count_x, group_count_y, group_count_z],
-            limit: groups_size_limit,
-            max_total: max_groups,
-        }));
-    }
-
-    if group_count_x > 0 && group_count_y > 0 && group_count_z > 0 {
+    let total_count = check_workgroup_sizes(
+        &[group_count_x, group_count_y, group_count_z],
+        &[
+            state.device.limits.max_task_mesh_workgroups_per_dimension,
+            state.device.limits.max_task_mesh_workgroups_per_dimension,
+            state.device.limits.max_task_mesh_workgroups_per_dimension,
+        ],
+        "max_task_mesh_workgroups_per_dimension",
+        state.device.limits.max_task_mesh_workgroup_total_count,
+        "max_task_mesh_workgroup_total_count",
+    )
+    .map_err(|err| RenderBundleErrorInner::Draw(err.into()))?;
+
+    if total_count > 0 {
         state.flush_bindings();
         state.commands.push(ArcRenderCommand::DrawMeshTasks {
             group_count_x,
@@ -895,6 +897,7 @@ fn multi_draw_indirect(
     let vertex_limits = super::VertexLimits::new(state.vertex_buffer_sizes(), &pipeline.steps);
 
     let stride = super::get_src_stride_of_indirect_args(family);
+    assert!(offset <= wgt::BufferAddress::MAX - stride);
     state
         .buffer_memory_init_actions
         .extend(buffer.initialization_status.read().create_action(

wgpu-core/src/command/compute.rs

@@ -885,11 +885,11 @@ fn dispatch_indirect(
         return Err(ComputePassErrorInner::UnalignedIndirectBufferOffset(offset));
     }
 
-    let end_offset = offset + size_of::<wgt::DispatchIndirectArgs>() as u64;
-    if end_offset > buffer.size {
+    let args_size = size_of::<wgt::DispatchIndirectArgs>() as u64;
+    if buffer.size < args_size || buffer.size - args_size < offset {
         return Err(ComputePassErrorInner::IndirectBufferOverrun {
             offset,
-            end_offset,
+            end_offset: offset + args_size,
             buffer_size: buffer.size,
         });
     }

wgpu-core/src/command/draw.rs

@@ -6,6 +6,7 @@ use wgt::error::{ErrorType, WebGpuError};
 
 use super::bind::BinderError;
 use crate::command::pass;
+use crate::validation::InvalidWorkgroupSizeError;
 use crate::{
     binding_model::{BindingError, ImmediateUploadError, LateMinBufferBindingSizeMismatch},
     resource::{
@@ -60,14 +61,8 @@ pub enum DrawError {
         if *wanted_mesh_pipeline {"standard"} else {"mesh shader"},
     )]
     WrongPipelineType { wanted_mesh_pipeline: bool },
-    #[error(
-        "Each current draw group size dimension ({current:?}) must be less or equal to {limit}, and the product must be less or equal to {max_total}"
-    )]
-    InvalidGroupSize {
-        current: [u32; 3],
-        limit: u32,
-        max_total: u32,
-    },
+    #[error(transparent)]
+    InvalidGroupSize(#[from] InvalidWorkgroupSizeError),
     #[error(
         "Mesh shader calls in multiview render passes require enabling the `EXPERIMENTAL_MESH_SHADER_MULTIVIEW` feature, and the highest bit ({highest_view_index}) in the multiview mask must be <= `Limits::max_multiview_view_count` ({max_multiviews})"
     )]

wgpu-core/src/command/ray_tracing.rs

@@ -721,9 +721,10 @@ fn iter_blas<'snatch_guard: 'buffers, 'buffers>(
                         }) {
                             input_barriers.push(barrier);
                         }
-                        if vertex_buffer.size
-                            < (mesh.size.vertex_count + mesh.first_vertex) as u64
-                                * mesh.vertex_stride
+                        if u64::from(mesh.size.vertex_count)
+                            .checked_add(u64::from(mesh.first_vertex))
+                            .and_then(|end_index| end_index.checked_mul(mesh.vertex_stride))
+                            .is_none_or(|end| vertex_buffer.size < end)
                         {
                             return Err(BuildAccelerationStructureError::InsufficientBufferSize(
                                 vertex_buffer.error_ident(),

wgpu-core/src/command/render.rs

@@ -44,7 +44,8 @@ use crate::{
     },
     snatch::SnatchGuard,
     track::{ResourceUsageCompatibilityError, Tracker, UsageScope},
-    validation, Label,
+    validation::{self, check_workgroup_sizes},
+    Label,
 };
 
 #[cfg(feature = "serde")]
@@ -2739,21 +2740,17 @@ fn draw_mesh_tasks(
         .device
         .limits
         .max_task_mesh_workgroup_total_count;
-    if group_count_x > groups_size_limit
-        || group_count_y > groups_size_limit
-        || group_count_z > groups_size_limit
-        || group_count_x * group_count_y * group_count_z > max_groups
-    {
-        return Err(DrawError::InvalidGroupSize {
-            current: [group_count_x, group_count_y, group_count_z],
-            limit: groups_size_limit,
-            max_total: max_groups,
-        }
-        .into());
-    }
+    let total_count = check_workgroup_sizes(
+        &[group_count_x, group_count_y, group_count_z],
+        &[groups_size_limit, groups_size_limit, groups_size_limit],
+        "max_task_mesh_workgroups_per_dimension",
+        max_groups,
+        "max_task_mesh_workgroup_total_count",
+    )
+    .map_err(|err| RenderPassErrorInner::Draw(err.into()))?;
 
     unsafe {
-        if group_count_x > 0 && group_count_y > 0 && group_count_z > 0 {
+        if total_count > 0 {
             state.pass.base.raw_encoder.draw_mesh_tasks(
                 group_count_x,
                 group_count_y,

wgpu-core/src/device/queue.rs

@@ -868,6 +868,7 @@ impl Queue {
             self.device.alignments.buffer_copy_pitch.get() as u32,
             block_size,
         );
+        assert!(u32::MAX - bytes_in_last_row >= bytes_per_row_alignment);
         let stage_bytes_per_row = wgt::math::align_to(bytes_in_last_row, bytes_per_row_alignment);
 
         // Platform validation requires that the staging buffer always be
@@ -883,17 +884,21 @@ impl Queue {
         } else {
             profiling::scope!("copy chunked");
             // Copy row by row into the optimal alignment.
-            let block_rows_in_copy =
-                (size.depth_or_array_layers - 1) * rows_per_image + height_in_blocks;
-            let stage_size =
-                wgt::BufferSize::new(stage_bytes_per_row as u64 * block_rows_in_copy as u64)
-                    .unwrap();
+            let block_rows_in_copy = u64::from(size.depth_or_array_layers - 1)
+                * u64::from(rows_per_image)
+                + u64::from(height_in_blocks);
+            // The copy size was validated against the source buffer, however,
+            // `stage_bytes_per_row` can differ, so let's be paranoid.
+            let stage_size = u64::from(stage_bytes_per_row)
+                .checked_mul(block_rows_in_copy)
+                .and_then(wgt::BufferSize::new)
+                .unwrap();
             let mut staging_buffer = StagingBuffer::new(&self.device, stage_size)?;
-            for layer in 0..size.depth_or_array_layers {
-                let rows_offset = layer * rows_per_image;
-                for row in rows_offset..rows_offset + height_in_blocks {
-                    let src_offset = data_layout.offset as u32 + row * bytes_per_row;
-                    let dst_offset = row * stage_bytes_per_row;
+            for layer in 0..u64::from(size.depth_or_array_layers) {
+                let rows_offset = layer * u64::from(rows_per_image);
+                for row in rows_offset..rows_offset + u64::from(height_in_blocks) {
+                    let src_offset = data_layout.offset + row * u64::from(bytes_per_row);
+                    let dst_offset = row * u64::from(stage_bytes_per_row);
                     unsafe {
                         staging_buffer.write_with_offset(
                             data,