Compilation flag with zcash's implementation

mithril-core/Cargo.toml

@@ -9,10 +9,12 @@ crate-type = ["lib", "cdylib", "staticlib"]
 
 [dependencies]
 blake2      = "0.9.2"
-blst = "0.3.7"
+bls12_381 = { version = "0.7.0", features = ["experimental"], optional = true }
+blst = { version = "0.3.7", optional = true }
 bytes = "1"
 digest      = { version = "0.9.0", features = ["alloc"] }
 generic-array = "0.14"
+group = { version = "0.12.0", features = ["tests"], optional = true }
 libc        = "0.2.40"
 num-bigint  = "0.4.0"
 num-integer = { version = "0.1.44", optional = true }
@@ -52,6 +54,8 @@ name = "size_benches"
 harness = false
 
 [features]
-default = ["rug"]
+default = ["rug", "zcash"]
 rug-backend = ["rug/default"]
 num-integer-backend = ["num-integer/default"]
+zcash = ["bls12_381", "group"]
+blast = ["blst"]

mithril-core/src/error.rs

@@ -1,10 +1,15 @@
 //! Crate specific errors
 
 use crate::merkle_tree::Path;
-use crate::multi_sig::{Signature, VerificationKey, VerificationKeyPoP};
+#[cfg(feature = "zcash")]
+use crate::multi_sig_zcash::{Signature, VerificationKey, VerificationKeyPoP};
 use crate::stm::PartyId;
-use blst::BLST_ERROR;
 use digest::{Digest, FixedOutput};
+#[cfg(not(feature = "zcash"))]
+use {
+    crate::multi_sig::{Signature, VerificationKey, VerificationKeyPoP},
+    blst::BLST_ERROR,
+};
 
 // todo: better organise these errors.
 
@@ -81,6 +86,9 @@ pub enum AggregationFailure {
     /// Not enough signatures were collected, got this many instead.
     #[error("Not enough signatures. Got only {0} out of {1}.")]
     NotEnoughSignatures(u64, u64),
+    /// This error happens when we try to convert a u64 to a usize and it does not fit
+    #[error("Invalid usize conversion")]
+    InvalidUsizeConversion,
 }
 
 /// Error types for single signature verification
@@ -116,7 +124,7 @@ pub enum MerkleTreeError {
 pub enum RegisterError {
     /// This key has already been registered by a participant
     #[error("This key has already been registered.")]
-    KeyRegistered(VerificationKey),
+    KeyRegistered(Box<VerificationKey>),
     /// This participant has already been registered
     #[error("Participant {0} has already been registered.")]
     PartyRegistered(PartyId),
@@ -168,6 +176,7 @@ impl<D: Digest + Clone + FixedOutput> From<VerificationFailure<D>> for MithrilWi
     }
 }
 
+#[cfg(not(feature = "zcash"))]
 pub(crate) fn blst_err_to_atms(e: BLST_ERROR) -> Result<(), MultiSignatureError> {
     match e {
         BLST_ERROR::BLST_SUCCESS => Ok(()),

mithril-core/src/key_reg.rs

@@ -5,11 +5,13 @@ use digest::{Digest, FixedOutput};
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
-use super::multi_sig::VerificationKeyPoP;
-use super::stm::{PartyId, Stake};
+#[cfg(not(feature = "zcash"))]
+use crate::multi_sig::{VerificationKey, VerificationKeyPoP};
+#[cfg(feature = "zcash")]
+use crate::multi_sig_zcash::{VerificationKey, VerificationKeyPoP};
 
+use super::stm::{PartyId, Stake};
 use crate::merkle_tree::{MTLeaf, MerkleTree};
-use crate::multi_sig::VerificationKey;
 
 /// Struct that collects public keys and stakes of parties. Each participant (both the
 /// signers and the clerks) need to run their own instance of the key registration.
@@ -73,7 +75,7 @@ impl KeyReg {
         pk: VerificationKeyPoP,
     ) -> Result<(), RegisterError> {
         if self.keys.contains(&pk.vk) {
-            return Err(RegisterError::KeyRegistered(pk.vk));
+            return Err(RegisterError::KeyRegistered(Box::new(pk.vk)));
         }
 
         if let Some(mut party) = self.parties.get_mut(&party_id) {
@@ -95,13 +97,17 @@ impl KeyReg {
     where
         D: Digest + FixedOutput,
     {
-        let mut total_stake = 0;
+        let mut total_stake: Stake = 0;
         let mut reg_parties = self
             .parties
             .iter()
             .filter_map(|(_, party)| {
                 if let Some(vk) = party.vk {
-                    total_stake += party.stake;
+                    let (res, overflow) = total_stake.overflowing_add(party.stake);
+                    if overflow {
+                        panic!("Total stake overflow");
+                    }
+                    total_stake = res;
                     return Some(MTLeaf(vk, party.stake));
                 }
                 None
@@ -126,15 +132,19 @@ impl Default for KeyReg {
 #[cfg(test)]
 mod tests {
     use super::*;
+    #[cfg(not(feature = "zcash"))]
     use crate::multi_sig::SigningKey;
+    #[cfg(feature = "zcash")]
+    use crate::multi_sig_zcash::SigningKey;
     use blake2::Blake2b;
     use proptest::collection::vec;
     use proptest::prelude::*;
     use rand_chacha::ChaCha20Rng;
     use rand_core::SeedableRng;
 
     fn arb_participants(min: usize, max: usize) -> impl Strategy<Value = Vec<(PartyId, Stake)>> {
-        vec(any::<Stake>(), min..=max).prop_map(|v| {
+        vec(1..1u64 << 60, min..=max).prop_map(|v| {
+            // 1<<60 to avoid overflows
             v.into_iter()
                 .enumerate()
                 .map(|(index, value)| (index as u64, value))
@@ -184,7 +194,7 @@ mod tests {
                         assert!(parties.insert(p.0));
                     },
                     Err(RegisterError::KeyRegistered(pk1)) => {
-                        assert!(pk1 == pk.vk);
+                        assert!(pk1.as_ref() == &pk.vk);
                         assert!(keys.contains(&pk.vk));
                     }
                     Err(RegisterError::PartyRegistered(party)) => {

mithril-core/src/lib.rs

@@ -7,5 +7,8 @@ mod dense_mapping;
 pub mod error;
 pub mod key_reg;
 pub mod merkle_tree;
+#[cfg(not(feature = "zcash"))]
 mod multi_sig;
+#[cfg(feature = "zcash")]
+mod multi_sig_zcash;
 pub mod stm;

mithril-core/src/merkle_tree.rs

@@ -1,6 +1,9 @@
 //! Creation and verification of Merkle Trees
 use crate::error::MerkleTreeError;
+#[cfg(not(feature = "zcash"))]
 use crate::multi_sig::VerificationKey;
+#[cfg(feature = "zcash")]
+use crate::multi_sig_zcash::VerificationKey;
 use crate::stm::Stake;
 use digest::{Digest, FixedOutput};
 use serde::{Deserialize, Serialize};
@@ -240,11 +243,7 @@ where
     /// * All nodes of the merkle tree (starting with the root)
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut result = Vec::with_capacity(8 + self.nodes.len() * D::output_size());
-        result.extend_from_slice(
-            &u64::try_from(self.n)
-                .expect("Length must fit in u32")
-                .to_be_bytes(),
-        );
+        result.extend_from_slice(&u64::try_from(self.n).unwrap().to_be_bytes());
         for node in self.nodes.iter() {
             result.extend_from_slice(node);
         }
@@ -255,7 +254,8 @@ where
     pub fn from_bytes(bytes: &[u8]) -> Result<Self, MerkleTreeError> {
         let mut u64_bytes = [0u8; 8];
         u64_bytes.copy_from_slice(&bytes[..8]);
-        let n = usize::try_from(u64::from_be_bytes(u64_bytes)).unwrap(); // todo: handle the conversion
+        let n = usize::try_from(u64::from_be_bytes(u64_bytes))
+            .map_err(|_| MerkleTreeError::SerializationError)?;
         let num_nodes = n + n.next_power_of_two() - 1;
         let mut nodes = Vec::with_capacity(num_nodes);
         for i in 0..num_nodes {