Merge pull request #32 from anoma/xuyang/error_handling

refactor error handling and add more tests
anoma · Nov 25, 2024 · 0e5dd1a · 0e5dd1a
2 parents dee4558 + 65d246c
commit 0e5dd1a
Show file tree

Hide file tree

Showing 20 changed files with 448 additions and 510 deletions.
diff --git a/lib/cairo.ex b/lib/cairo.ex
@@ -88,7 +88,7 @@ defmodule Cairo do
     to: Cairo.CairoProver,
     as: :program_hash
 
-  @spec felt_to_string(list(byte())) :: binary()
+  @spec felt_to_string(list(byte())) :: binary() | {:error, term()}
   defdelegate felt_to_string(felt),
     to: Cairo.CairoProver,
     as: :cairo_felt_to_string

diff --git a/lib/cairo/cairo.ex b/lib/cairo/cairo.ex
@@ -51,6 +51,7 @@ defmodule Cairo.CairoProver do
   @spec program_hash(list(byte())) :: nif_result(list(byte()))
   def program_hash(_public_inputs), do: error()
 
+  @spec cairo_felt_to_string(list(byte())) :: nif_result(binary())
   def cairo_felt_to_string(_felt), do: error()
 
   def cairo_generate_compliance_input_json(

diff --git a/native/cairo_prover/Cargo.lock b/native/cairo_prover/Cargo.lock
diff --git a/native/cairo_prover/Cargo.toml b/native/cairo_prover/Cargo.toml
@@ -11,9 +11,9 @@ crate-type = ["cdylib"]
 
 [dependencies]
 rustler = "0.31.0"
-cairo-platinum-prover = { git = "https://github.com/lambdaclass/lambdaworks", version = "0.9.0"}
-stark-platinum-prover = { git = "https://github.com/lambdaclass/lambdaworks", version = "0.9.0"}
-lambdaworks-math = { git = "https://github.com/lambdaclass/lambdaworks", version = "0.9.0"}
+cairo-platinum-prover = { git = "https://github.com/heliaxdev/lambdaworks", branch = "cairo_rm"}
+stark-platinum-prover = { git = "https://github.com/heliaxdev/lambdaworks", branch = "cairo_rm"}
+lambdaworks-math = { git = "https://github.com/heliaxdev/lambdaworks", branch = "cairo_rm"}
 bincode = "2.0.0-rc.3"
 serde_json = { version = "1.0", features = ["preserve_order"] }
 hashbrown = { version = "0.14.0", features = ["serde"] }
@@ -26,3 +26,4 @@ num-integer = { version = "0.1.45", default-features = false }
 rand = "0.8.5"
 lazy_static = "1.4"
 serde = { version = "1.0.160", features = ["derive"] }
+thiserror = "1.0"
diff --git a/native/cairo_prover/src/compliance_input.rs b/native/cairo_prover/src/compliance_input.rs
@@ -1,4 +1,4 @@
-use crate::utils::felt_to_string;
+use crate::{error::CairoError, utils::felt_to_string};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]
@@ -31,32 +31,31 @@ struct PathNode {
 
 impl ComplianceInputJson {
     pub fn to_json_string(
-        input_resource: &Vec<u8>,
-        output_resource: &Vec<u8>,
-        path: &Vec<Vec<u8>>,
+        input_resource: Vec<u8>,
+        output_resource: Vec<u8>,
+        path: Vec<Vec<u8>>,
         pos: u64,
-        input_nf_key: &Vec<u8>,
-        eph_root: &Vec<u8>,
-        rcv: &Vec<u8>,
-    ) -> String {
-        let input = ResourceJson::from_bytes(input_resource);
-        let output = ResourceJson::from_bytes(output_resource);
+        input_nf_key: Vec<u8>,
+        eph_root: Vec<u8>,
+        rcv: Vec<u8>,
+    ) -> Result<String, CairoError> {
+        let input = ResourceJson::from_bytes(input_resource)?;
+        let output = ResourceJson::from_bytes(output_resource)?;
 
-        let rcv = felt_to_string(rcv);
-        let eph_root = felt_to_string(eph_root);
-        let input_nf_key = felt_to_string(input_nf_key);
+        let rcv = felt_to_string(rcv)?;
+        let eph_root = felt_to_string(eph_root)?;
+        let input_nf_key = felt_to_string(input_nf_key)?;
         let mut next_pos = pos;
-        let merkle_path = path
-            .iter()
-            .map(|v| {
-                let snd = if next_pos % 2 == 0 { false } else { true };
-                next_pos >>= 1;
-                PathNode {
-                    fst: felt_to_string(v),
-                    snd,
-                }
-            })
-            .collect();
+        let mut merkle_path = Vec::new();
+        for node in path.into_iter() {
+            let snd = next_pos % 2 != 0;
+            next_pos >>= 1;
+            let node = PathNode {
+                fst: felt_to_string(node)?,
+                snd,
+            };
+            merkle_path.push(node);
+        }
 
         let compliance_input = Self {
             input,
@@ -66,22 +65,22 @@ impl ComplianceInputJson {
             rcv,
             eph_root,
         };
-        serde_json::to_string(&compliance_input).unwrap()
+        Ok(serde_json::to_string(&compliance_input)?)
     }
 }
 
 impl ResourceJson {
-    pub fn from_bytes(bytes: &Vec<u8>) -> Self {
-        Self {
-            logic: felt_to_string(&bytes[0..32].to_vec()),
-            label: felt_to_string(&bytes[32..64].to_vec()),
-            quantity: felt_to_string(&bytes[64..96].to_vec()),
-            data: felt_to_string(&bytes[96..128].to_vec()),
-            nonce: felt_to_string(&bytes[128..160].to_vec()),
-            npk: felt_to_string(&bytes[160..192].to_vec()),
-            rseed: felt_to_string(&bytes[192..224].to_vec()),
-            eph: if bytes[224] == 0 { false } else { true },
-        }
+    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, CairoError> {
+        Ok(Self {
+            logic: felt_to_string(bytes[0..32].to_vec())?,
+            label: felt_to_string(bytes[32..64].to_vec())?,
+            quantity: felt_to_string(bytes[64..96].to_vec())?,
+            data: felt_to_string(bytes[96..128].to_vec())?,
+            nonce: felt_to_string(bytes[128..160].to_vec())?,
+            npk: felt_to_string(bytes[160..192].to_vec())?,
+            rseed: felt_to_string(bytes[192..224].to_vec())?,
+            eph: bytes[224] != 0,
+        })
     }
 }
 
@@ -97,14 +96,15 @@ fn test_compliance_input_json() {
     let path = (0..32).map(|_| random_felt()).collect();
 
     let json = ComplianceInputJson::to_json_string(
-        &random_resouce.to_vec(),
-        &random_resouce.to_vec(),
-        &path,
+        random_resouce.to_vec(),
+        random_resouce.to_vec(),
+        path,
         0,
-        &random_felt(),
-        &random_felt(),
-        &random_felt(),
-    );
+        random_felt(),
+        random_felt(),
+        random_felt(),
+    )
+    .unwrap();
 
     println!("compliance_input_json: {}", json);
 }
diff --git a/native/cairo_prover/src/encryption.rs b/native/cairo_prover/src/encryption.rs
@@ -1,3 +1,4 @@
+use crate::{error::CairoError, utils::bytes_to_felt_vec};
 use starknet_crypto::{poseidon_hash, poseidon_hash_many};
 use starknet_curve::curve_params::GENERATOR;
 use starknet_types_core::{
@@ -29,9 +30,15 @@ impl Ciphertext {
         &self.0
     }
 
-    pub fn encrypt(messages: &[Felt], pk: &AffinePoint, sk: &Felt, encrypt_nonce: &Felt) -> Self {
+    pub fn encrypt(
+        messages: &[Felt],
+        pk: &AffinePoint,
+        sk: &Felt,
+        encrypt_nonce: &Felt,
+    ) -> Result<Self, CairoError> {
         // Generate the secret key
-        let (secret_key_x, secret_key_y) = SecretKey::from_dh_exchange(pk, sk).get_coordinates();
+        let secret_key = SecretKey::from_dh_exchange(pk, sk)?;
+        let (secret_key_x, secret_key_y) = secret_key.get_coordinates();
 
         // Pad the messages
         let plaintext = Plaintext::padding(messages);
@@ -56,33 +63,34 @@ impl Ciphertext {
         cipher.push(poseidon_state);
 
         // Add sender's public key
-        let generator = ProjectivePoint::from_affine(GENERATOR.x(), GENERATOR.y()).unwrap();
-        let sender_pk = (&generator * *sk).to_affine().unwrap();
+        let generator = ProjectivePoint::from_affine(GENERATOR.x(), GENERATOR.y())
+            .map_err(|_| CairoError::InvalidAffinePoint)?;
+        let sender_pk = (&generator * *sk)
+            .to_affine()
+            .map_err(|_| CairoError::InvalidAffinePoint)?;
         cipher.push(sender_pk.x());
         cipher.push(sender_pk.y());
 
         // Add encrypt_nonce
         cipher.push(*encrypt_nonce);
 
-        cipher.into()
-    }
+        let ret: [Felt; CIPHERTEXT_NUM] = cipher
+            .try_into()
+            .map_err(|_| CairoError::InvalidCiphertextLength)?;
 
-    pub fn decrypt(&self, sk: &Felt) -> Option<Vec<Felt>> {
-        let cipher_text = self.inner();
-        let cipher_len = cipher_text.len();
-        if cipher_len != CIPHERTEXT_NUM {
-            return None;
-        }
+        Ok(Self(ret))
+    }
 
-        let mac = cipher_text[CIPHERTEXT_MAC];
-        let pk_x = cipher_text[CIPHERTEXT_PK_X];
-        let pk_y = cipher_text[CIPHERTEXT_PK_Y];
-        let encrypt_nonce = cipher_text[CIPHERTEXT_NONCE];
+    pub fn decrypt(&self, sk: &Felt) -> Result<Vec<Felt>, CairoError> {
+        let mac = self.inner()[CIPHERTEXT_MAC];
+        let pk_x = self.inner()[CIPHERTEXT_PK_X];
+        let pk_y = self.inner()[CIPHERTEXT_PK_Y];
+        let encrypt_nonce = self.inner()[CIPHERTEXT_NONCE];
 
         if let Ok(pk) = AffinePoint::new(pk_x, pk_y) {
             // Generate the secret key
-            let (secret_key_x, secret_key_y) =
-                SecretKey::from_dh_exchange(&pk, sk).get_coordinates();
+            let sk = SecretKey::from_dh_exchange(&pk, sk)?;
+            let (secret_key_x, secret_key_y) = sk.get_coordinates();
 
             // Init poseidon sponge state
             let mut poseidon_state = poseidon_hash_many(&vec![
@@ -94,30 +102,28 @@ impl Ciphertext {
 
             // Decrypt
             let mut msg = vec![];
-            for cipher_element in &cipher_text[0..PLAINTEXT_NUM] {
+            for cipher_element in &self.inner()[0..PLAINTEXT_NUM] {
                 let msg_element = *cipher_element - poseidon_state;
                 msg.push(msg_element);
                 poseidon_state = poseidon_hash(*cipher_element, secret_key_x);
             }
 
             if mac != poseidon_state {
-                return None;
+                return Err(CairoError::DecryptionFailure);
             }
 
-            Some(msg)
+            Ok(msg)
         } else {
-            return None;
+            Err(CairoError::InvalidPublicKey)
         }
     }
-}
 
-impl From<Vec<Felt>> for Ciphertext {
-    fn from(input_vec: Vec<Felt>) -> Self {
-        Ciphertext(
-            input_vec
-                .try_into()
-                .expect("public input with incorrect length"),
-        )
+    pub fn from_bytes(input_vec: Vec<Vec<u8>>) -> Result<Self, CairoError> {
+        let cipher_felt = bytes_to_felt_vec(input_vec)?;
+        let cipher: [Felt; CIPHERTEXT_NUM] = cipher_felt
+            .try_into()
+            .map_err(|_| CairoError::InvalidCiphertextLength)?;
+        Ok(Self(cipher))
     }
 }
 
@@ -149,12 +155,13 @@ impl From<Vec<Felt>> for Plaintext {
 }
 
 impl SecretKey {
-    pub fn from_dh_exchange(pk: &AffinePoint, sk: &Felt) -> Self {
-        Self(
-            (&ProjectivePoint::try_from(pk.clone()).unwrap() * *sk)
-                .to_affine()
-                .unwrap(),
-        )
+    pub fn from_dh_exchange(pk: &AffinePoint, sk: &Felt) -> Result<Self, CairoError> {
+        let pk_projective =
+            ProjectivePoint::try_from(pk.clone()).map_err(|_| CairoError::InvalidAffinePoint)?;
+        let key = (&pk_projective * *sk)
+            .to_affine()
+            .map_err(|_| CairoError::InvalidDHKey)?;
+        Ok(Self(key))
     }
 
     pub fn get_coordinates(&self) -> (Felt, Felt) {
@@ -173,7 +180,7 @@ fn test_encryption() {
     let encrypt_nonce = Felt::ONE;
 
     // Encryption
-    let cipher = Ciphertext::encrypt(&messages, &pk, &sender_sk, &encrypt_nonce);
+    let cipher = Ciphertext::encrypt(&messages, &pk, &sender_sk, &encrypt_nonce).unwrap();
 
     // Decryption
     let decryption = cipher.decrypt(&Felt::ONE).unwrap();