Fix builtin security check.

lambdaclass · Jan 12, 2023 · c1f7e6b · c1f7e6b
1 parent 23536dc
commit c1f7e6b
Showing 1 changed file with 33 additions and 27 deletions.
diff --git 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 = match offsets.iter().max() {
+            Some(x) => div_floor(*x, cells_per_instance as usize) + 1,
+            None => 0,
+        };
         if n > div_floor(offsets.len(), n_input_cells as usize) {
             return Err(MemoryError::MissingMemoryCells(match self {
                 BuiltinRunner::Bitwise(_) => "bitwise",
@@ -303,17 +306,14 @@ impl BuiltinRunner {
             .into());
         }
 
-        // 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 missing_offsets = Vec::new();
         for i in 0..n as usize {
-            let offset = cells_per_instance as usize * i;
+            let 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 => {
+                let offset = offset_base + j;
+                match offsets.contains(&offset) {
+                    true => {}
+                    false => {
                         missing_offsets.push(offset);
                     }
                 }
@@ -937,7 +937,6 @@ mod tests {
             mayberelocatable!(0, 2).into(),
             mayberelocatable!(0, 3).into(),
             mayberelocatable!(0, 4).into(),
-            mayberelocatable!(0, 5).into(),
         ]];
 
         assert_eq!(
@@ -1020,8 +1019,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![
@@ -1041,28 +1045,32 @@ mod tests {
 
     #[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 +1089,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 +1108,6 @@ mod tests {
             mayberelocatable!(0, 4).into(),
             mayberelocatable!(0, 5).into(),
             mayberelocatable!(0, 6).into(),
-            mayberelocatable!(0, 7).into(),
         ]];
 
         assert_eq!(