From a35a151f472b318b6b87897b90d2b26d3413f19d Mon Sep 17 00:00:00 2001
From: Stepan Koltsov <stepan.koltsov@gmail.com>
Date: Sat, 6 Jan 2024 17:31:01 +0000
Subject: [PATCH] Fix integer overflow in BlobVec::reserve_exact (#11234)

# Objective

When `BlobVec::reserve` is called with an argument causing capacity
overflow, in release build capacity overflow is ignored, and capacity is
decreased.

I'm not sure it is possible to exploit this issue using public API of
`bevy_ecs`, but better fix it anyway.

## Solution

Check for capacity overflow.
---
 crates/bevy_ecs/src/storage/blob_vec.rs | 28 ++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/crates/bevy_ecs/src/storage/blob_vec.rs b/crates/bevy_ecs/src/storage/blob_vec.rs
index bcbf0c6bae..e8cf7c074b 100644
--- a/crates/bevy_ecs/src/storage/blob_vec.rs
+++ b/crates/bevy_ecs/src/storage/blob_vec.rs
@@ -132,10 +132,10 @@ impl BlobVec {
     /// For ZST it panics unconditionally because ZST `BlobVec` capacity
     /// is initialized to `usize::MAX` and always stays that way.
     fn grow_exact(&mut self, increment: NonZeroUsize) {
-        // If we got here with ZST, we requested total capacity > usize::MAX.
-        assert!(self.item_layout.size() != 0, "ZST capacity overflow");
-
-        let new_capacity = self.capacity + increment.get();
+        let new_capacity = self
+            .capacity
+            .checked_add(increment.get())
+            .expect("capacity overflow");
         let new_layout =
             array_layout(&self.item_layout, new_capacity).expect("array layout should be valid");
         let new_data = if self.capacity == 0 {
@@ -621,7 +621,7 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(expected = "ZST capacity overflow")]
+    #[should_panic(expected = "capacity overflow")]
     fn blob_vec_zst_size_overflow() {
         // SAFETY: no drop is correct drop for `()`.
         let mut blob_vec = unsafe { BlobVec::new(Layout::new::<()>(), None, 0) };
@@ -644,6 +644,24 @@ mod tests {
         }
     }
 
+    #[test]
+    #[should_panic(expected = "capacity overflow")]
+    fn blob_vec_capacity_overflow() {
+        // SAFETY: no drop is correct drop for `u32`.
+        let mut blob_vec = unsafe { BlobVec::new(Layout::new::<u32>(), None, 0) };
+
+        assert_eq!(0, blob_vec.capacity(), "Self-check");
+
+        OwningPtr::make(17u32, |ptr| {
+            // SAFETY: we push the value of correct type.
+            unsafe {
+                blob_vec.push(ptr);
+            }
+        });
+
+        blob_vec.reserve_exact(usize::MAX);
+    }
+
     #[test]
     fn aligned_zst() {
         // NOTE: This test is explicitly for uncovering potential UB with miri.