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

2024-10-29 21:41:36 +00:00 · 2024-10-21 18:12:28 +02:00 · 2024-10-21 18:12:28 +02:00 · 207747cab5
commit 207747cab5
parent 3199a3a6b0
17 changed files with 104 additions and 27 deletions
--- 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,
        }
    }
 }
--- 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, "]")?;
                }
            }
--- 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
 	),
 )
--- 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),
--- 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
 	),
 )
--- 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),
--- 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),
--- 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,
--- 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
 	),
 )
--- 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
 	),
 )
--- 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
 	),
 )
--- 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
 	),
 )
--- 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);
 }
--- 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();
--- 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
--- 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")
            }
--- 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 {}