mithril_stm/protocol/key_registration/
register.rs

1use digest::{Digest, FixedOutput};
2use std::collections::BTreeSet;
3
4use crate::{
5    RegisterError, SignerIndex, Stake, StmResult, VerificationKeyProofOfPossessionForConcatenation,
6    membership_commitment::{MerkleTree, MerkleTreeLeaf},
7};
8
9use super::RegistrationEntry;
10
11/// Key Registration
12#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Debug)]
13pub struct KeyRegistration {
14    registration_entries: BTreeSet<RegistrationEntry>,
15}
16
17impl KeyRegistration {
18    /// Initialize an empty registration
19    pub fn initialize() -> Self {
20        Self {
21            registration_entries: Default::default(),
22        }
23    }
24
25    /// Check whether the given `RegistrationEntry` is already registered.
26    /// Insert the new entry if all checks pass.
27    /// # Error
28    /// The function fails when the entry is already registered.
29    pub fn register_by_entry(&mut self, entry: &RegistrationEntry) -> StmResult<()> {
30        if !self.registration_entries.contains(entry) {
31            self.registration_entries.insert(*entry);
32            return Ok(());
33        }
34        Err(RegisterError::KeyRegistered(Box::new(entry.get_bls_verification_key())).into())
35    }
36
37    /// Registers a new signer with the given verification key proof of possession and stake.
38    /// This function only works for concatenation proof system.
39    /// The purpose of this function is to simplify the process for the rest of the codebase.
40    pub fn register(
41        &mut self,
42        stake: Stake,
43        vk_pop: &VerificationKeyProofOfPossessionForConcatenation,
44    ) -> StmResult<()> {
45        let entry = RegistrationEntry::new(*vk_pop, stake)?;
46        self.register_by_entry(&entry)
47    }
48
49    /// Gets the index of a signer registration entry
50    pub fn get_signer_index_for_registration(
51        &self,
52        entry: &RegistrationEntry,
53    ) -> Option<SignerIndex> {
54        self.registration_entries
55            .iter()
56            .position(|r| r == entry)
57            .map(|s| s as u64)
58    }
59
60    /// Get the registration entry for a given signer index.
61    pub fn get_registration_entry_for_index(
62        &self,
63        signer_index: &SignerIndex,
64    ) -> StmResult<RegistrationEntry> {
65        self.registration_entries
66            .iter()
67            .nth(*signer_index as usize)
68            .cloned()
69            .ok_or_else(|| RegisterError::UnregisteredIndex.into())
70    }
71
72    /// Converts the KeyRegistration into a Merkle tree
73    pub fn into_merkle_tree<
74        D: Digest + FixedOutput,
75        L: From<RegistrationEntry> + MerkleTreeLeaf,
76    >(
77        &self,
78    ) -> MerkleTree<D, L> {
79        MerkleTree::new(
80            &self
81                .registration_entries
82                .iter()
83                .map(|entry| (*entry).into())
84                .collect::<Vec<L>>(),
85        )
86    }
87
88    /// Closes the registration by computing the total stake registered.
89    pub fn close_registration(self) -> ClosedKeyRegistration {
90        let total_stake: Stake = self.registration_entries.iter().fold(0, |acc, entry| {
91            let (res, overflow) = acc.overflowing_add(entry.get_stake());
92            if overflow {
93                panic!(
94                    "Total stake overflow accumulated stake: {}, adding stake: {}",
95                    acc,
96                    entry.get_stake()
97                );
98            }
99            res
100        });
101        ClosedKeyRegistration {
102            key_registration: self,
103            total_stake,
104        }
105    }
106}
107
108/// Closed Key Registration
109#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Debug)]
110pub struct ClosedKeyRegistration {
111    /// The key registration entries
112    pub key_registration: KeyRegistration,
113
114    /// The total stake registered
115    pub total_stake: Stake,
116}
117
118#[cfg(test)]
119mod tests {
120    use proptest::{collection::vec, prelude::*};
121    use rand_chacha::ChaCha20Rng;
122    use rand_core::SeedableRng;
123
124    use crate::{
125        VerificationKeyProofOfPossessionForConcatenation, signature_scheme::BlsSigningKey,
126    };
127
128    use super::*;
129
130    proptest! {
131        #[test]
132        fn test_keyreg(stake in vec(1..1u64 << 60, 2..=10),
133                       nkeys in 2..10_usize,
134                       fake_it in 0..4usize,
135                       seed in any::<[u8;32]>()) {
136            let mut rng = ChaCha20Rng::from_seed(seed);
137            let mut kr = KeyRegistration::initialize();
138
139            let gen_keys = (1..nkeys).map(|_| {
140                let sk = BlsSigningKey::generate(&mut rng);
141                VerificationKeyProofOfPossessionForConcatenation::from(&sk)
142            }).collect::<Vec<_>>();
143
144            let fake_key = {
145                let sk = BlsSigningKey::generate(&mut rng);
146                VerificationKeyProofOfPossessionForConcatenation::from(&sk)
147            };
148
149            // Record successful registrations
150            let mut keys = BTreeSet::new();
151
152            for (i, &stake) in stake.iter().enumerate() {
153                let mut pk = gen_keys[i % gen_keys.len()];
154
155                if fake_it == 0 {
156                    pk.pop = fake_key.pop;
157                }
158
159                let entry_result = RegistrationEntry::new(pk, stake);
160
161                match entry_result {
162                    Ok(entry) => {
163                        let reg = kr.register_by_entry(&entry);
164                        match reg {
165                            Ok(_) => {
166                                assert!(keys.insert(entry));
167                            },
168                            Err(error) => match error.downcast_ref::<RegisterError>(){
169                                Some(RegisterError::KeyRegistered(e1)) => {
170                                    assert!(e1.as_ref() == &entry.get_bls_verification_key());
171                                    assert!(keys.contains(&entry));
172                                },
173                                _ => {panic!("Unexpected error: {error}")}
174                            }
175                        }
176                    },
177                    Err(error) =>  match error.downcast_ref::<RegisterError>(){
178                        Some(RegisterError::KeyInvalid(a)) => {
179                            assert_eq!(fake_it, 0);
180                            assert!(pk.verify_proof_of_possession().is_err());
181                            assert!(a.as_ref() == &pk.vk);
182                        },
183                        _ => {panic!("Unexpected error: {error}")}
184                    }
185                }
186            }
187
188            if !kr.registration_entries.is_empty() {
189                let closed = kr.close_registration();
190                let retrieved_keys = closed.key_registration.registration_entries;
191                assert!(retrieved_keys == keys);
192            }
193        }
194    }
195
196    mod golden {
197        use blake2::{Blake2b, digest::consts::U32};
198
199        use crate::{
200            Initializer, Parameters,
201            membership_commitment::{MerkleTreeBatchCommitment, MerkleTreeConcatenationLeaf},
202        };
203
204        use super::*;
205
206        const GOLDEN_JSON: &str = r#"
207        {
208            "root":[4,3,108,183,145,65,166,69,250,202,51,64,90,232,45,103,56,138,102,63,209,245,81,22,120,16,6,96,140,204,210,55],
209            "nr_leaves":4,
210            "hasher":null
211        }"#;
212
213        fn golden_value() -> MerkleTreeBatchCommitment<Blake2b<U32>, MerkleTreeConcatenationLeaf> {
214            let mut rng = ChaCha20Rng::from_seed([0u8; 32]);
215            let params = Parameters {
216                m: 10,
217                k: 5,
218                phi_f: 0.8,
219            };
220            let number_of_parties = 4;
221
222            let mut key_reg = KeyRegistration::initialize();
223            for stake in 0..number_of_parties {
224                let initializer = Initializer::new(params, stake, &mut rng);
225                key_reg.register_by_entry(&initializer.clone().into()).unwrap();
226            }
227
228            let closed_key_reg: ClosedKeyRegistration = key_reg.close_registration();
229            closed_key_reg
230                .key_registration
231                .into_merkle_tree()
232                .to_merkle_tree_batch_commitment()
233        }
234
235        #[test]
236        fn golden_conversions() {
237            let value = serde_json::from_str(GOLDEN_JSON)
238                .expect("This JSON deserialization should not fail");
239            assert_eq!(golden_value(), value);
240
241            let serialized =
242                serde_json::to_string(&value).expect("This JSON serialization should not fail");
243            let golden_serialized = serde_json::to_string(&golden_value())
244                .expect("This JSON serialization should not fail");
245            assert_eq!(golden_serialized, serialized);
246        }
247    }
248}
mithril_stm/protocol/key_registration/register.rs

mithril_stm/protocol/key_registration/
register.rs
