From e711a3599857231284f3b64efb2e64db5beba2dc Mon Sep 17 00:00:00 2001
From: "N.E.C." <beholdnec@gmail.com>
Date: Tue, 22 Oct 2024 06:33:05 -0700
Subject: [PATCH] Add bounds check to Buffer slice method (#6432)

* Add bounds check to Buffer slice method

* avoid unwrap / double option check by using `map_or`

---------

Co-authored-by: Andreas Reich <r_andreas2@web.de>
---
 CHANGELOG.md           |  1 +
 wgpu/src/api/buffer.rs | 56 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index db56245c1..16d160d5b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -160,6 +160,7 @@ By @bradwerth [#6216](https://github.com/gfx-rs/wgpu/pull/6216).
 - `Rg11b10Float` is renamed to `Rg11b10Ufloat`. By @sagudev in [#6108](https://github.com/gfx-rs/wgpu/pull/6108)
 - Invalidate the device when we encounter driver-induced device loss or on unexpected errors. By @teoxoy in [#6229](https://github.com/gfx-rs/wgpu/pull/6229).
 - Make Vulkan error handling more robust. By @teoxoy in [#6119](https://github.com/gfx-rs/wgpu/pull/6119).
+- Add bounds checking to Buffer slice method. By @beholdnec in [#6432](https://github.com/gfx-rs/wgpu/pull/6432).
 
 #### Internal
 
diff --git a/wgpu/src/api/buffer.rs b/wgpu/src/api/buffer.rs
index 9d490616d..fa9c7f9ec 100644
--- a/wgpu/src/api/buffer.rs
+++ b/wgpu/src/api/buffer.rs
@@ -242,6 +242,7 @@ impl Buffer {
     /// end of the buffer.
     pub fn slice<S: RangeBounds<BufferAddress>>(&self, bounds: S) -> BufferSlice<'_> {
         let (offset, size) = range_to_offset_size(bounds);
+        check_buffer_bounds(self.size, offset, size);
         BufferSlice {
             buffer: self,
             offset,
@@ -673,6 +674,31 @@ impl Drop for Buffer {
     }
 }
 
+fn check_buffer_bounds(
+    buffer_size: BufferAddress,
+    offset: BufferAddress,
+    size: Option<BufferSize>,
+) {
+    // A slice of length 0 is invalid, so the offset must not be equal to or greater than the buffer size.
+    if offset >= buffer_size {
+        panic!(
+            "slice offset {} is out of range for buffer of size {}",
+            offset, buffer_size
+        );
+    }
+
+    if let Some(size) = size {
+        // Detect integer overflow.
+        let end = offset.checked_add(size.get());
+        if end.map_or(true, |end| end > buffer_size) {
+            panic!(
+                "slice offset {} size {} is out of range for buffer of size {}",
+                offset, size, buffer_size
+            );
+        }
+    }
+}
+
 fn range_to_offset_size<S: RangeBounds<BufferAddress>>(
     bounds: S,
 ) -> (BufferAddress, Option<BufferSize>) {
@@ -690,9 +716,10 @@ fn range_to_offset_size<S: RangeBounds<BufferAddress>>(
 
     (offset, size)
 }
+
 #[cfg(test)]
 mod tests {
-    use super::{range_to_offset_size, BufferSize};
+    use super::{check_buffer_bounds, range_to_offset_size, BufferSize};
 
     #[test]
     fn range_to_offset_size_works() {
@@ -715,4 +742,31 @@ mod tests {
     fn range_to_offset_size_panics_for_unbounded_empty_range() {
         range_to_offset_size(..0);
     }
+
+    #[test]
+    #[should_panic]
+    fn check_buffer_bounds_panics_for_offset_at_size() {
+        check_buffer_bounds(100, 100, None);
+    }
+
+    #[test]
+    fn check_buffer_bounds_works_for_end_in_range() {
+        check_buffer_bounds(200, 100, BufferSize::new(50));
+        check_buffer_bounds(200, 100, BufferSize::new(100));
+        check_buffer_bounds(u64::MAX, u64::MAX - 100, BufferSize::new(100));
+        check_buffer_bounds(u64::MAX, 0, BufferSize::new(u64::MAX));
+        check_buffer_bounds(u64::MAX, 1, BufferSize::new(u64::MAX - 1));
+    }
+
+    #[test]
+    #[should_panic]
+    fn check_buffer_bounds_panics_for_end_over_size() {
+        check_buffer_bounds(200, 100, BufferSize::new(101));
+    }
+
+    #[test]
+    #[should_panic]
+    fn check_buffer_bounds_panics_for_end_wraparound() {
+        check_buffer_bounds(u64::MAX, 1, BufferSize::new(u64::MAX));
+    }
 }