From f96a4ffed38ee5a494fd7e235ca24bdb8404b81c Mon Sep 17 00:00:00 2001
From: aniampio <anna.piotrowska.15@ucl.ac.uk>
Date: Tue, 19 Apr 2022 16:55:57 +0300
Subject: [PATCH] Update get range proof signature function to return an error

---
 common/nym-compact-ecash/src/error.rs              |  3 +++
 common/nym-compact-ecash/src/proofs/proof_spend.rs |  5 +++--
 common/nym-compact-ecash/src/scheme/mod.rs         |  2 +-
 common/nym-compact-ecash/src/scheme/setup.rs       | 13 +++++++++++--
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/common/nym-compact-ecash/src/error.rs b/common/nym-compact-ecash/src/error.rs
index 04ce046cbd..0937ffa10e 100644
--- a/common/nym-compact-ecash/src/error.rs
+++ b/common/nym-compact-ecash/src/error.rs
@@ -25,6 +25,9 @@ pub enum CompactEcashError {
     #[error("Spend Verification related error: {0}")]
     Spend(String),
 
+    #[error("ZKP Proof related error: {0}")]
+    RangeProofOutOfBound(String),
+
     #[error("Tried to deserialize {object} with bytes of invalid length. Expected {actual} < {} or {modulus_target} % {modulus} == 0")]
     DeserializationInvalidLength {
         actual: usize,
diff --git a/common/nym-compact-ecash/src/proofs/proof_spend.rs b/common/nym-compact-ecash/src/proofs/proof_spend.rs
index e0976b3317..18801e884f 100644
--- a/common/nym-compact-ecash/src/proofs/proof_spend.rs
+++ b/common/nym-compact-ecash/src/proofs/proof_spend.rs
@@ -319,7 +319,7 @@ impl SpendProof {
 mod tests {
     use bls12_381::{G1Projective, G2Projective, Scalar};
     use group::Curve;
-    use rand::thread_rng;
+    use rand::{Rng, thread_rng};
 
     use crate::proofs::proof_spend::{SpendInstance, SpendProof, SpendWitness};
     use crate::scheme::{pseudorandom_fgt, pseudorandom_fgv};
@@ -350,6 +350,7 @@ mod tests {
         let v = grparams.random_scalar();
         let t = grparams.random_scalar();
         let attributes = vec![sk, v, t];
+        // the below value must be from range 0 to params.L()
         let l = 5;
         let gamma1 = *grparams.gamma1();
         let g1 = *grparams.gen1();
@@ -387,7 +388,7 @@ mod tests {
         let o_lambda = ((o_a + o_d) * lambda).neg();
 
         // parse the signature associated with value l
-        let sign_l = params.get_sign_by_idx(l);
+        let sign_l = params.get_sign_by_idx(l).unwrap();
         // randomise the signature associated with value l
         let (sign_l_prime, r_l) = sign_l.randomise(grparams);
         // compute kappa_l
diff --git a/common/nym-compact-ecash/src/scheme/mod.rs b/common/nym-compact-ecash/src/scheme/mod.rs
index 9c976dbf4e..55838744f3 100644
--- a/common/nym-compact-ecash/src/scheme/mod.rs
+++ b/common/nym-compact-ecash/src/scheme/mod.rs
@@ -121,7 +121,7 @@ impl Wallet {
         let o_lambda = ((o_a + o_d) * lambda).neg();
 
         // parse the signature associated with value l
-        let sign_l = params.get_sign_by_idx(self.l());
+        let sign_l = params.get_sign_by_idx(self.l())?;
         // randomise the signature associated with value l
         let (sign_l_prime, sign_l_blinding_factor) = sign_l.randomise(grparams);
         // compute kappa_l
diff --git a/common/nym-compact-ecash/src/scheme/setup.rs b/common/nym-compact-ecash/src/scheme/setup.rs
index 229ef54f60..a5b78c21f3 100644
--- a/common/nym-compact-ecash/src/scheme/setup.rs
+++ b/common/nym-compact-ecash/src/scheme/setup.rs
@@ -5,7 +5,7 @@ use ff::Field;
 use group::{Curve, GroupEncoding};
 use rand::thread_rng;
 
-use crate::error::Result;
+use crate::error::{CompactEcashError, Result};
 use crate::utils::{hash_g1, Signature};
 
 const ATTRIBUTES_LEN: usize = 3;
@@ -112,7 +112,16 @@ impl Parameters {
     pub fn pkRP(&self) -> &PublicKeyRP { &self.pkRP }
     pub fn L(&self) -> u64 { self.L }
     pub fn signs(&self) -> &HashMap<u64, Signature> { &self.signs }
-    pub fn get_sign_by_idx(&self, idx: u64) -> &Signature { self.signs.get(&idx).unwrap() }
+    pub fn get_sign_by_idx(&self, idx: u64) -> Result<&Signature> {
+        match self.signs.get(&idx) {
+            Some(val) => return Ok(val),
+            None => return
+                Err(CompactEcashError::RangeProofOutOfBound
+                    ("Cannot find the range proof signature for the given value. \
+                        Check if the requested value is within the bound 0..L".to_string()
+                    ))
+        }
+    }
 }
 
 pub fn setup() -> Parameters {