From bb17218ab1e6254a0afe43c4beb4087cfdb6c9e2 Mon Sep 17 00:00:00 2001
From: Alon Haramati <alonh@starkware.co>
Date: Thu, 12 Jan 2023 12:37:01 +0200
Subject: [PATCH] Fix builtin security check.

---
 src/vm/runners/builtin_runner/mod.rs | 76 ++++++++++++++++++----------
 1 file changed, 48 insertions(+), 28 deletions(-)

diff --git a/src/vm/runners/builtin_runner/mod.rs b/src/vm/runners/builtin_runner/mod.rs
index 5eafead68b..5107468d33 100644
--- a/src/vm/runners/builtin_runner/mod.rs
+++ b/src/vm/runners/builtin_runner/mod.rs
@@ -289,7 +289,10 @@ impl BuiltinRunner {
             })
             .collect::<Vec<_>>();
 
-        let n = div_floor(offsets.len(), cells_per_instance as usize);
+        let n = offsets
+            .iter()
+            .max()
+            .map_or(0, |x| div_floor(*x, cells_per_instance as usize) + 1);
         if n > div_floor(offsets.len(), n_input_cells as usize) {
             return Err(MemoryError::MissingMemoryCells(match self {
                 BuiltinRunner::Bitwise(_) => "bitwise",
@@ -305,20 +308,28 @@ impl BuiltinRunner {
 
         // Since both offsets and this iterator are ordered, a simple pointer is
         // enough to check if the values are present.
-        let mut offsets_iter = offsets.iter().copied().peekable();
+        let mut offsets_iter = offsets.into_iter().peekable();
         let mut missing_offsets = Vec::new();
         for i in 0..n {
-            let offset = cells_per_instance as usize * i;
+            let expected_offset_base = cells_per_instance as usize * i;
             for j in 0..n_input_cells as usize {
-                let offset = offset + j;
-                match offsets_iter.next_if_eq(&offset) {
-                    Some(_) => {}
-                    None => {
-                        missing_offsets.push(offset);
+                let expected_offset = expected_offset_base + j;
+                let current_offset = loop {
+                    match offsets_iter.peek() {
+                        None => break None,
+                        Some(offset) if offset >= &expected_offset => break Some(offset),
+                        _ => {
+                            offsets_iter.next();
+                        }
                     }
+                };
+                match current_offset {
+                    Some(offset) if offset == &expected_offset => {}
+                    _ => missing_offsets.push(expected_offset),
                 }
             }
         }
+
         if !missing_offsets.is_empty() {
             return Err(MemoryError::MissingMemoryCellsWithOffsets(
                 match self {
@@ -937,7 +948,6 @@ mod tests {
             mayberelocatable!(0, 2).into(),
             mayberelocatable!(0, 3).into(),
             mayberelocatable!(0, 4).into(),
-            mayberelocatable!(0, 5).into(),
         ]];
 
         assert_eq!(
@@ -1020,8 +1030,13 @@ mod tests {
 
     #[test]
     fn run_security_checks_range_check_missing_memory_cells_with_offsets() {
-        let builtin: BuiltinRunner =
-            BuiltinRunner::RangeCheck(RangeCheckBuiltinRunner::new(8, 8, true));
+        let mut range_check_builtin = RangeCheckBuiltinRunner::new(8, 8, true);
+
+        range_check_builtin.cells_per_instance = 3;
+        range_check_builtin.n_input_cells = 2;
+
+        let builtin: BuiltinRunner = range_check_builtin.into();
+
         let mut vm = vm!();
 
         vm.memory.data = vec![vec![
@@ -1029,40 +1044,47 @@ mod tests {
             mayberelocatable!(0, 1).into(),
             mayberelocatable!(0, 2).into(),
             mayberelocatable!(0, 3).into(),
-            mayberelocatable!(0, 4).into(),
+            None,
             mayberelocatable!(0, 5).into(),
+            mayberelocatable!(0, 17).into(),
+            mayberelocatable!(0, 22).into(),
+            None,
         ]];
 
         assert_eq!(
             builtin.run_security_checks(&mut vm),
-            Err(MemoryError::MissingMemoryCellsWithOffsets("range_check", vec![0],).into()),
+            Err(MemoryError::MissingMemoryCellsWithOffsets("range_check", vec![0, 4],).into()),
         );
     }
 
     #[test]
     fn run_security_checks_range_check_missing_memory_cells() {
+        let builtin: BuiltinRunner =
+            BuiltinRunner::RangeCheck(RangeCheckBuiltinRunner::new(8, 8, true));
+        let mut vm = vm!();
+
+        vm.memory.data = vec![vec![None, mayberelocatable!(0, 0).into()]];
+
+        assert_eq!(
+            builtin.run_security_checks(&mut vm),
+            Err(MemoryError::MissingMemoryCells("range_check").into()),
+        );
+    }
+
+    #[test]
+    fn run_security_checks_range_check_empty() {
         let mut range_check_builtin = RangeCheckBuiltinRunner::new(8, 8, true);
 
-        range_check_builtin.cells_per_instance = 1;
+        range_check_builtin.cells_per_instance = 3;
         range_check_builtin.n_input_cells = 2;
 
         let builtin: BuiltinRunner = range_check_builtin.into();
 
         let mut vm = vm!();
 
-        vm.memory.data = vec![vec![
-            mayberelocatable!(0, 0).into(),
-            mayberelocatable!(0, 1).into(),
-            mayberelocatable!(0, 2).into(),
-            mayberelocatable!(0, 3).into(),
-            mayberelocatable!(0, 4).into(),
-            mayberelocatable!(0, 5).into(),
-        ]];
+        vm.memory.data = vec![vec![None, None, None]];
 
-        assert_eq!(
-            builtin.run_security_checks(&mut vm),
-            Err(MemoryError::MissingMemoryCells("range_check").into()),
-        );
+        assert_eq!(builtin.run_security_checks(&mut vm), Ok(()),);
     }
 
     #[test]
@@ -1081,7 +1103,6 @@ mod tests {
             mayberelocatable!(0, 2).into(),
             mayberelocatable!(0, 3).into(),
             mayberelocatable!(0, 4).into(),
-            mayberelocatable!(0, 5).into(),
         ]];
 
         assert_eq!(builtin.run_security_checks(&mut vm), Ok(()));
@@ -1101,7 +1122,6 @@ mod tests {
             mayberelocatable!(0, 4).into(),
             mayberelocatable!(0, 5).into(),
             mayberelocatable!(0, 6).into(),
-            mayberelocatable!(0, 7).into(),
         ]];
 
         assert_eq!(