mithril_common/test_utils/
fixture_builder.rs

1use kes_summed_ed25519::{kes::Sum6Kes, traits::KesSk, PublicKey as KesPublicKey};
2use rand_chacha::ChaCha20Rng;
3use rand_core::{RngCore, SeedableRng};
4
5use crate::{
6    crypto_helper::{
7        tests_setup, tests_setup::setup_temp_directory_for_signer, ColdKeyGenerator, OpCert,
8        ProtocolStakeDistribution, SerDeShelleyFileFormat, Sum6KesBytes,
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 {:?}.",
160                kes_key_seed
161            );
162            MithrilFixtureBuilder::generate_kes_key(kes_key_seed)
163        }
164    }
165
166    fn cached_kes_key(kes_key_seed: &[u8]) -> Option<(Sum6KesBytes, KesPublicKey)> {
167        precomputed_kes_key::cached_kes_key(kes_key_seed).map(
168            |(kes_bytes, kes_verification_key)| {
169                let kes_verification_key = KesPublicKey::from_bytes(&kes_verification_key).unwrap();
170                let kes_bytes = Sum6KesBytes(kes_bytes);
171
172                (kes_bytes, kes_verification_key)
173            },
174        )
175    }
176
177    fn generate_kes_key(kes_key_seed: &mut [u8]) -> (Sum6KesBytes, KesPublicKey) {
178        let mut key_buffer = [0u8; Sum6Kes::SIZE + 4];
179
180        let (kes_secret_key, kes_verification_key) = Sum6Kes::keygen(&mut key_buffer, kes_key_seed);
181        let mut kes_bytes = Sum6KesBytes([0u8; Sum6Kes::SIZE + 4]);
182        kes_bytes.0.copy_from_slice(&kes_secret_key.clone_sk());
183
184        (kes_bytes, kes_verification_key)
185    }
186
187    fn generate_cold_key_seed(&self, party_index: usize) -> Vec<u8> {
188        let mut cold_key_seed: Vec<_> = (party_index)
189            .to_le_bytes()
190            .iter()
191            .zip(self.party_id_seed)
192            .map(|(v1, v2)| v1 + v2)
193            .collect();
194        cold_key_seed.resize(32, 0);
195
196        cold_key_seed
197    }
198
199    fn build_party_with_operational_certificate(&self, party_index: usize) -> PartyId {
200        let cold_key_seed = self.generate_cold_key_seed(party_index).to_vec();
201        let mut kes_key_seed = cold_key_seed.clone();
202
203        let keypair =
204            ColdKeyGenerator::create_deterministic_keypair(cold_key_seed.try_into().unwrap());
205        let (kes_bytes, kes_verification_key) =
206            MithrilFixtureBuilder::provide_kes_key(&mut kes_key_seed);
207        let operational_certificate = OpCert::new(kes_verification_key, 0, 0, keypair);
208        let party_id = operational_certificate
209            .compute_protocol_party_id()
210            .expect("compute protocol party id should not fail");
211        let temp_dir = setup_temp_directory_for_signer(&party_id, true)
212            .expect("setup temp directory should return a value");
213        if !temp_dir.join("kes.sk").exists() {
214            kes_bytes
215                .to_file(temp_dir.join("kes.sk"))
216                .expect("KES secret key file export should not fail");
217        }
218        if !temp_dir.join("opcert.cert").exists() {
219            operational_certificate
220                .to_file(temp_dir.join("opcert.cert"))
221                .expect("operational certificate file export should not fail");
222        }
223        party_id
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use std::collections::BTreeSet;
231
232    #[test]
233    fn with_protocol_params() {
234        let protocol_parameters = ProtocolParameters::new(1, 10, 0.56);
235        let result = MithrilFixtureBuilder::default()
236            .with_protocol_parameters(protocol_parameters.clone())
237            .build();
238
239        assert_eq!(protocol_parameters, result.protocol_parameters());
240    }
241
242    #[test]
243    fn with_signers() {
244        let result = MithrilFixtureBuilder::default().with_signers(4).build();
245
246        assert_eq!(4, result.signers_with_stake().len());
247    }
248
249    #[test]
250    fn random_stake_distribution_generates_as_many_signers_as_parties() {
251        let result = MithrilFixtureBuilder::default()
252            .with_stake_distribution(StakeDistributionGenerationMethod::RandomDistribution {
253                seed: [0u8; 32],
254                min_stake: 1,
255            })
256            .with_signers(4)
257            .build();
258
259        assert_eq!(4, result.stake_distribution().len());
260    }
261
262    #[test]
263    fn uniform_stake_distribution() {
264        let expected_stake = 10;
265        let stake_distribution = MithrilFixtureBuilder::default()
266            .with_stake_distribution(StakeDistributionGenerationMethod::Uniform(expected_stake))
267            .with_signers(5)
268            .build()
269            .stake_distribution();
270
271        assert!(
272            stake_distribution
273                .iter()
274                .all(|(_, stake)| *stake == expected_stake),
275            "Generated stake distribution doesn't have uniform stakes: {stake_distribution:?}"
276        );
277    }
278
279    #[test]
280    fn each_parties_generated_with_random_stake_distribution_have_different_stakes() {
281        let result = MithrilFixtureBuilder::default()
282            .with_stake_distribution(StakeDistributionGenerationMethod::RandomDistribution {
283                seed: [0u8; 32],
284                min_stake: 1,
285            })
286            .with_signers(5)
287            .build();
288        let stakes = result.stake_distribution();
289
290        // BtreeSet dedup values
291        assert_eq!(stakes.len(), BTreeSet::from_iter(stakes.values()).len());
292    }
293
294    #[test]
295    fn dont_generate_party_ids_for_custom_stake_distribution() {
296        let stake_distribution = StakeDistribution::from_iter([("party".to_owned(), 4)]);
297        let builder = MithrilFixtureBuilder::default()
298            .with_stake_distribution(StakeDistributionGenerationMethod::Custom(
299                stake_distribution,
300            ))
301            .with_signers(5);
302
303        assert_eq!(Vec::<PartyId>::new(), builder.generate_party_ids());
304    }
305
306    #[test]
307    fn changing_party_id_seed_change_all_builded_party_ids() {
308        let first_signers = MithrilFixtureBuilder::default()
309            .with_signers(10)
310            .build()
311            .signers_with_stake();
312        let different_party_id_seed_signers = MithrilFixtureBuilder::default()
313            .with_signers(10)
314            .with_party_id_seed([1u8; 32])
315            .build()
316            .signers_with_stake();
317        let first_party_ids: Vec<&PartyId> = first_signers.iter().map(|s| &s.party_id).collect();
318
319        for party_id in different_party_id_seed_signers.iter().map(|s| &s.party_id) {
320            assert!(!first_party_ids.contains(&party_id));
321        }
322    }
323
324    /// Verify that there is a cached kes key for a number of party id.
325    /// 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.
326    /// The number of party id that should be in cache is defined with `precomputed_number`
327    #[test]
328    fn verify_kes_key_cache_content() {
329        // Generate code that should be in the `cached_kes_key` function of the `precomputed_kes_key.rs` file.
330        // It can be copied and pasted to update the cache.
331        fn generate_code(party_ids: &Vec<(&[u8], [u8; 612], KesPublicKey)>) -> String {
332            party_ids
333                .iter()
334                .map(|(key, i, p)| format!("{:?} => ({:?}, {:?}),", key, i, p.as_bytes()))
335                .collect::<Vec<_>>()
336                .join("\n")
337        }
338
339        let precomputed_number = 10;
340
341        let fixture = MithrilFixtureBuilder::default();
342        let cold_keys: Vec<_> = (0..precomputed_number)
343            .map(|party_index| fixture.generate_cold_key_seed(party_index))
344            .collect();
345
346        let computed_keys_key: Vec<_> = cold_keys
347            .iter()
348            .map(|cold_key| {
349                let mut kes_key_seed: Vec<u8> = cold_key.clone();
350                let (kes_bytes, kes_verification_key) =
351                    MithrilFixtureBuilder::generate_kes_key(&mut kes_key_seed);
352
353                (cold_key.as_slice(), kes_bytes.0, kes_verification_key)
354            })
355            .collect();
356
357        let cached_kes_key: Vec<_> = cold_keys
358            .iter()
359            .filter_map(|cold_key| {
360                MithrilFixtureBuilder::cached_kes_key(cold_key).map(
361                    |(kes_bytes, kes_verification_key)| {
362                        (cold_key.as_slice(), kes_bytes.0, kes_verification_key)
363                    },
364                )
365            })
366            .collect();
367
368        let expected_code = generate_code(&computed_keys_key);
369        let actual_code = generate_code(&cached_kes_key);
370
371        assert_eq!(
372            computed_keys_key, cached_kes_key,
373            "Precomputed KES keys should be:\n{}\nbut seems to be:\n{}",
374            expected_code, actual_code
375        );
376
377        let kes_key_seed = fixture.generate_cold_key_seed(precomputed_number);
378        assert!(
379            MithrilFixtureBuilder::cached_kes_key(kes_key_seed.as_slice()).is_none(),
380            "We checked precomputed KES keys up to {} but it seems to be more.",
381            precomputed_number
382        );
383    }
384}