mithril_stm/protocol/
key_registration.rs

1//! Key registration functionality.
2use std::{
3    collections::{HashMap, hash_map::Entry},
4    sync::Arc,
5};
6
7use anyhow::anyhow;
8
9use crate::{
10    MembershipDigest, RegisterError, Stake, StmResult,
11    membership_commitment::{MerkleTree, MerkleTreeConcatenationLeaf},
12    signature_scheme::{BlsVerificationKey, BlsVerificationKeyProofOfPossession},
13};
14
15/// Stores a registered party with its public key and the associated stake.
16pub type RegisteredParty = MerkleTreeConcatenationLeaf;
17
18/// Struct that collects public keys and stakes of parties.
19/// Each participant (both the signers and the clerks) need to run their own instance of the key registration.
20// todo: replace with KeyReg
21#[derive(Clone, Debug, Default, PartialEq, Eq)]
22pub struct KeyRegistration {
23    keys: HashMap<BlsVerificationKey, Stake>,
24}
25
26impl KeyRegistration {
27    /// Initialize an empty `KeyRegistration`.
28    pub fn init() -> Self {
29        Self::default()
30    }
31
32    /// Verify and register a public key and stake for a particular party.
33    /// # Error
34    /// The function fails when the proof of possession is invalid or when the key is already registered.
35    pub fn register(
36        &mut self,
37        stake: Stake,
38        pk: BlsVerificationKeyProofOfPossession,
39    ) -> StmResult<()> {
40        if let Entry::Vacant(e) = self.keys.entry(pk.vk) {
41            pk.verify_proof_of_possession()
42                .map_err(|_| RegisterError::KeyInvalid(Box::new(pk)))?;
43            e.insert(stake);
44            return Ok(());
45        }
46        Err(anyhow!(RegisterError::KeyRegistered(Box::new(pk.vk))))
47    }
48
49    /// Finalize the key registration.
50    /// This function disables `KeyReg::register`, consumes the instance of `self`, and returns a `ClosedKeyRegistration`.
51    pub fn close<D>(self) -> ClosedKeyRegistration<D>
52    where
53        D: MembershipDigest,
54    {
55        let mut total_stake: Stake = 0;
56        let mut reg_parties = self
57            .keys
58            .iter()
59            .map(|(&vk, &stake)| {
60                let (res, overflow) = total_stake.overflowing_add(stake);
61                if overflow {
62                    panic!("Total stake overflow");
63                }
64                total_stake = res;
65                MerkleTreeConcatenationLeaf(vk, stake)
66            })
67            .collect::<Vec<RegisteredParty>>();
68        reg_parties.sort();
69
70        ClosedKeyRegistration {
71            merkle_tree: Arc::new(MerkleTree::new(&reg_parties)),
72            reg_parties,
73            total_stake,
74        }
75    }
76}
77
78/// Structure generated out of a closed registration containing the registered parties, total stake, and the merkle tree.
79/// One can only get a global `avk` out of a closed key registration.
80#[derive(Clone, Debug, PartialEq, Eq)]
81pub struct ClosedKeyRegistration<D: MembershipDigest> {
82    /// Ordered list of registered parties.
83    pub reg_parties: Vec<RegisteredParty>,
84    /// Total stake of the registered parties.
85    pub total_stake: Stake,
86    /// Unique public key out of the key registration instance.
87    pub merkle_tree: Arc<MerkleTree<D::ConcatenationHash, MerkleTreeConcatenationLeaf>>,
88}
89
90#[cfg(test)]
91mod tests {
92    use proptest::{collection::vec, prelude::*};
93    use rand_chacha::ChaCha20Rng;
94    use rand_core::SeedableRng;
95
96    use crate::{MithrilMembershipDigest, signature_scheme::BlsSigningKey};
97
98    use super::*;
99
100    proptest! {
101        #[test]
102        fn test_keyreg(stake in vec(1..1u64 << 60, 2..=10),
103                       nkeys in 2..10_usize,
104                       fake_it in 0..4usize,
105                       seed in any::<[u8;32]>()) {
106            let mut rng = ChaCha20Rng::from_seed(seed);
107            let mut kr = KeyRegistration::init();
108
109            let gen_keys = (1..nkeys).map(|_| {
110                let sk = BlsSigningKey::generate(&mut rng);
111                BlsVerificationKeyProofOfPossession::from(&sk)
112            }).collect::<Vec<_>>();
113
114            let fake_key = {
115                let sk = BlsSigningKey::generate(&mut rng);
116                BlsVerificationKeyProofOfPossession::from(&sk)
117            };
118
119            // Record successful registrations
120            let mut keys = HashMap::new();
121
122            for (i, &stake) in stake.iter().enumerate() {
123                let mut pk = gen_keys[i % gen_keys.len()];
124
125                if fake_it == 0 {
126                    pk.pop = fake_key.pop;
127                }
128
129                let reg = kr.register(stake, pk);
130
131                match reg {
132                    Ok(_) => {
133                        assert!(keys.insert(pk.vk, stake).is_none());
134                    },
135                    Err(error) =>  match error.downcast_ref::<RegisterError>(){
136                        Some(RegisterError::KeyRegistered(pk1)) => {
137                            assert!(pk1.as_ref() == &pk.vk);
138                            assert!(keys.contains_key(&pk.vk));
139                        },
140                        Some(RegisterError::KeyInvalid(a)) => {
141                            assert_eq!(fake_it, 0);
142                            assert!(a.verify_proof_of_possession().is_err());
143                        },
144                        _ => {panic!("Unexpected error: {error}")}
145                    }
146                }
147            }
148            if !kr.keys.is_empty() {
149                let closed = kr.close::<MithrilMembershipDigest>();
150                let retrieved_keys = closed.reg_parties.iter().map(|r| (r.0, r.1)).collect::<HashMap<_,_>>();
151                assert!(retrieved_keys == keys);
152            }
153        }
154    }
155
156    mod golden {
157        use blake2::{Blake2b, digest::consts::U32};
158
159        use crate::{
160            Initializer, MithrilMembershipDigest, Parameters,
161            membership_commitment::MerkleTreeBatchCommitment,
162        };
163
164        use super::*;
165
166        const GOLDEN_JSON: &str = r#"
167        {
168            "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],
169            "nr_leaves":4,
170            "hasher":null
171        }"#;
172
173        fn golden_value() -> MerkleTreeBatchCommitment<Blake2b<U32>, MerkleTreeConcatenationLeaf> {
174            let mut rng = ChaCha20Rng::from_seed([0u8; 32]);
175            let params = Parameters {
176                m: 10,
177                k: 5,
178                phi_f: 0.8,
179            };
180            let number_of_parties = 4;
181
182            let mut key_reg = KeyRegistration::init();
183            for stake in 0..number_of_parties {
184                let initializer = Initializer::new(params, stake, &mut rng);
185                key_reg.register(initializer.stake, initializer.pk).unwrap();
186            }
187
188            let closed_key_reg: ClosedKeyRegistration<MithrilMembershipDigest> = key_reg.close();
189            closed_key_reg.merkle_tree.to_merkle_tree_batch_commitment()
190        }
191
192        #[test]
193        fn golden_conversions() {
194            let value = serde_json::from_str(GOLDEN_JSON)
195                .expect("This JSON deserialization should not fail");
196            assert_eq!(golden_value(), value);
197
198            let serialized =
199                serde_json::to_string(&value).expect("This JSON serialization should not fail");
200            let golden_serialized = serde_json::to_string(&golden_value())
201                .expect("This JSON serialization should not fail");
202            assert_eq!(golden_serialized, serialized);
203        }
204    }
205}
mithril_stm/protocol/key_registration.rs

mithril_stm/protocol/
key_registration.rs
