[hlsl-out] add support for restricting indexing to avoid OOB accesses

gfx-rs · Oct 23, 2024 · 207747c · 207747c
1 parent 3199a3a
commit 207747c
Show file tree

Hide file tree

Showing 17 changed files with 104 additions and 27 deletions.
diff --git a/naga/src/back/hlsl/mod.rs b/naga/src/back/hlsl/mod.rs
@@ -207,6 +207,8 @@ pub struct Options {
     pub push_constants_target: Option<BindTarget>,
     /// Should workgroup variables be zero initialized (by polyfilling)?
     pub zero_initialize_workgroup_memory: bool,
+    /// Should we restrict indexing of vectors, matrices and arrays?
+    pub restrict_indexing: bool,
 }
 
 impl Default for Options {
@@ -218,6 +220,7 @@ impl Default for Options {
             special_constants_binding: None,
             push_constants_target: None,
             zero_initialize_workgroup_memory: true,
+            restrict_indexing: true,
         }
     }
 }

diff --git a/naga/src/back/hlsl/writer.rs b/naga/src/back/hlsl/writer.rs
@@ -8,7 +8,7 @@ use super::{
 };
 use crate::{
     back::{self, Baked},
-    proc::{self, ExpressionKindTracker, NameKey},
+    proc::{self, index, ExpressionKindTracker, NameKey},
     valid, Handle, Module, Scalar, ScalarKind, ShaderStage, TypeInner,
 };
 use std::{fmt, mem};
@@ -2587,24 +2587,66 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
                     let resolved = func_ctx.resolve_type(base, &module.types);
 
-                    let non_uniform_qualifier = match *resolved {
+                    let (indexing_binding_array, non_uniform_qualifier) = match *resolved {
                         TypeInner::BindingArray { .. } => {
                             let uniformity = &func_ctx.info[index].uniformity;
 
-                            uniformity.non_uniform_result.is_some()
+                            (true, uniformity.non_uniform_result.is_some())
                         }
-                        _ => false,
+                        _ => (false, false),
                     };
 
                     self.write_expr(module, base, func_ctx)?;
                     write!(self.out, "[")?;
-                    if non_uniform_qualifier {
-                        write!(self.out, "NonUniformResourceIndex(")?;
-                    }
-                    self.write_expr(module, index, func_ctx)?;
-                    if non_uniform_qualifier {
+
+                    let needs_bound_check = self.options.restrict_indexing
+                        && !indexing_binding_array
+                        && match resolved.pointer_space() {
+                            Some(
+                                crate::AddressSpace::Function
+                                | crate::AddressSpace::Private
+                                | crate::AddressSpace::WorkGroup
+                                | crate::AddressSpace::PushConstant,
+                            )
+                            | None => true,
+                            Some(crate::AddressSpace::Uniform) => false, // TODO: needs checks for dynamic uniform buffers, see https://github.com/gfx-rs/wgpu/issues/4483
+                            Some(
+                                crate::AddressSpace::Handle | crate::AddressSpace::Storage { .. },
+                            ) => unreachable!(),
+                        };
+                    // Decide whether this index needs to be clamped to fall within range.
+                    let restriction_needed = if needs_bound_check {
+                        index::access_needs_check(
+                            base,
+                            index::GuardedIndex::Expression(index),
+                            module,
+                            func_ctx.expressions,
+                            func_ctx.info,
+                        )
+                    } else {
+                        None
+                    };
+                    if let Some(limit) = restriction_needed {
+                        write!(self.out, "min(uint(")?;
+                        self.write_expr(module, index, func_ctx)?;
+                        write!(self.out, "), ")?;
+                        match limit {
+                            index::IndexableLength::Known(limit) => {
+                                write!(self.out, "{}u", limit - 1)?;
+                            }
+                            index::IndexableLength::Dynamic => unreachable!(),
+                        }
                         write!(self.out, ")")?;
+                    } else {
+                        if non_uniform_qualifier {
+                            write!(self.out, "NonUniformResourceIndex(")?;
+                        }
+                        self.write_expr(module, index, func_ctx)?;
+                        if non_uniform_qualifier {
+                            write!(self.out, ")")?;
+                        }
                     }
+
                     write!(self.out, "]")?;
                 }
             }

diff --git a/naga/tests/in/atomicCompareExchange-int64.param.ron b/naga/tests/in/atomicCompareExchange-int64.param.ron
@@ -11,5 +11,6 @@
 		special_constants_binding: Some((space: 1, register: 0)),
 		push_constants_target: Some((space: 0, register: 0)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 )
diff --git a/naga/tests/in/atomicOps-int64-min-max.param.ron b/naga/tests/in/atomicOps-int64-min-max.param.ron
@@ -11,6 +11,7 @@
 		special_constants_binding: Some((space: 1, register: 0)),
 		push_constants_target: Some((space: 0, register: 0)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 	msl: (
 		lang_version: (2, 4),

diff --git a/naga/tests/in/atomicOps-int64.param.ron b/naga/tests/in/atomicOps-int64.param.ron
@@ -11,5 +11,6 @@
 		special_constants_binding: Some((space: 1, register: 0)),
 		push_constants_target: Some((space: 0, register: 0)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 )
diff --git a/naga/tests/in/binding-arrays.param.ron b/naga/tests/in/binding-arrays.param.ron
@@ -16,6 +16,7 @@
 		fake_missing_bindings: true,
 		special_constants_binding: None,
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 	msl: (
 		lang_version: (2, 0),

diff --git a/naga/tests/in/int64.param.ron b/naga/tests/in/int64.param.ron
@@ -11,6 +11,7 @@
 		special_constants_binding: Some((space: 1, register: 0)),
 		push_constants_target: Some((space: 0, register: 0)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 	msl: (
 		lang_version: (2, 3),

diff --git a/naga/tests/in/interface.param.ron b/naga/tests/in/interface.param.ron
@@ -13,6 +13,7 @@
 		fake_missing_bindings: false,
 		special_constants_binding: Some((space: 1, register: 0)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 	wgsl: (
 		explicit_types: true,

diff --git a/naga/tests/in/push-constants.param.ron b/naga/tests/in/push-constants.param.ron
@@ -16,5 +16,6 @@
 		special_constants_binding: Some((space: 1, register: 0)),
 		push_constants_target: Some((space: 0, register: 0)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 )
diff --git a/naga/tests/in/skybox.param.ron b/naga/tests/in/skybox.param.ron
@@ -59,5 +59,6 @@
 		fake_missing_bindings: false,
 		special_constants_binding: Some((space: 0, register: 1)),
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 )
diff --git a/naga/tests/in/spv/subgroup-operations-s.param.ron b/naga/tests/in/spv/subgroup-operations-s.param.ron
@@ -23,5 +23,6 @@
 		fake_missing_bindings: true,
 		special_constants_binding: None,
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 )
diff --git a/naga/tests/in/subgroup-operations.param.ron b/naga/tests/in/subgroup-operations.param.ron
@@ -23,5 +23,6 @@
 		fake_missing_bindings: true,
 		special_constants_binding: None,
 		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
 	),
 )
diff --git a/naga/tests/out/hlsl/access.hlsl b/naga/tests/out/hlsl/access.hlsl
@@ -196,7 +196,7 @@ void test_matrix_within_array_within_struct_accesses()
     __set_col_of_mat4x2(t_1.am[0], _e77, (90.0).xx);
     t_1.am[0]._0.y = 10.0;
     int _e89 = idx_1;
-    t_1.am[0]._0[_e89] = 20.0;
+    t_1.am[0]._0[min(uint(_e89), 1u)] = 20.0;
     int _e94 = idx_1;
     __set_el_of_mat4x2(t_1.am[0], _e94, 1, 30.0);
     int _e100 = idx_1;
@@ -298,8 +298,8 @@ float4 foo_vert(uint vi : SV_VertexID) : SV_Position
     int2 c = asint(qux.Load2(0));
     const float _e33 = read_from_private(foo);
     c2_ = Constructarray5_int_(a_1, int(b), 3, 4, 5);
-    c2_[(vi + 1u)] = 42;
-    int value = c2_[vi];
+    c2_[min(uint((vi + 1u)), 4u)] = 42;
+    int value = c2_[min(uint(vi), 4u)];
     const float _e47 = test_arr_as_arg(ZeroValuearray5_array10_float__());
     return float4(mul(float4((value).xxxx), _matrix), 2.0);
 }

diff --git a/naga/tests/out/hlsl/workgroup-uniform-load.hlsl b/naga/tests/out/hlsl/workgroup-uniform-load.hlsl
@@ -10,7 +10,7 @@ void test_workgroupUniformLoad(uint3 workgroup_id : SV_GroupID, uint3 __local_in
     }
     GroupMemoryBarrierWithGroupSync();
     GroupMemoryBarrierWithGroupSync();
-    int _e4 = arr_i32_[workgroup_id.x];
+    int _e4 = arr_i32_[min(uint(workgroup_id.x), 127u)];
     GroupMemoryBarrierWithGroupSync();
     if ((_e4 > 10)) {
         GroupMemoryBarrierWithGroupSync();

diff --git a/tests/tests/oob_indexing.rs b/tests/tests/oob_indexing.rs
@@ -1,5 +1,5 @@
 use wgpu_test::{gpu_test, FailureCase, GpuTestConfiguration, TestParameters, TestingContext};
-use wgt::Backends;
+use wgt::{Backend, Backends};
 
 /// Tests that writing and reading to the max length of a container (vec, mat, array)
 /// in the workgroup, private and function address spaces + let declarations
@@ -10,7 +10,7 @@ static RESTRICT_WORKGROUP_PRIVATE_FUNCTION_LET: GpuTestConfiguration = GpuTestCo
         TestParameters::default()
             .downlevel_flags(wgpu::DownlevelFlags::COMPUTE_SHADERS)
             .limits(wgpu::Limits::downlevel_defaults())
-            .skip(FailureCase::backend(Backends::DX12 | Backends::GL)),
+            .skip(FailureCase::backend(Backends::GL)),
     )
     .run_async(|ctx| async move {
         let test_resources = TestResources::new(&ctx);
@@ -51,7 +51,11 @@ static RESTRICT_WORKGROUP_PRIVATE_FUNCTION_LET: GpuTestConfiguration = GpuTestCo
         drop(view);
         test_resources.readback_buffer.unmap();
 
-        assert_eq!([1; 12], current_res);
+        if ctx.adapter_info.backend == Backend::Dx12 {
+            assert_eq!([1, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0], current_res);
+        } else {
+            assert_eq!([1; 12], current_res);
+        }
     });
 
 struct TestResources {
@@ -64,7 +68,16 @@ struct TestResources {
 
 impl TestResources {
     fn new(ctx: &TestingContext) -> Self {
-        const SHADER_SRC: &str = "
+        // FXC doesn't support dynamically indexing and writing to vectors and matrices, it errors with:
+        // error X3500: array reference cannot be used as an l-value; not natively addressable
+        // see also: https://github.com/gfx-rs/wgpu/issues/4460
+        let opt = if ctx.adapter_info.backend == Backend::Dx12 {
+            "//"
+        } else {
+            ""
+        };
+        let shader_src = format!(
+            "
             @group(0) @binding(0)
             var<storage, read_write> in: u32;
             @group(0) @binding(1)
@@ -79,7 +92,7 @@ impl TestResources {
             var<private> private_mat: mat3x3f;
 
             @compute @workgroup_size(1)
-            fn main() {
+            fn main() {{
                 let i = in;
 
                 var var_array = array<u32, 3>();
@@ -95,8 +108,8 @@ impl TestResources {
 
                 var var_vec = vec3u();
                 wg_vec[i] = 1u;
-                private_vec[i] = 1u;
-                var_vec[i] = 1u;
+                {opt} private_vec[i] = 1u;
+                {opt} var_vec[i] = 1u;
                 let let_vec = var_vec;
 
                 out[4] = wg_vec[i];
@@ -106,22 +119,23 @@ impl TestResources {
 
                 var var_mat = mat3x3f();
                 wg_mat[i][0] = 1f;
-                private_mat[i][0] = 1f;
-                var_mat[i][0] = 1f;
+                {opt} private_mat[i][0] = 1f;
+                {opt} var_mat[i][0] = 1f;
                 let let_mat = var_mat;
 
                 out[8] = u32(wg_mat[i][0]);
                 out[9] = u32(private_mat[i][0]);
                 out[10] = u32(var_mat[i][0]);
                 out[11] = u32(let_mat[i][0]);
-            }
-        ";
+            }}
+        "
+        );
 
         let module = ctx
             .device
             .create_shader_module(wgpu::ShaderModuleDescriptor {
                 label: None,
-                source: wgpu::ShaderSource::Wgsl(SHADER_SRC.into()),
+                source: wgpu::ShaderSource::Wgsl(shader_src.into()),
             });
 
         let bgl = ctx

diff --git a/wgpu-hal/src/dx12/device.rs b/wgpu-hal/src/dx12/device.rs
@@ -271,11 +271,13 @@ impl super::Device {
         .map_err(|e| crate::PipelineError::PipelineConstants(stage_bit, format!("HLSL: {e:?}")))?;
 
         let needs_temp_options = stage.zero_initialize_workgroup_memory
-            != layout.naga_options.zero_initialize_workgroup_memory;
+            != layout.naga_options.zero_initialize_workgroup_memory
+            || stage.module.runtime_checks != layout.naga_options.restrict_indexing;
         let mut temp_options;
         let naga_options = if needs_temp_options {
             temp_options = layout.naga_options.clone();
             temp_options.zero_initialize_workgroup_memory = stage.zero_initialize_workgroup_memory;
+            temp_options.restrict_indexing = stage.module.runtime_checks;
             &temp_options
         } else {
             &layout.naga_options
@@ -1223,6 +1225,7 @@ impl crate::Device for super::Device {
                 special_constants_binding,
                 push_constants_target,
                 zero_initialize_workgroup_memory: true,
+                restrict_indexing: true,
             },
         })
     }
@@ -1438,7 +1441,11 @@ impl crate::Device for super::Device {
 
         let raw_name = desc.label.and_then(|label| ffi::CString::new(label).ok());
         match shader {
-            crate::ShaderInput::Naga(naga) => Ok(super::ShaderModule { naga, raw_name }),
+            crate::ShaderInput::Naga(naga) => Ok(super::ShaderModule {
+                naga,
+                raw_name,
+                runtime_checks: desc.runtime_checks,
+            }),
             crate::ShaderInput::SpirV(_) => {
                 panic!("SPIRV_SHADER_PASSTHROUGH is not enabled for this backend")
             }

diff --git a/wgpu-hal/src/dx12/mod.rs b/wgpu-hal/src/dx12/mod.rs
@@ -953,6 +953,7 @@ impl crate::DynPipelineLayout for PipelineLayout {}
 pub struct ShaderModule {
     naga: crate::NagaShader,
     raw_name: Option<ffi::CString>,
+    runtime_checks: bool,
 }
 
 impl crate::DynShaderModule for ShaderModule {}