From 3294f0cc1bded1038bc4452d38de909ddacdd4e4 Mon Sep 17 00:00:00 2001
From: Joy Wang <108701016+joyqvq@users.noreply.github.com>
Date: Mon, 2 Oct 2023 17:11:52 -0400
Subject: [PATCH] fix: add validations and improve as_arkworks

---
 fastcrypto-zkp/src/bn254/zk_login.rs     |  9 ++++++---
 fastcrypto-zkp/src/bn254/zk_login_api.rs | 10 +++++-----
 fastcrypto-zkp/src/circom.rs             |  6 +++---
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/fastcrypto-zkp/src/bn254/zk_login.rs b/fastcrypto-zkp/src/bn254/zk_login.rs
index 679e436fd0..f16a9fc665 100644
--- a/fastcrypto-zkp/src/bn254/zk_login.rs
+++ b/fastcrypto-zkp/src/bn254/zk_login.rs
@@ -341,6 +341,9 @@ impl ZkLoginInputs {
 
     /// Initialize JWTDetails by parsing header_base64 and iss_base64_details.
     pub fn init(&mut self) -> Result<Self, FastCryptoError> {
+        if BigUint::from_str(&self.address_seed).is_err() {
+            return Err(FastCryptoError::InvalidInput);
+        }
         self.jwt_details = JWTDetails::new(&self.header_base64, &self.iss_base64_details)?;
         Ok(self.to_owned())
     }
@@ -420,9 +423,9 @@ impl ZkLoginProof {
     /// Convert the Circom G1/G2/GT to arkworks G1/G2/GT
     pub fn as_arkworks(&self) -> Result<Proof<Bn254>, FastCryptoError> {
         Ok(Proof {
-            a: g1_affine_from_str_projective(self.a.clone())?,
-            b: g2_affine_from_str_projective(self.b.clone())?,
-            c: g1_affine_from_str_projective(self.c.clone())?,
+            a: g1_affine_from_str_projective(&self.a)?,
+            b: g2_affine_from_str_projective(&self.b)?,
+            c: g1_affine_from_str_projective(&self.c)?,
         })
     }
 }
diff --git a/fastcrypto-zkp/src/bn254/zk_login_api.rs b/fastcrypto-zkp/src/bn254/zk_login_api.rs
index b8a12b5796..30f2b7de29 100644
--- a/fastcrypto-zkp/src/bn254/zk_login_api.rs
+++ b/fastcrypto-zkp/src/bn254/zk_login_api.rs
@@ -38,13 +38,13 @@ static GLOBAL_VERIFYING_KEY: Lazy<PreparedVerifyingKey<Bn254>> = Lazy::new(globa
 /// Load a fixed verifying key from zkLogin.vkey output. This is based on a local setup and should not use in production.
 fn global_pvk() -> PreparedVerifyingKey<Bn254> {
     // Convert the Circom G1/G2/GT to arkworks G1/G2/GT
-    let vk_alpha_1 = g1_affine_from_str_projective(vec![
+    let vk_alpha_1 = g1_affine_from_str_projective(&vec![
         "21529901943976716921335152104180790524318946701278905588288070441048877064089".to_string(),
         "7775817982019986089115946956794180159548389285968353014325286374017358010641".to_string(),
         "1".to_string(),
     ])
     .unwrap();
-    let vk_beta_2 = g2_affine_from_str_projective(vec![
+    let vk_beta_2 = g2_affine_from_str_projective(&vec![
         vec![
             "6600437987682835329040464538375790690815756241121776438004683031791078085074"
                 .to_string(),
@@ -60,7 +60,7 @@ fn global_pvk() -> PreparedVerifyingKey<Bn254> {
         vec!["1".to_string(), "0".to_string()],
     ])
     .unwrap();
-    let vk_gamma_2 = g2_affine_from_str_projective(vec![
+    let vk_gamma_2 = g2_affine_from_str_projective(&vec![
         vec![
             "10857046999023057135944570762232829481370756359578518086990519993285655852781"
                 .to_string(),
@@ -76,7 +76,7 @@ fn global_pvk() -> PreparedVerifyingKey<Bn254> {
         vec!["1".to_string(), "0".to_string()],
     ])
     .unwrap();
-    let vk_delta_2 = g2_affine_from_str_projective(vec![
+    let vk_delta_2 = g2_affine_from_str_projective(&vec![
         vec![
             "19260309516619721648285279557078789954438346514188902804737557357941293711874"
                 .to_string(),
@@ -111,7 +111,7 @@ fn global_pvk() -> PreparedVerifyingKey<Bn254> {
             "1".to_string(),
         ],
     ] {
-        let g1 = g1_affine_from_str_projective(e).unwrap();
+        let g1 = g1_affine_from_str_projective(&e).unwrap();
         vk_gamma_abc_g1.push(g1);
     }
 
diff --git a/fastcrypto-zkp/src/circom.rs b/fastcrypto-zkp/src/circom.rs
index 0291f2566b..96b5363154 100644
--- a/fastcrypto-zkp/src/circom.rs
+++ b/fastcrypto-zkp/src/circom.rs
@@ -7,7 +7,7 @@ use fastcrypto::error::FastCryptoError;
 pub type CircomG1 = Vec<String>;
 pub type CircomG2 = Vec<Vec<String>>;
 
-pub fn g1_affine_from_str_projective(s: CircomG1) -> Result<G1Affine, FastCryptoError> {
+pub fn g1_affine_from_str_projective(s: &CircomG1) -> Result<G1Affine, FastCryptoError> {
     if s.len() != 3 {
         return Err(FastCryptoError::InvalidInput);
     }
@@ -22,13 +22,13 @@ pub fn g1_affine_from_str_projective(s: CircomG1) -> Result<G1Affine, FastCrypto
     .into())
 }
 
-pub fn g2_affine_from_str_projective(s: CircomG2) -> Result<G2Affine, FastCryptoError> {
+pub fn g2_affine_from_str_projective(s: &CircomG2) -> Result<G2Affine, FastCryptoError> {
     use ark_bn254::G2Projective;
     if s.len() != 3 {
         return Err(FastCryptoError::InvalidInput);
     }
 
-    for x in &s {
+    for x in s {
         if x.len() != 2 {
             return Err(FastCryptoError::InvalidInput);
         }