Fix builtin security check.

lambdaclass · Jan 15, 2023 · 322ad1b · 322ad1b
1 parent 23536dc
commit 322ad1b
Showing 1 changed file with 48 additions and 28 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",
@@ -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 as usize {
-            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,49 +1030,61 @@ 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![
             None,
             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!(