check expected error and correct box usage for errors

Author: curiecrypt

mithril-stm/src/error.rs

@@ -36,7 +36,7 @@ pub enum MultiSignatureError {
 
     /// Verification key is the infinity
     #[error("Verification key is the infinity")]
-    VerificationKeyInfinity(VerificationKey),
+    VerificationKeyInfinity(Box<VerificationKey>),
 }
 
 /// Errors which can be output by Mithril single signature verification.
@@ -146,6 +146,10 @@ pub enum RegisterError {
     #[error("This key has already been registered.")]
     KeyRegistered(Box<VerificationKey>),
 
+    /// Verification key is the infinity
+    #[error("Verification key is the infinity")]
+    VerificationKeyInfinity(Box<VerificationKey>),
+
     /// The supplied key is not valid
     #[error("The verification of correctness of the supplied key is invalid.")]
     KeyInvalid(Box<VerificationKeyPoP>),
@@ -263,6 +267,7 @@ impl From<MultiSignatureError> for RegisterError {
         match e {
             MultiSignatureError::SerializationError => Self::SerializationError,
             MultiSignatureError::KeyInvalid(e) => Self::KeyInvalid(e),
+            MultiSignatureError::VerificationKeyInfinity(e) => Self::VerificationKeyInfinity(e),
             _ => unreachable!(),
         }
     }
@@ -273,7 +278,7 @@ impl From<MultiSignatureError> for RegisterError {
 pub(crate) fn blst_err_to_mithril(
     e: BLST_ERROR,
     sig: Option<Signature>,
-    key: Option<VerificationKey>
+    key: Option<VerificationKey>,
 ) -> Result<(), MultiSignatureError> {
     match e {
         BLST_ERROR::BLST_SUCCESS => Ok(()),
@@ -282,7 +287,7 @@ pub(crate) fn blst_err_to_mithril(
                 return Err(MultiSignatureError::SignatureInfinity(s));
             }
             if let Some(vk) = key {
-                return Err(MultiSignatureError::VerificationKeyInfinity(vk));
+                return Err(MultiSignatureError::VerificationKeyInfinity(Box::new(vk)));
             }
             Err(MultiSignatureError::SerializationError)
         }

mithril-stm/src/key_reg.rs

@@ -42,12 +42,9 @@ impl KeyReg {
     /// The function fails when the proof of possession is invalid or when the key is already registered.
     pub fn register(&mut self, stake: Stake, pk: VerificationKeyPoP) -> Result<(), RegisterError> {
         if let Entry::Vacant(e) = self.keys.entry(pk.vk) {
-            if pk.check().is_ok() {
-                e.insert(stake);
-                return Ok(());
-            } else {
-                return Err(RegisterError::KeyInvalid(Box::new(pk)));
-            }
+            pk.check()?;
+            e.insert(stake);
+            return Ok(());
         }
         Err(RegisterError::KeyRegistered(Box::new(pk.vk)))
     }

mithril-stm/src/multi_sig.rs

@@ -406,7 +406,8 @@ impl Signature {
 
         blst_err_to_mithril(
             aggr_sig.0.verify(false, msg, &[], &[], &aggr_vk.0, false),
-            Some(aggr_sig), None,
+            Some(aggr_sig),
+            None,
         )
     }
 
@@ -434,7 +435,7 @@ impl Signature {
         blst_err_to_mithril(
             batched_sig.aggregate_verify(false, &slice_msgs, &[], &p2_vks, false),
             None,
-            None
+            None,
         )
         .map_err(|_| MultiSignatureError::BatchInvalid)
     }
@@ -636,6 +637,7 @@ mod unsafe_helpers {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::error::RegisterError;
     use crate::key_reg::KeyReg;
     use proptest::prelude::*;
     use rand_chacha::ChaCha20Rng;
@@ -674,7 +676,8 @@ mod tests {
             let p1 = blst_p1::default();
             let sig_infinity = Signature(p1_affine_to_sig(&p1));
 
-            assert!(sig_infinity.verify(&msg, &vk).is_err());
+            let result = sig_infinity.verify(&msg, &vk);
+            assert_eq!(result, Err(MultiSignatureError::SignatureInfinity(sig_infinity)));
         }
 
         #[test]
@@ -687,7 +690,8 @@ mod tests {
             let vk_infinity = VerificationKey(p2_affine_to_vk(&p2));
             let vkpop_infinity = VerificationKeyPoP { vk: vk_infinity, pop };
 
-            assert!(vkpop_infinity.check().is_err());
+            let result = vkpop_infinity.check();
+            assert_eq!(result, Err(MultiSignatureError::VerificationKeyInfinity(Box::new(vkpop_infinity.vk))));
         }
 
         #[test]
@@ -707,7 +711,8 @@ mod tests {
                 let _ = kr.register(1, vkpop);
             }
 
-            assert!(kr.register(1, vkpop_infinity).is_err());
+            let result = kr.register(1, vkpop_infinity);
+            assert_eq!(result, Err(RegisterError::VerificationKeyInfinity(Box::new(vkpop_infinity.vk))));
         }
 
         #[test]