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

mithril_stm/protocol/key_registration/
register.rs
