mithril_common/test_utils/
fixture_builder.rs

1use kes_summed_ed25519::{PublicKey as KesPublicKey, kes::Sum6Kes, traits::KesSk};
2use rand_chacha::ChaCha20Rng;
3use rand_core::{RngCore, SeedableRng};
4
5use crate::{
6    crypto_helper::{
7        ColdKeyGenerator, OpCert, ProtocolStakeDistribution, SerDeShelleyFileFormat, Sum6KesBytes,
8        tests_setup, tests_setup::setup_temp_directory_for_signer,
9    },
10    entities::{PartyId, ProtocolParameters, Stake, StakeDistribution},
11    test_utils::{fake_data, mithril_fixture::MithrilFixture},
12};
13
14use super::precomputed_kes_key;
15
16/// A builder of mithril types.
17pub struct MithrilFixtureBuilder {
18    protocol_parameters: ProtocolParameters,
19    enable_signers_certification: bool,
20    number_of_signers: usize,
21    stake_distribution_generation_method: StakeDistributionGenerationMethod,
22    party_id_seed: [u8; 32],
23}
24
25impl Default for MithrilFixtureBuilder {
26    fn default() -> Self {
27        Self {
28            protocol_parameters: fake_data::protocol_parameters(),
29            enable_signers_certification: true,
30            number_of_signers: 5,
31            stake_distribution_generation_method:
32                StakeDistributionGenerationMethod::RandomDistribution {
33                    seed: [0u8; 32],
34                    min_stake: 1,
35                },
36            party_id_seed: [0u8; 32],
37        }
38    }
39}
40
41/// Methods that can be used to generate the stake distribution.
42pub enum StakeDistributionGenerationMethod {
43    /// Each party will have a random stake.
44    RandomDistribution {
45        /// The randomizer seed
46        seed: [u8; 32],
47        /// The minimum stake
48        min_stake: Stake,
49    },
50
51    /// Use a custom stake distribution
52    ///
53    /// Important: this will overwrite the number of signers set by with_signers.
54    Custom(StakeDistribution),
55
56    /// Make a stake distribution where all parties will have the given stake
57    Uniform(Stake),
58}
59
60impl MithrilFixtureBuilder {
61    /// Set the protocol_parameters.
62    pub fn with_protocol_parameters(mut self, protocol_parameters: ProtocolParameters) -> Self {
63        self.protocol_parameters = protocol_parameters;
64        self
65    }
66
67    /// Set the number of signers that will be generated.
68    pub fn with_signers(mut self, number_of_signers: usize) -> Self {
69        self.number_of_signers = number_of_signers;
70        self
71    }
72
73    /// If set the generated signers won't be certified (meaning that they won't
74    /// have a operational certificate).
75    pub fn disable_signers_certification(mut self) -> Self {
76        self.enable_signers_certification = false;
77        self
78    }
79
80    /// Set the generation method used to compute the stake distribution.
81    pub fn with_stake_distribution(
82        mut self,
83        stake_distribution_generation_method: StakeDistributionGenerationMethod,
84    ) -> Self {
85        self.stake_distribution_generation_method = stake_distribution_generation_method;
86        self
87    }
88
89    /// Set the seed used to generated the party ids
90    pub fn with_party_id_seed(mut self, seed: [u8; 32]) -> Self {
91        self.party_id_seed = seed;
92        self
93    }
94
95    /// Transform the specified parameters to a [MithrilFixture].
96    pub fn build(self) -> MithrilFixture {
97        let protocol_stake_distribution = self.generate_stake_distribution();
98        let signers = tests_setup::setup_signers_from_stake_distribution(
99            &protocol_stake_distribution,
100            &self.protocol_parameters.clone().into(),
101        );
102
103        MithrilFixture::new(
104            self.protocol_parameters,
105            signers,
106            protocol_stake_distribution,
107        )
108    }
109
110    fn generate_stake_distribution(&self) -> ProtocolStakeDistribution {
111        let signers_party_ids = self.generate_party_ids();
112
113        match &self.stake_distribution_generation_method {
114            StakeDistributionGenerationMethod::RandomDistribution { seed, min_stake } => {
115                let mut stake_rng = ChaCha20Rng::from_seed(*seed);
116
117                signers_party_ids
118                    .into_iter()
119                    .map(|party_id| {
120                        let stake = min_stake + stake_rng.next_u64() % 999;
121                        (party_id, stake)
122                    })
123                    .collect::<Vec<_>>()
124            }
125            StakeDistributionGenerationMethod::Custom(stake_distribution) => stake_distribution
126                .clone()
127                .into_iter()
128                .collect::<ProtocolStakeDistribution>(),
129            StakeDistributionGenerationMethod::Uniform(stake) => signers_party_ids
130                .into_iter()
131                .map(|party_id| (party_id, *stake))
132                .collect::<ProtocolStakeDistribution>(),
133        }
134    }
135
136    fn generate_party_ids(&self) -> Vec<PartyId> {
137        match self.stake_distribution_generation_method {
138            StakeDistributionGenerationMethod::Custom(_) => vec![],
139            _ => {
140                let signers_party_ids = (0..self.number_of_signers).map(|party_index| {
141                    if self.enable_signers_certification {
142                        self.build_party_with_operational_certificate(party_index)
143                    } else {
144                        party_index.to_string()
145                    }
146                });
147                signers_party_ids.collect::<Vec<_>>()
148            }
149        }
150    }
151
152    fn provide_kes_key(kes_key_seed: &mut [u8]) -> (Sum6KesBytes, KesPublicKey) {
153        if let Some((kes_bytes, kes_verification_key)) =
154            MithrilFixtureBuilder::cached_kes_key(kes_key_seed)
155        {
156            (kes_bytes, kes_verification_key)
157        } else {
158            println!(
159                "KES key not found in test cache, generating a new one for the seed {kes_key_seed:?}."
160            );
161            MithrilFixtureBuilder::generate_kes_key(kes_key_seed)
162        }
163    }
164
165    fn cached_kes_key(kes_key_seed: &[u8]) -> Option<(Sum6KesBytes, KesPublicKey)> {
166        precomputed_kes_key::cached_kes_key(kes_key_seed).map(
167            |(kes_bytes, kes_verification_key)| {
168                let kes_verification_key = KesPublicKey::from_bytes(&kes_verification_key).unwrap();
169                let kes_bytes = Sum6KesBytes(kes_bytes);
170
171                (kes_bytes, kes_verification_key)
172            },
173        )
174    }
175
176    fn generate_kes_key(kes_key_seed: &mut [u8]) -> (Sum6KesBytes, KesPublicKey) {
177        let mut key_buffer = [0u8; Sum6Kes::SIZE + 4];
178
179        let (kes_secret_key, kes_verification_key) = Sum6Kes::keygen(&mut key_buffer, kes_key_seed);
180        let mut kes_bytes = Sum6KesBytes([0u8; Sum6Kes::SIZE + 4]);
181        kes_bytes.0.copy_from_slice(&kes_secret_key.clone_sk());
182
183        (kes_bytes, kes_verification_key)
184    }
185
186    fn generate_cold_key_seed(&self, party_index: usize) -> Vec<u8> {
187        let mut cold_key_seed: Vec<_> = (party_index)
188            .to_le_bytes()
189            .iter()
190            .zip(self.party_id_seed)
191            .map(|(v1, v2)| v1 + v2)
192            .collect();
193        cold_key_seed.resize(32, 0);
194
195        cold_key_seed
196    }
197
198    fn build_party_with_operational_certificate(&self, party_index: usize) -> PartyId {
199        let cold_key_seed = self.generate_cold_key_seed(party_index).to_vec();
200        let mut kes_key_seed = cold_key_seed.clone();
201
202        let keypair =
203            ColdKeyGenerator::create_deterministic_keypair(cold_key_seed.try_into().unwrap());
204        let (kes_bytes, kes_verification_key) =
205            MithrilFixtureBuilder::provide_kes_key(&mut kes_key_seed);
206        let operational_certificate = OpCert::new(kes_verification_key, 0, 0, keypair);
207        let party_id = operational_certificate
208            .compute_protocol_party_id()
209            .expect("compute protocol party id should not fail");
210        let temp_dir = setup_temp_directory_for_signer(&party_id, true)
211            .expect("setup temp directory should return a value");
212        if !temp_dir.join("kes.sk").exists() {
213            kes_bytes
214                .to_file(temp_dir.join("kes.sk"))
215                .expect("KES secret key file export should not fail");
216        }
217        if !temp_dir.join("opcert.cert").exists() {
218            operational_certificate
219                .to_file(temp_dir.join("opcert.cert"))
220                .expect("operational certificate file export should not fail");
221        }
222        party_id
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use std::collections::BTreeSet;
230
231    #[test]
232    fn with_protocol_params() {
233        let protocol_parameters = ProtocolParameters::new(1, 10, 0.56);
234        let result = MithrilFixtureBuilder::default()
235            .with_protocol_parameters(protocol_parameters.clone())
236            .build();
237
238        assert_eq!(protocol_parameters, result.protocol_parameters());
239    }
240
241    #[test]
242    fn with_signers() {
243        let result = MithrilFixtureBuilder::default().with_signers(4).build();
244
245        assert_eq!(4, result.signers_with_stake().len());
246    }
247
248    #[test]
249    fn random_stake_distribution_generates_as_many_signers_as_parties() {
250        let result = MithrilFixtureBuilder::default()
251            .with_stake_distribution(StakeDistributionGenerationMethod::RandomDistribution {
252                seed: [0u8; 32],
253                min_stake: 1,
254            })
255            .with_signers(4)
256            .build();
257
258        assert_eq!(4, result.stake_distribution().len());
259    }
260
261    #[test]
262    fn uniform_stake_distribution() {
263        let expected_stake = 10;
264        let stake_distribution = MithrilFixtureBuilder::default()
265            .with_stake_distribution(StakeDistributionGenerationMethod::Uniform(expected_stake))
266            .with_signers(5)
267            .build()
268            .stake_distribution();
269
270        assert!(
271            stake_distribution.iter().all(|(_, stake)| *stake == expected_stake),
272            "Generated stake distribution doesn't have uniform stakes: {stake_distribution:?}"
273        );
274    }
275
276    #[test]
277    fn each_parties_generated_with_random_stake_distribution_have_different_stakes() {
278        let result = MithrilFixtureBuilder::default()
279            .with_stake_distribution(StakeDistributionGenerationMethod::RandomDistribution {
280                seed: [0u8; 32],
281                min_stake: 1,
282            })
283            .with_signers(5)
284            .build();
285        let stakes = result.stake_distribution();
286
287        // BtreeSet dedup values
288        assert_eq!(stakes.len(), BTreeSet::from_iter(stakes.values()).len());
289    }
290
291    #[test]
292    fn dont_generate_party_ids_for_custom_stake_distribution() {
293        let stake_distribution = StakeDistribution::from_iter([("party".to_owned(), 4)]);
294        let builder = MithrilFixtureBuilder::default()
295            .with_stake_distribution(StakeDistributionGenerationMethod::Custom(
296                stake_distribution,
297            ))
298            .with_signers(5);
299
300        assert_eq!(Vec::<PartyId>::new(), builder.generate_party_ids());
301    }
302
303    #[test]
304    fn changing_party_id_seed_change_all_builded_party_ids() {
305        let first_signers = MithrilFixtureBuilder::default()
306            .with_signers(10)
307            .build()
308            .signers_with_stake();
309        let different_party_id_seed_signers = MithrilFixtureBuilder::default()
310            .with_signers(10)
311            .with_party_id_seed([1u8; 32])
312            .build()
313            .signers_with_stake();
314        let first_party_ids: Vec<&PartyId> = first_signers.iter().map(|s| &s.party_id).collect();
315
316        for party_id in different_party_id_seed_signers.iter().map(|s| &s.party_id) {
317            assert!(!first_party_ids.contains(&party_id));
318        }
319    }
320
321    /// Verify that there is a cached kes key for a number of party id.
322    /// If the cache is not up to date, the test will generate the code that can be copied/pasted into the [precomputed_kes_key] module.
323    /// The number of party id that should be in cache is defined with `precomputed_number`
324    #[test]
325    fn verify_kes_key_cache_content() {
326        // Generate code that should be in the `cached_kes_key` function of the `precomputed_kes_key.rs` file.
327        // It can be copied and pasted to update the cache.
328        fn generate_code(party_ids: &Vec<(&[u8], [u8; 612], KesPublicKey)>) -> String {
329            party_ids
330                .iter()
331                .map(|(key, i, p)| format!("{:?} => ({:?}, {:?}),", key, i, p.as_bytes()))
332                .collect::<Vec<_>>()
333                .join("\n")
334        }
335
336        let precomputed_number = 10;
337
338        let fixture = MithrilFixtureBuilder::default();
339        let cold_keys: Vec<_> = (0..precomputed_number)
340            .map(|party_index| fixture.generate_cold_key_seed(party_index))
341            .collect();
342
343        let computed_keys_key: Vec<_> = cold_keys
344            .iter()
345            .map(|cold_key| {
346                let mut kes_key_seed: Vec<u8> = cold_key.clone();
347                let (kes_bytes, kes_verification_key) =
348                    MithrilFixtureBuilder::generate_kes_key(&mut kes_key_seed);
349
350                (cold_key.as_slice(), kes_bytes.0, kes_verification_key)
351            })
352            .collect();
353
354        let cached_kes_key: Vec<_> = cold_keys
355            .iter()
356            .filter_map(|cold_key| {
357                MithrilFixtureBuilder::cached_kes_key(cold_key).map(
358                    |(kes_bytes, kes_verification_key)| {
359                        (cold_key.as_slice(), kes_bytes.0, kes_verification_key)
360                    },
361                )
362            })
363            .collect();
364
365        let expected_code = generate_code(&computed_keys_key);
366        let actual_code = generate_code(&cached_kes_key);
367
368        assert_eq!(
369            computed_keys_key, cached_kes_key,
370            "Precomputed KES keys should be:\n{expected_code}\nbut seems to be:\n{actual_code}"
371        );
372
373        let kes_key_seed = fixture.generate_cold_key_seed(precomputed_number);
374        assert!(
375            MithrilFixtureBuilder::cached_kes_key(kes_key_seed.as_slice()).is_none(),
376            "We checked precomputed KES keys up to {precomputed_number} but it seems to be more."
377        );
378    }
379}