From a8078e067e3a5b18eb26ae07d28ad4cf345bd9bf Mon Sep 17 00:00:00 2001
From: Mario Rugiero <mario.rugiero@lambdaclass.com>
Date: Thu, 4 May 2023 17:21:29 -0300
Subject: [PATCH] fix(security)!: avoid DoS on malicious insert to memory
 (#1099)

Co-authored-by: Pedro Fontana <fontana.pedro93@gmail.com>
---
 CHANGELOG.md                   |  6 ++++++
 src/vm/errors/memory_errors.rs |  2 ++
 src/vm/vm_memory/memory.rs     | 25 +++++++++++++++++++++++--
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b18f115316..6bdf3d7166 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -144,6 +144,12 @@
     %}
     ```
 
+* fix(security)!: avoid DoS on malicious insertion to memory [#1099](https://github.com/lambdaclass/cairo-rs/pull/1099)
+    * A program could crash the library by attempting to insert a value at an address with a big offset; fixed by trying to reserve to check for allocation failure
+    * A program could crash the program by exploiting an integer overflow when attempting to insert a value at an address with offset `usize::MAX`
+
+    BREAKING: added a new error variant `MemoryError::VecCapacityExceeded`
+
 * fix(starknet-crypto): bump version to `0.5.0` [#1088](https://github.com/lambdaclass/cairo-rs/pull/1088)
     * This includes the fix for a `panic!` in `ecdsa::verify`.
       See: [#365](https://github.com/xJonathanLEI/starknet-rs/issues/365) and [#366](https://github.com/xJonathanLEI/starknet-rs/pulls/366)
diff --git a/src/vm/errors/memory_errors.rs b/src/vm/errors/memory_errors.rs
index 2fc00ae7ea..df9acbb990 100644
--- a/src/vm/errors/memory_errors.rs
+++ b/src/vm/errors/memory_errors.rs
@@ -94,6 +94,8 @@ pub enum MemoryError {
     // SegmentArenaBuiltin
     #[error("segment_arena_builtin: assert used >= INITIAL_SEGMENT_SIZE")]
     InvalidUsedSizeSegmentArena,
+    #[error("Vector capacity exceeded")]
+    VecCapacityExceeded,
 }
 
 #[derive(Debug, PartialEq, Eq, Error)]
diff --git a/src/vm/vm_memory/memory.rs b/src/vm/vm_memory/memory.rs
index 67c64750af..c5c41400b6 100644
--- a/src/vm/vm_memory/memory.rs
+++ b/src/vm/vm_memory/memory.rs
@@ -136,8 +136,15 @@ impl Memory {
 
         //Check if the element is inserted next to the last one on the segment
         //Forgoing this check would allow data to be inserted in a different index
-        if segment.len() <= value_offset {
-            segment.resize(value_offset + 1, None);
+        let (len, capacity) = (segment.len(), segment.capacity());
+        if len <= value_offset {
+            let new_len = value_offset
+                .checked_add(1)
+                .ok_or(MemoryError::VecCapacityExceeded)?;
+            segment
+                .try_reserve(new_len.saturating_sub(capacity))
+                .map_err(|_| MemoryError::VecCapacityExceeded)?;
+            segment.resize(new_len, None);
         }
         // At this point there's *something* in there
 
@@ -1523,6 +1530,20 @@ mod memory_tests {
         assert_eq!((ord, pos), mem.memcmp(lhs.into(), rhs.into(), len));
     }
 
+    #[test]
+    fn insert_alloc_fails_gracefully() {
+        let mut mem = memory![((0, 0), 1)];
+        let err = mem.insert((0, usize::MAX >> 1).into(), Felt252::one());
+        assert_eq!(err, Err(MemoryError::VecCapacityExceeded));
+    }
+
+    #[test]
+    fn insert_overflow_fails_gracefully() {
+        let mut mem = memory![((0, 0), 1)];
+        let err = mem.insert((0, usize::MAX).into(), Felt252::one());
+        assert_eq!(err, Err(MemoryError::VecCapacityExceeded));
+    }
+
     #[test]
     fn memcmp() {
         check_memcmp((0, 0), (0, 0), 3, Equal, 3);