mithril_common/test_utils/
fixture_builder.rs1use 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
16pub 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
41pub enum StakeDistributionGenerationMethod {
43 RandomDistribution {
45 seed: [u8; 32],
47 min_stake: Stake,
49 },
50
51 Custom(StakeDistribution),
55
56 Uniform(Stake),
58}
59
60impl MithrilFixtureBuilder {
61 pub fn with_protocol_parameters(mut self, protocol_parameters: ProtocolParameters) -> Self {
63 self.protocol_parameters = protocol_parameters;
64 self
65 }
66
67 pub fn with_signers(mut self, number_of_signers: usize) -> Self {
69 self.number_of_signers = number_of_signers;
70 self
71 }
72
73 pub fn disable_signers_certification(mut self) -> Self {
76 self.enable_signers_certification = false;
77 self
78 }
79
80 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 pub fn with_party_id_seed(mut self, seed: [u8; 32]) -> Self {
91 self.party_id_seed = seed;
92 self
93 }
94
95 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 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 #[test]
328 fn verify_kes_key_cache_content() {
329 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}