Refactor SignatureBuiltinRunner::add_relocation_rule + add tests fo…

…r signature.rs (#793)

* Refactor code to remove unrechable code

* Refactore add_relocation_rules

* Fix code

* Remove unused method as_any

* Add test

* Add tests

* Update src/vm/runners/builtin_runner/signature.rs

Co-authored-by: Mario Rugiero <mario.rugiero@lambdaclass.com>

* Remove unused trait

---------

Co-authored-by: Mario Rugiero <mario.rugiero@lambdaclass.com>

src/vm/errors/memory_errors.rs

@@ -1,3 +1,4 @@
+use felt::Felt;
 use thiserror::Error;
 
 use crate::types::relocatable::{MaybeRelocatable, Relocatable};
@@ -50,10 +51,16 @@ pub enum MemoryError {
     MissingMemoryCellsWithOffsets(&'static str, Vec<usize>),
     #[error("ErrorInitializing Verifying Key from public key: {0:?}")]
     InitializingVerifyingKey(Vec<u8>),
-    #[error("Invalid Signature")]
-    InvalidSignature,
-    #[error("Signature not found")]
-    SignatureNotFound,
+    #[error(
+        "Signature {0}, is invalid, with respect to the public key {1}, 
+    and the message hash {2}."
+    )]
+    InvalidSignature(String, Felt, Felt),
+    #[error(
+        "Signature hint is missing for ECDSA builtin at address {0}.
+    Add it using 'ecdsa_builtin.add_signature'."
+    )]
+    SignatureNotFound(Relocatable),
     #[error("Could not create pubkey from: {0:?}")]
     ErrorParsingPubKey(String),
     #[error("Could not retrieve message from: {0:?}")]
@@ -62,6 +69,10 @@ pub enum MemoryError {
     ErrorVerifyingSignature,
     #[error("Couldn't obtain a mutable accessed offset")]
     CantGetMutAccessedOffset,
+    #[error("ECDSA builtin: Expected public key at address {0} to be an integer")]
+    PubKeyNonInt(Relocatable),
+    #[error("ECDSA builtin: Expected message hash at address {0} to be an integer")]
+    MsgNonInt(Relocatable),
     #[error("Failed to convert String: {0} to FieldElement")]
     FailedStringToFieldElementConversion(String),
 }

src/vm/runners/builtin_runner/signature.rs

@@ -14,10 +14,10 @@ use crate::{
     },
 };
 use felt::Felt;
-use num_integer::{div_ceil, Integer};
+use num_integer::div_ceil;
 use num_traits::ToPrimitive;
 use starknet_crypto::{verify, FieldElement, Signature};
-use std::{any::Any, cell::RefCell, collections::HashMap, rc::Rc};
+use std::{cell::RefCell, collections::HashMap, rc::Rc};
 
 #[derive(Debug, Clone)]
 pub struct SignatureBuiltinRunner {
@@ -99,57 +99,47 @@ impl SignatureBuiltinRunner {
         let cells_per_instance = self.cells_per_instance;
         let signatures = Rc::clone(&self.signatures);
         let rule: ValidationRule = ValidationRule(Box::new(
-            move |memory: &Memory,
-                  address: &Relocatable|
-                  -> Result<Vec<Relocatable>, MemoryError> {
-                let address_offset = address.offset.mod_floor(&(cells_per_instance as usize));
-                let mem_addr_sum = memory.get(&(address + 1_i32));
-                let mem_addr_less = if address.offset > 0 {
-                    memory.get(
-                        &address
-                            .sub_usize(1)
-                            .map_err(|_| MemoryError::NumOutOfBounds)?,
-                    )
-                } else {
-                    Ok(None)
+            move |memory: &Memory, addr: &Relocatable| -> Result<Vec<Relocatable>, MemoryError> {
+                let cell_index = addr.offset % cells_per_instance as usize;
+
+                let (pubkey_addr, message_addr) = match cell_index {
+                    0 => (*addr, addr + 1),
+                    1 => match addr.sub_usize(1) {
+                        Ok(prev_addr) => (prev_addr, *addr),
+                        Err(_) => return Ok(vec![]),
+                    },
+                    _ => return Ok(vec![]),
                 };
-                let (pubkey_addr, msg_addr) = match (address_offset, mem_addr_sum, mem_addr_less) {
-                    (0, Ok(Some(_element)), _) => {
-                        let pubkey_addr = address;
-                        let msg_addr = address + 1_i32;
-                        (*pubkey_addr, msg_addr)
-                    }
-                    (1, _, Ok(Some(_element))) if address.offset > 0 => {
-                        let pubkey_addr = address
-                            .sub_usize(1)
-                            .map_err(|_| MemoryError::EffectiveSizesNotCalled)?;
-                        let msg_addr = address;
-                        (pubkey_addr, *msg_addr)
-                    }
-                    _ => return Ok(Vec::new()),
+
+                let pubkey = match memory.get_integer(&pubkey_addr) {
+                    Ok(num) => num,
+                    Err(_) if cell_index == 1 => return Ok(vec![]),
+                    _ => return Err(MemoryError::PubKeyNonInt(pubkey_addr)),
+                };
+
+                let msg = match memory.get_integer(&message_addr) {
+                    Ok(num) => num,
+                    Err(_) if cell_index == 0 => return Ok(vec![]),
+                    _ => return Err(MemoryError::MsgNonInt(message_addr)),
                 };
 
-                let msg = memory
-                    .get_integer(&msg_addr)
-                    .map_err(|_| MemoryError::FoundNonInt)?;
-                let pub_key = memory
-                    .get_integer(&pubkey_addr)
-                    .map_err(|_| MemoryError::FoundNonInt)?;
                 let signatures_map = signatures.borrow();
                 let signature = signatures_map
                     .get(&pubkey_addr)
-                    .ok_or(MemoryError::SignatureNotFound)?;
-                let public_key = FieldElement::from_dec_str(&pub_key.to_str_radix(10))
-                    .map_err(|_| MemoryError::ErrorParsingPubKey(pub_key.to_str_radix(10)))?;
+                    .ok_or(MemoryError::SignatureNotFound(pubkey_addr))?;
+
+                let public_key = FieldElement::from_dec_str(&pubkey.to_str_radix(10))
+                    .map_err(|_| MemoryError::ErrorParsingPubKey(pubkey.to_str_radix(10)))?;
                 let (r, s) = (signature.r, signature.s);
                 let message = FieldElement::from_dec_str(&msg.to_str_radix(10))
                     .map_err(|_| MemoryError::ErrorRetrievingMessage(msg.to_str_radix(10)))?;
-                let was_verified = verify(&public_key, &message, &r, &s)
-                    .map_err(|_| MemoryError::ErrorVerifyingSignature)?;
-                if was_verified {
-                    Ok(vec![])
-                } else {
-                    Err(MemoryError::InvalidSignature)
+                match verify(&public_key, &message, &r, &s) {
+                    Ok(true) => Ok(vec![]),
+                    _ => Err(MemoryError::InvalidSignature(
+                        signature.to_string(),
+                        pubkey.into_owned(),
+                        msg.into_owned(),
+                    )),
                 }
             },
         ));
@@ -170,10 +160,6 @@ impl SignatureBuiltinRunner {
         Ok(None)
     }
 
-    pub fn as_any(&self) -> &dyn Any {
-        self
-    }
-
     pub fn ratio(&self) -> u32 {
         self.ratio
     }
@@ -531,4 +517,59 @@ mod tests {
         let result = builtin.deduce_memory_cell(&Relocatable::from((0, 5)), &memory);
         assert_eq!(result, Ok(None));
     }
+
+    #[test]
+    fn get_allocated_memory_units_safe_div_fail() {
+        let builtin = SignatureBuiltinRunner::new(&EcdsaInstanceDef::default(), true);
+        let mut vm = vm!();
+        vm.current_step = 500;
+        assert_eq!(
+            builtin.get_allocated_memory_units(&vm),
+            Err(MemoryError::ErrorCalculatingMemoryUnits)
+        )
+    }
+
+    #[test]
+    fn get_used_cells_and_allocated_size_safe_div_fail() {
+        let builtin = SignatureBuiltinRunner::new(&EcdsaInstanceDef::default(), true);
+        let mut vm = vm!();
+        vm.current_step = 500;
+        assert_eq!(
+            builtin.get_used_cells_and_allocated_size(&vm),
+            Err(MemoryError::InsufficientAllocatedCells)
+        )
+    }
+
+    #[test]
+    fn get_used_cells_and_allocated_size_insufficient_allocated() {
+        let builtin = SignatureBuiltinRunner::new(&EcdsaInstanceDef::default(), true);
+        let mut vm = vm!();
+        vm.segments.segment_used_sizes = Some(vec![50]);
+        assert_eq!(
+            builtin.get_used_cells_and_allocated_size(&vm),
+            Err(MemoryError::InsufficientAllocatedCells)
+        )
+    }
+
+    #[test]
+    fn final_stack_invalid_stop_pointer() {
+        let mut builtin = SignatureBuiltinRunner::new(&EcdsaInstanceDef::default(), true);
+        let mut vm = vm!();
+        vm.memory = memory![((0, 0), (1, 0))];
+        assert_eq!(
+            builtin.final_stack(&vm.segments, &vm.memory, (0, 1).into()),
+            Err(RunnerError::InvalidStopPointer("ecdsa".to_string()))
+        )
+    }
+
+    #[test]
+    fn final_stack_no_used_instances() {
+        let mut builtin = SignatureBuiltinRunner::new(&EcdsaInstanceDef::default(), true);
+        let mut vm = vm!();
+        vm.memory = memory![((0, 0), (0, 0))];
+        assert_eq!(
+            builtin.final_stack(&vm.segments, &vm.memory, (0, 1).into()),
+            Err(RunnerError::FinalStack)
+        )
+    }
 }

src/vm/vm_memory/memory.rs

@@ -621,7 +621,7 @@ mod memory_tests {
         builtin.initialize_segments(&mut segments, &mut memory);
         builtin.add_validation_rule(&mut memory).unwrap();
         let error = memory.validate_existing_memory();
-        assert_eq!(error, Err(MemoryError::SignatureNotFound));
+        assert_eq!(error, Err(MemoryError::SignatureNotFound((1, 0).into())));
     }
 
     #[test]