mithril_common/crypto_helper/
tests_setup.rs

1//! Test data builders for Mithril STM types, for testing purpose.
2use super::{ed25519_alias::genesis::*, types::*, OpCert, SerDeShelleyFileFormat};
3use crate::{
4    entities::{Certificate, ProtocolMessage, ProtocolMessagePartKey, SignerWithStake, Stake},
5    test_utils::{CertificateChainBuilder, SignerFixture, TempDir},
6};
7
8use rand_chacha::ChaCha20Rng;
9use rand_core::SeedableRng;
10use std::{fs, path::PathBuf};
11
12/// Create or retrieve a temporary directory for storing cryptographic material for a signer, use this for tests only.
13pub fn setup_temp_directory_for_signer(
14    party_id: &ProtocolPartyId,
15    auto_create: bool,
16) -> Option<PathBuf> {
17    let temp_dir = TempDir::new("tests_setup", "mithril_crypto_helper_material")
18        .build_path()
19        .join(party_id);
20
21    if auto_create {
22        fs::create_dir_all(&temp_dir).expect("temp dir creation should not fail");
23    }
24    temp_dir.exists().then_some(temp_dir)
25}
26
27/// Instantiate a [ProtocolMessage] using fake data, use this for tests only.
28pub fn setup_message() -> ProtocolMessage {
29    let mut protocol_message = ProtocolMessage::new();
30    protocol_message.set_message_part(
31        ProtocolMessagePartKey::SnapshotDigest,
32        "message_to_sign_123".to_string(),
33    );
34    protocol_message.set_message_part(
35        ProtocolMessagePartKey::NextAggregateVerificationKey,
36        "next-avk-123".to_string(),
37    );
38    protocol_message
39}
40
41/// Instantiate a [ProtocolParameters], use this for tests only.
42pub fn setup_protocol_parameters() -> ProtocolParameters {
43    ProtocolParameters {
44        m: 100,
45        k: 5,
46        phi_f: 0.65,
47    }
48}
49
50fn setup_protocol_initializer(
51    party_id: &str,
52    kes_secret_key_path: Option<PathBuf>,
53    stake: Stake,
54    protocol_parameters: &ProtocolParameters,
55) -> ProtocolInitializer {
56    let protocol_initializer_seed: [u8; 32] = format!("{party_id:<032}").as_bytes()[..32]
57        .try_into()
58        .unwrap();
59    let mut protocol_initializer_rng = ChaCha20Rng::from_seed(protocol_initializer_seed);
60    let kes_period = kes_secret_key_path.as_ref().map(|_| 0);
61    let protocol_initializer: ProtocolInitializer = ProtocolInitializer::setup(
62        *protocol_parameters,
63        kes_secret_key_path,
64        kes_period,
65        stake,
66        &mut protocol_initializer_rng,
67    )
68    .expect("protocol initializer setup should not fail");
69
70    protocol_initializer
71}
72
73fn setup_signer_with_stake(
74    party_id: &str,
75    stake: Stake,
76    protocol_initializer: &ProtocolInitializer,
77    operational_certificate: Option<ProtocolOpCert>,
78    kes_period: u32,
79) -> SignerWithStake {
80    let kes_period = operational_certificate.as_ref().and(Some(kes_period));
81
82    SignerWithStake::new(
83        party_id.to_owned(),
84        protocol_initializer.verification_key().into(),
85        protocol_initializer.verification_key_signature(),
86        operational_certificate,
87        kes_period,
88        stake,
89    )
90}
91
92fn decode_op_cert_in_dir(dir: Option<PathBuf>) -> Option<ProtocolOpCert> {
93    dir.as_ref().map(|dir| {
94        OpCert::from_file(dir.join("opcert.cert"))
95            .expect("operational certificate decoding should not fail")
96            .into()
97    })
98}
99
100/// Instantiate a list of protocol signers based on the given [ProtocolStakeDistribution] and [ProtocolParameters], use this for tests only.
101pub fn setup_signers_from_stake_distribution(
102    stake_distribution: &ProtocolStakeDistribution,
103    protocol_parameters: &ProtocolParameters,
104) -> Vec<SignerFixture> {
105    let mut key_registration = ProtocolKeyRegistration::init(stake_distribution);
106    let mut signers: Vec<(SignerWithStake, ProtocolInitializer, Option<PathBuf>)> = vec![];
107
108    for (party_id, stake) in stake_distribution {
109        let kes_period = 0;
110        let temp_dir = setup_temp_directory_for_signer(party_id, false);
111        let kes_secret_key_path: Option<PathBuf> = temp_dir.as_ref().map(|dir| dir.join("kes.sk"));
112        let protocol_initializer = setup_protocol_initializer(
113            party_id,
114            kes_secret_key_path.clone(),
115            *stake,
116            protocol_parameters,
117        );
118        let operational_certificate = decode_op_cert_in_dir(temp_dir);
119        let signer_with_stake = setup_signer_with_stake(
120            party_id,
121            *stake,
122            &protocol_initializer,
123            operational_certificate.clone(),
124            kes_period,
125        );
126
127        key_registration
128            .register(
129                Some(signer_with_stake.party_id.to_owned()),
130                operational_certificate,
131                protocol_initializer.verification_key_signature(),
132                Some(kes_period),
133                protocol_initializer.verification_key().into(),
134            )
135            .expect("key registration should have succeeded");
136
137        signers.push((signer_with_stake, protocol_initializer, kes_secret_key_path));
138    }
139
140    let closed_key_registration = key_registration.close();
141
142    signers
143        .into_iter()
144        .map(
145            |(signer_with_stake, protocol_initializer, kes_secret_key_path)| {
146                let protocol_closed_key_registration = closed_key_registration.clone();
147                let protocol_signer = protocol_initializer
148                    .clone()
149                    .new_signer(protocol_closed_key_registration.clone())
150                    .expect("creating a new protocol signer should not fail");
151
152                SignerFixture {
153                    signer_with_stake,
154                    protocol_signer,
155                    protocol_initializer,
156                    protocol_closed_key_registration,
157                    kes_secret_key_path,
158                }
159            },
160        )
161        .collect::<_>()
162}
163
164/// Instantiate a certificate chain, use this for tests only.
165pub fn setup_certificate_chain(
166    total_certificates: u64,
167    certificates_per_epoch: u64,
168) -> (Vec<Certificate>, ProtocolGenesisVerifier) {
169    let certificate_chain_builder = CertificateChainBuilder::new()
170        .with_total_certificates(total_certificates)
171        .with_certificates_per_epoch(certificates_per_epoch)
172        .with_protocol_parameters(setup_protocol_parameters());
173
174    certificate_chain_builder.build()
175}
mithril_common/crypto_helper/tests_setup.rs

mithril_common/crypto_helper/
tests_setup.rs
