refactor(stm): update PoP, protocol parameters and decoding functions

mithril-stm/benches/size_benches.rs

@@ -32,7 +32,14 @@ where
     let mut key_reg = KeyRegistration::initialize();
     for stake in parties {
         let p = Initializer::new(params, stake, &mut rng);
-        key_reg.register_by_entry(&p.clone().try_into().unwrap()).unwrap();
+        key_reg
+            .register(
+                stake,
+                &p.get_verification_key_proof_of_possession_for_concatenation(),
+                #[cfg(feature = "future_snark")]
+                p.schnorr_verification_key,
+            )
+            .unwrap();
         ps.push(p);
     }
 

mithril-stm/benches/stm.rs

@@ -44,7 +44,14 @@ fn stm_benches<D: MembershipDigest>(
             // We need to initialise the key_reg at each iteration
             key_reg = KeyRegistration::initialize();
             for p in initializers.iter() {
-                key_reg.register_by_entry(&p.clone().try_into().unwrap()).unwrap();
+                key_reg
+                    .register(
+                        p.stake,
+                        &p.get_verification_key_proof_of_possession_for_concatenation(),
+                        #[cfg(feature = "future_snark")]
+                        p.schnorr_verification_key,
+                    )
+                    .unwrap();
             }
         })
     });
@@ -136,7 +143,14 @@ fn batch_benches<D>(
             }
             let mut key_reg = KeyRegistration::initialize();
             for p in initializers.iter() {
-                key_reg.register_by_entry(&p.clone().try_into().unwrap()).unwrap();
+                key_reg
+                    .register(
+                        p.stake,
+                        &p.get_verification_key_proof_of_possession_for_concatenation(),
+                        #[cfg(feature = "future_snark")]
+                        p.schnorr_verification_key,
+                    )
+                    .unwrap();
             }
 
             let closed_reg = key_reg.close_registration(&params).unwrap();

mithril-stm/examples/key_registration.rs

@@ -7,7 +7,7 @@ use rand_core::{RngCore, SeedableRng};
 use mithril_stm::{
     AggregateSignature, AggregateSignatureType, AncillaryGenesisData, AncillaryProofInput, Clerk,
     ClosedKeyRegistration, Initializer, KeyRegistration, MithrilMembershipDigest, Parameters,
-    RegistrationEntry, Stake, VerificationKeyProofOfPossessionForConcatenation,
+    Stake, VerificationKeyProofOfPossessionForConcatenation,
 };
 
 type D = MithrilMembershipDigest;
@@ -236,14 +236,14 @@ fn local_reg(
     };
     // data, such as the public key, stake and id.
     for (pk, _) in pks.iter().zip(ids.iter()) {
-        let entry = RegistrationEntry::new(
-            *pk,
-            1,
-            #[cfg(feature = "future_snark")]
-            None,
-        )
-        .unwrap();
-        local_keyreg.register_by_entry(&entry).unwrap();
+        local_keyreg
+            .register(
+                1,
+                pk,
+                #[cfg(feature = "future_snark")]
+                None,
+            )
+            .unwrap();
     }
     local_keyreg.close_registration(&params).unwrap()
 }

mithril-stm/src/lib.rs

@@ -15,7 +15,7 @@
 //!
 //! use mithril_stm::{
 //!    AggregateSignatureType, AggregationError, AncillaryGenesisData, AncillaryProofInput, Clerk,
-//!    Initializer, KeyRegistration, Parameters, RegistrationEntry, Signer, SingleSignature,
+//!    Initializer, KeyRegistration, Parameters, Signer, SingleSignature,
 //!    MithrilMembershipDigest,
 //! };
 //!
@@ -59,13 +59,12 @@
 //!     // Create keys for this party
 //!     let p = Initializer::new(params, stake, &mut rng);
 //!     // Register keys with the KeyRegistration service
-//!     let entry = RegistrationEntry::new(
-//!         p.get_verification_key_proof_of_possession_for_concatenation(),
+//!     key_reg.register(
 //!         p.stake,
+//!         &p.get_verification_key_proof_of_possession_for_concatenation(),
 //!        #[cfg(feature = "future_snark")] p.schnorr_verification_key,
 //!     )
 //!     .unwrap();
-//!     key_reg.register_by_entry(&entry).unwrap();
 //!     ps.push(p);
 //! }
 //!
@@ -152,11 +151,12 @@ pub use protocol::{
     AggregateSignature, AggregateSignatureError, AggregateSignatureType, AggregateVerificationKey,
     AggregationError, AncillaryGenesisData, AncillaryProofInput, AncillaryProofOutput,
     AncillaryProverData, AncillaryVerifierData, Clerk, ClosedKeyRegistration,
-    ClosedRegistrationEntry, GenesisVerificationKeyBundle, Initializer, KeyRegistration,
-    Parameters, RegisterError, RegistrationEntry, RegistrationEntryForConcatenation,
-    SignatureError, Signer, SingleSignature, SingleSignatureWithRegisteredParty,
-    VerificationKeyForConcatenation, VerificationKeyProofOfPossessionForConcatenation,
+    GenesisVerificationKeyBundle, Initializer, KeyRegistration, Parameters, RegisterError,
+    RegistrationEntryForConcatenation, SignatureError, Signer, SingleSignature,
+    SingleSignatureWithRegisteredParty, VerificationKeyForConcatenation,
+    VerificationKeyProofOfPossessionForConcatenation,
 };
+pub(crate) use protocol::{ClosedRegistrationEntry, RegistrationEntry};
 pub use signature_scheme::BlsSignatureError;
 
 use blake2::{Blake2b, digest::consts::U32};

mithril-stm/src/membership_commitment/merkle_tree/commitment.rs

@@ -176,12 +176,23 @@ impl<D: Digest + FixedOutput, L: MerkleTreeLeaf> MerkleTreeBatchCommitment<D, L>
             )));
         }
 
-        let nr_nodes = self.nr_leaves + self.nr_leaves.next_power_of_two() - 1;
+        let nr_nodes = self
+            .nr_leaves
+            .checked_next_power_of_two()
+            .and_then(|nr_nodes| nr_nodes.checked_add(self.nr_leaves))
+            .and_then(|nr_nodes| nr_nodes.checked_sub(1))
+            .ok_or(MerkleTreeError::SerializationError)?;
 
         ordered_indices = ordered_indices
             .into_iter()
-            .map(|i| i + self.nr_leaves.next_power_of_two() - 1)
-            .collect();
+            .map(|i| {
+                self.nr_leaves
+                    .checked_next_power_of_two()
+                    .and_then(|idx| idx.checked_add(i))
+                    .and_then(|idx| idx.checked_sub(1))
+                    .ok_or(MerkleTreeError::SerializationError)
+            })
+            .collect::<Result<Vec<_>, _>>()?;
 
         let mut idx = ordered_indices[0];
         // First we need to hash the leave values

mithril-stm/src/membership_commitment/merkle_tree/tree.rs

@@ -3,12 +3,14 @@ use std::marker::PhantomData;
 use digest::{Digest, FixedOutput};
 use serde::{Deserialize, Serialize};
 
-use crate::StmResult;
-use crate::codec;
+#[cfg(test)]
+use crate::{StmResult, codec};
 
+#[cfg(test)]
+use super::MerkleTreeError;
 use super::{
-    MerkleBatchPath, MerkleTreeBatchCommitment, MerkleTreeError, MerkleTreeLeaf, left_child,
-    parent, right_child, sibling,
+    MerkleBatchPath, MerkleTreeBatchCommitment, MerkleTreeLeaf, left_child, parent, right_child,
+    sibling,
 };
 #[cfg(feature = "future_snark")]
 // TODO: remove this allow dead_code directive when function is called or future_snark is activated
@@ -147,13 +149,15 @@ impl<D: Digest + FixedOutput, L: MerkleTreeLeaf> MerkleTree<D, L> {
     }
 
     /// Convert a `MerkleTree` into a byte string.
+    #[cfg(test)]
     pub fn to_bytes(&self) -> StmResult<Vec<u8>> {
         codec::to_cbor_bytes(self)
     }
 
     /// Try to convert a byte string into a `MerkleTree`.
     /// # Error
     /// It returns error if conversion fails.
+    #[cfg(test)]
     pub fn from_bytes(bytes: &[u8]) -> StmResult<Self> {
         codec::from_versioned_bytes(bytes, Self::from_bytes_legacy)
     }
@@ -162,27 +166,38 @@ impl<D: Digest + FixedOutput, L: MerkleTreeLeaf> MerkleTree<D, L> {
     /// # Layout
     /// * Number of leaves committed in the Merkle Tree (as u64)
     /// * All nodes of the merkle tree (starting with the root)
+    #[cfg(test)]
     fn from_bytes_legacy(bytes: &[u8]) -> StmResult<Self> {
         let mut u64_bytes = [0u8; 8];
         u64_bytes.copy_from_slice(bytes.get(..8).ok_or(MerkleTreeError::SerializationError)?);
         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);
+        let num_nodes = n
+            .checked_next_power_of_two()
+            .and_then(|num_nodes| num_nodes.checked_add(n))
+            .and_then(|num_nodes| num_nodes.checked_sub(1))
+            .ok_or(MerkleTreeError::SerializationError)?;
+        let mut nodes = Vec::new();
         for i in 0..num_nodes {
+            let range_low = i
+                .checked_mul(<D as Digest>::output_size())
+                .and_then(|rl| rl.checked_add(16))
+                .ok_or(MerkleTreeError::SerializationError)?;
+            let range_high = i
+                .checked_add(1)
+                .and_then(|rh| rh.checked_mul(<D as Digest>::output_size()))
+                .and_then(|rh| rh.checked_add(16))
+                .ok_or(MerkleTreeError::SerializationError)?;
             nodes.push(
                 bytes
-                    .get(
-                        8 + i * <D as Digest>::output_size()
-                            ..8 + (i + 1) * <D as Digest>::output_size(),
-                    )
+                    .get(range_low..range_high)
                     .ok_or(MerkleTreeError::SerializationError)?
                     .to_vec(),
             );
         }
         Ok(Self {
             nodes,
-            leaf_off: num_nodes - n,
+            leaf_off: num_nodes.checked_sub(n).ok_or(MerkleTreeError::SerializationError)?,
             n,
             hasher: PhantomData,
             leaves: PhantomData,

mithril-stm/src/proof_system/concatenation/aggregate_key.rs

@@ -77,7 +77,7 @@ impl<D: MembershipDigest> From<&ClosedKeyRegistration>
             mt_commitment: reg
                 .to_merkle_tree::<D::ConcatenationHash, RegistrationEntryForConcatenation>()
                 .to_merkle_tree_batch_commitment(),
-            total_stake: reg.total_stake,
+            total_stake: reg.get_total_stake(),
         }
     }
 }

mithril-stm/src/proof_system/concatenation/eligibility.rs

@@ -1,4 +1,6 @@
-use crate::{PhiFValue, Stake};
+use anyhow::anyhow;
+
+use crate::{PhiFValue, Stake, StmResult};
 
 cfg_num_integer! {
     use num_bigint::{BigInt, Sign};
@@ -29,23 +31,25 @@ cfg_num_integer! {
     ///                 1 - p    1 - (ev / evMax)    (evMax - ev)
     ///
     /// Used to determine winning lottery tickets.
-    pub(crate) fn is_lottery_won(phi_f: PhiFValue, ev: [u8; 64], stake: Stake, total_stake: Stake) -> bool {
+    pub(crate) fn is_lottery_won(phi_f: PhiFValue, ev: [u8; 64], stake: Stake, total_stake: Stake) -> StmResult<bool> {
         // If phi_f = 1, then we automatically break with true
         if (phi_f - 1.0).abs() < PhiFValue::EPSILON {
-            return true;
+            return Ok(true);
+        } else if !(phi_f > 0.0 && phi_f <= 1.0) {
+            return Err(anyhow!("phi_f must be in the range (0, 1], got {phi_f}"));
         }
 
         let ev_max = BigInt::from(2u8).pow(512);
         let ev = BigInt::from_bytes_le(Sign::Plus, &ev);
         let q = Ratio::new_raw(ev_max.clone(), ev_max - ev);
 
         let c =
-            Ratio::from_float((1.0 - phi_f).ln()).expect("Only fails if the float is infinite or NaN.");
+            Ratio::from_float((1.0 - phi_f).ln()).ok_or(anyhow!("phi_f must not be infinite or NaN, got {phi_f}"))?;
         let w = Ratio::new_raw(BigInt::from(stake), BigInt::from(total_stake));
         let x = (w * c).neg();
 
         // Now we compute a taylor function that breaks when the result is known.
-        taylor_comparison(1000, q, x)
+        Ok(taylor_comparison(1000, q, x))
     }
 
     /// Checks if cmp < exp(x). Uses error approximation for an early stop. Whenever the value being
@@ -92,10 +96,12 @@ cfg_rug! {
     /// order to keep the error in the 1e-17 range, we need to carry out the computations with 34
     /// decimal digits (in order to represent the 4.5e16 ada without any rounding errors, we need
     /// double that precision).
-    pub(crate) fn is_lottery_won(phi_f: PhiFValue, ev: [u8; 64], stake: Stake, total_stake: Stake) -> bool {
+    pub(crate) fn is_lottery_won(phi_f: PhiFValue, ev: [u8; 64], stake: Stake, total_stake: Stake) -> StmResult<bool> {
         // If phi_f = 1, then we automatically break with true
         if (phi_f - 1.0).abs() < PhiFValue::EPSILON {
-            return true;
+            return Ok(true);
+        } else if !(phi_f > 0.0 && phi_f <= 1.0) {
+            return Err(anyhow!("phi_f must be in the range (0, 1], got {phi_f}"));
         }
         let ev = rug::Integer::from_digits(&ev, Order::LsfLe);
         let ev_max: Float = Float::with_val(117, 2).pow(512);
@@ -104,7 +110,7 @@ cfg_rug! {
         let w = Float::with_val(117, stake) / Float::with_val(117, total_stake);
         let phi = Float::with_val(117, 1.0) - Float::with_val(117, 1.0 - phi_f).pow(w);
 
-        q < phi
+        Ok(q < phi)
     }
 }
 
@@ -142,7 +148,7 @@ mod tests {
             ev.copy_from_slice(&[&ev_1[..], &ev_2[..]].concat());
 
             let quick_result = trivial_is_lottery_won(phi_f, ev, stake, total_stake);
-            let result = is_lottery_won(phi_f, ev, stake, total_stake);
+            let result = is_lottery_won(phi_f, ev, stake, total_stake).unwrap();
             assert_eq!(quick_result, result);
         }
 
@@ -159,4 +165,18 @@ mod tests {
             assert!(!taylor_comparison(1000, cmp_p, Ratio::from_float(x).unwrap()));
         }
     }
+
+    #[test]
+    fn is_lottery_won_rejects_out_of_range_phi_f() {
+        let ev = [0u8; 64];
+        let stake = 100;
+        let total_stake = 1000;
+
+        for invalid_phi_f in [-0.5, 0.0, 1.5, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
+            assert!(
+                is_lottery_won(invalid_phi_f, ev, stake, total_stake).is_err(),
+                "phi_f = {invalid_phi_f} should be rejected"
+            );
+        }
+    }
 }

mithril-stm/src/proof_system/concatenation/proof.rs

@@ -282,21 +282,27 @@ impl<D: MembershipDigest> ConcatenationProof<D> {
             .map_err(|_| AggregateSignatureError::SerializationError)?;
         bytes_index += 8;
 
-        let mut sig_reg_list = Vec::with_capacity(total_sigs);
+        let mut sig_reg_list = Vec::new();
         for _ in 0..total_sigs {
+            let size_field_end = bytes_index
+                .checked_add(8)
+                .ok_or(AggregateSignatureError::SerializationError)?;
             u64_bytes.copy_from_slice(
                 bytes
-                    .get(bytes_index..bytes_index + 8)
+                    .get(bytes_index..size_field_end)
                     .ok_or(AggregateSignatureError::SerializationError)?,
             );
             let sig_reg_size = usize::try_from(u64::from_be_bytes(u64_bytes))
                 .map_err(|_| AggregateSignatureError::SerializationError)?;
+            let sig_reg_end = size_field_end
+                .checked_add(sig_reg_size)
+                .ok_or(AggregateSignatureError::SerializationError)?;
             let sig_reg = SingleSignatureWithRegisteredParty::from_bytes::<D>(
                 bytes
-                    .get(bytes_index + 8..bytes_index + 8 + sig_reg_size)
+                    .get(size_field_end..sig_reg_end)
                     .ok_or(AggregateSignatureError::SerializationError)?,
             )?;
-            bytes_index += 8 + sig_reg_size;
+            bytes_index = sig_reg_end;
             sig_reg_list.push(sig_reg);
         }
 

mithril-stm/src/proof_system/concatenation/signer.rs

@@ -60,7 +60,7 @@ impl<D: MembershipDigest> ConcatenationProofSigner<D> {
             self.key_registration_commitment.concatenate_with_message(message);
 
         let sigma = self.signing_key.sign(&message_with_commitment);
-        let indices = self.check_lottery(&message_with_commitment, &sigma);
+        let indices = self.check_lottery(&message_with_commitment, &sigma)?;
 
         if indices.is_empty() {
             Err(anyhow!(SignatureError::LotteryLost))
@@ -71,19 +71,23 @@ impl<D: MembershipDigest> ConcatenationProofSigner<D> {
 
     /// Checks the lottery for given `message_with_commitment` and the `sigma`.
     /// Returns a vector of winning indices.
-    pub fn check_lottery(&self, message_with_commitment: &[u8], sigma: &BlsSignature) -> Vec<u64> {
+    pub fn check_lottery(
+        &self,
+        message_with_commitment: &[u8],
+        sigma: &BlsSignature,
+    ) -> StmResult<Vec<u64>> {
         let mut indices = Vec::new();
         for index in 0..self.parameters.m {
             if is_lottery_won(
                 self.parameters.phi_f,
                 sigma.evaluate_dense_mapping(message_with_commitment, index),
                 self.stake,
                 self.total_stake,
-            ) {
+            )? {
                 indices.push(index);
             }
         }
-        indices
+        Ok(indices)
     }
 
     pub fn get_verification_key(&self) -> VerificationKeyForConcatenation {