From d17fce6e46026a3ffda7e44a9756d7c05259c77c Mon Sep 17 00:00:00 2001
From: Islam El-Ashi <islam.elashi@dfinity.org>
Date: Mon, 22 Jan 2024 11:00:21 +0100
Subject: [PATCH] fix: vec should panic when trying to insert element >
 `MAX_SIZE`

---
 src/base_vec.rs  |  4 ++--
 src/vec/tests.rs | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/base_vec.rs b/src/base_vec.rs
index cd286206..c99d4a70 100644
--- a/src/base_vec.rs
+++ b/src/base_vec.rs
@@ -179,7 +179,7 @@ impl<T: Storable, M: Memory> BaseVec<T, M> {
         assert!(index < self.len());
 
         let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
-        let bytes = item.to_bytes();
+        let bytes = item.to_bytes_checked();
         let data_offset = self
             .write_entry_size(offset, bytes.len() as u32)
             .expect("unreachable: cannot fail to write to pre-allocated area");
@@ -203,7 +203,7 @@ impl<T: Storable, M: Memory> BaseVec<T, M> {
     pub fn push(&self, item: &T) -> Result<(), GrowFailed> {
         let index = self.len();
         let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
-        let bytes = item.to_bytes();
+        let bytes = item.to_bytes_checked();
         let data_offset = self.write_entry_size(offset, bytes.len() as u32)?;
         safe_write(&self.memory, data_offset, bytes.borrow())?;
         // NB. We update the size only after we ensure that the data
diff --git a/src/vec/tests.rs b/src/vec/tests.rs
index de6c0de1..eb89d492 100644
--- a/src/vec/tests.rs
+++ b/src/vec/tests.rs
@@ -228,3 +228,41 @@ fn test_iter() {
     assert_eq!(sv.iter().skip(4).count(), 0);
     assert_eq!(sv.iter().skip(usize::MAX).count(), 0);
 }
+
+// A struct with a bugg implementation of storable where the max_size can
+// smaller than the serialized size.
+#[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
+struct BuggyStruct(Vec<u8>);
+impl crate::Storable for BuggyStruct {
+    fn to_bytes(&self) -> Cow<[u8]> {
+        Cow::Borrowed(&self.0)
+    }
+
+    fn from_bytes(_: Cow<[u8]>) -> Self {
+        unimplemented!();
+    }
+
+    const BOUND: Bound = Bound::Bounded {
+        max_size: 1,
+        is_fixed_size: false,
+    };
+}
+
+#[test]
+#[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
+fn push_element_bigger_than_max_size_panics() {
+    let sv = StableVec::<BuggyStruct, M>::new(M::default()).unwrap();
+    // Insert a struct where the serialized size is > `MAX_SIZE`. Should panic.
+    sv.push(&BuggyStruct(vec![1, 2, 3, 4])).unwrap();
+}
+
+#[test]
+#[should_panic(expected = "expected an element with length <= 1 bytes, but found 5")]
+fn set_element_bigger_than_max_size_panics() {
+    let sv = StableVec::<BuggyStruct, M>::new(M::default()).unwrap();
+    // Insert a struct where the serialized size is <= `MAX_SIZE`. This should succeed.
+    sv.push(&BuggyStruct(vec![1])).unwrap();
+
+    // Insert a struct where the serialized size is > `MAX_SIZE`. Should panic.
+    sv.set(0, &BuggyStruct(vec![1, 2, 3, 4, 5]));
+}