mithril_common/test/crypto_helper/
setup.rs

1//! Test data builders for Mithril STM types, for testing purpose.
2use std::{fs, path::PathBuf, sync::Arc};
3
4use rand_chacha::ChaCha20Rng;
5use rand_core::SeedableRng;
6
7use crate::{
8    crypto_helper::{
9        KesEvolutions, KesPeriod, KesSigner, KesSignerStandard, OpCert, ProtocolInitializer,
10        ProtocolKeyRegistration, ProtocolOpCert, ProtocolParameters, ProtocolPartyId,
11        ProtocolStakeDistribution, SerDeShelleyFileFormat, SignerRegistrationParameters,
12    },
13    entities::{ProtocolMessage, ProtocolMessagePartKey, SignerWithStake, Stake},
14    test::{
15        TempDir,
16        builder::{CertificateChainBuilder, CertificateChainFixture, SignerFixture},
17    },
18};
19
20/// Create or retrieve a temporary directory for storing cryptographic material for a signer, use this for tests only.
21pub fn setup_temp_directory_for_signer(
22    party_id: &ProtocolPartyId,
23    auto_create: bool,
24) -> Option<PathBuf> {
25    let temp_dir = TempDir::new("tests_setup", "mithril_crypto_helper_material")
26        .build_path()
27        .join(party_id);
28
29    if auto_create {
30        fs::create_dir_all(&temp_dir).expect("temp dir creation should not fail");
31    }
32    temp_dir.exists().then_some(temp_dir)
33}
34
35/// Instantiate a [ProtocolMessage] using fake data, use this for tests only.
36pub fn setup_message() -> ProtocolMessage {
37    let mut protocol_message = ProtocolMessage::new();
38    protocol_message.set_message_part(
39        ProtocolMessagePartKey::SnapshotDigest,
40        "message_to_sign_123".to_string(),
41    );
42    protocol_message.set_message_part(
43        ProtocolMessagePartKey::NextAggregateVerificationKey,
44        "next-avk-123".to_string(),
45    );
46    protocol_message
47}
48
49/// Instantiate a [ProtocolParameters], use this for tests only.
50pub fn setup_protocol_parameters() -> ProtocolParameters {
51    ProtocolParameters {
52        m: 100,
53        k: 5,
54        phi_f: 0.65,
55    }
56}
57
58fn setup_protocol_initializer(
59    party_id: &str,
60    kes_secret_key_path: Option<PathBuf>,
61    operational_certificate_path: Option<PathBuf>,
62    stake: Stake,
63    protocol_parameters: &ProtocolParameters,
64) -> ProtocolInitializer {
65    let protocol_initializer_seed: [u8; 32] =
66        format!("{party_id:<032}").as_bytes()[..32].try_into().unwrap();
67    let mut protocol_initializer_rng = ChaCha20Rng::from_seed(protocol_initializer_seed);
68    let kes_period = kes_secret_key_path.as_ref().map(|_| KesPeriod(0));
69    let kes_signer = kes_secret_key_path.map(|kes_secret_key_path| {
70        Arc::new(KesSignerStandard::new(
71            kes_secret_key_path,
72            operational_certificate_path.expect(
73                "Operational certificate path must be provided when a KES secret key exists",
74            ),
75        )) as Arc<dyn KesSigner>
76    });
77    let protocol_initializer: ProtocolInitializer = ProtocolInitializer::setup(
78        *protocol_parameters,
79        kes_signer,
80        kes_period,
81        stake,
82        &mut protocol_initializer_rng,
83    )
84    .expect("protocol initializer setup should not fail");
85
86    protocol_initializer
87}
88
89fn setup_signer_with_stake(
90    party_id: &str,
91    stake: Stake,
92    protocol_initializer: &ProtocolInitializer,
93    operational_certificate: Option<ProtocolOpCert>,
94    kes_evolutions: KesEvolutions,
95) -> SignerWithStake {
96    let kes_evolutions = operational_certificate.as_ref().and(Some(kes_evolutions));
97
98    SignerWithStake {
99        party_id: party_id.to_owned(),
100        verification_key_for_concatenation: protocol_initializer
101            .verification_key_for_concatenation()
102            .into(),
103        verification_key_signature_for_concatenation: protocol_initializer
104            .verification_key_signature_for_concatenation(),
105        operational_certificate,
106        kes_evolutions,
107        stake,
108        #[cfg(feature = "future_snark")]
109        verification_key_for_snark: protocol_initializer
110            .verification_key_for_snark()
111            .map(|vk| vk.into()),
112        #[cfg(feature = "future_snark")]
113        verification_key_signature_for_snark: protocol_initializer
114            .verification_key_signature_for_snark(),
115    }
116}
117
118fn decode_op_cert_in_dir(dir: Option<PathBuf>) -> Option<ProtocolOpCert> {
119    dir.as_ref().map(|dir| {
120        OpCert::from_file(dir.join("opcert.cert"))
121            .expect("operational certificate decoding should not fail")
122            .into()
123    })
124}
125
126/// Instantiate a list of protocol signers based on the given [ProtocolStakeDistribution] and [ProtocolParameters], use this for tests only.
127pub fn setup_signers_from_stake_distribution(
128    stake_distribution: &ProtocolStakeDistribution,
129    protocol_parameters: &ProtocolParameters,
130) -> Vec<SignerFixture> {
131    let mut key_registration = ProtocolKeyRegistration::init(stake_distribution);
132    let mut signers: Vec<(
133        SignerWithStake,
134        ProtocolInitializer,
135        Option<PathBuf>,
136        Option<PathBuf>,
137    )> = vec![];
138
139    for (party_id, stake) in stake_distribution {
140        let kes_evolutions = KesEvolutions(0);
141        let temp_dir = setup_temp_directory_for_signer(party_id, false);
142        let kes_secret_key_path: Option<PathBuf> = temp_dir.as_ref().map(|dir| dir.join("kes.sk"));
143        let operational_certificate_path = temp_dir.as_ref().map(|dir| dir.join("opcert.cert"));
144        let protocol_initializer = setup_protocol_initializer(
145            party_id,
146            kes_secret_key_path.clone(),
147            operational_certificate_path.clone(),
148            *stake,
149            protocol_parameters,
150        );
151        let operational_certificate = decode_op_cert_in_dir(temp_dir);
152        let signer_with_stake = setup_signer_with_stake(
153            party_id,
154            *stake,
155            &protocol_initializer,
156            operational_certificate.clone(),
157            kes_evolutions,
158        );
159
160        key_registration
161            .register(SignerRegistrationParameters {
162                party_id: Some(signer_with_stake.party_id.to_owned()),
163                operational_certificate,
164                verification_key_signature_for_concatenation: protocol_initializer
165                    .verification_key_signature_for_concatenation(),
166                kes_evolutions: Some(kes_evolutions),
167                verification_key_for_concatenation: protocol_initializer
168                    .verification_key_for_concatenation()
169                    .into(),
170                #[cfg(feature = "future_snark")]
171                verification_key_for_snark: protocol_initializer
172                    .verification_key_for_snark()
173                    .map(Into::into),
174                #[cfg(feature = "future_snark")]
175                verification_key_signature_for_snark: protocol_initializer
176                    .verification_key_signature_for_snark(),
177            })
178            .expect("key registration should have succeeded");
179
180        signers.push((
181            signer_with_stake,
182            protocol_initializer,
183            kes_secret_key_path,
184            operational_certificate_path,
185        ));
186    }
187
188    let closed_key_registration = key_registration.close();
189
190    signers
191        .into_iter()
192        .map(
193            |(
194                signer_with_stake,
195                protocol_initializer,
196                kes_secret_key_path,
197                operational_certificate_path,
198            )| {
199                let protocol_closed_key_registration = closed_key_registration.clone();
200                let protocol_signer = protocol_initializer
201                    .clone()
202                    .new_signer(protocol_closed_key_registration.clone())
203                    .expect("creating a new protocol signer should not fail");
204
205                SignerFixture {
206                    signer_with_stake,
207                    protocol_signer,
208                    protocol_initializer,
209                    protocol_closed_key_registration,
210                    kes_secret_key_path,
211                    operational_certificate_path,
212                }
213            },
214        )
215        .collect::<_>()
216}
217
218/// Instantiate a certificate chain, use this for tests only.
219pub fn setup_certificate_chain(
220    total_certificates: u64,
221    certificates_per_epoch: u64,
222) -> CertificateChainFixture {
223    CertificateChainBuilder::new()
224        .with_total_certificates(total_certificates)
225        .with_certificates_per_epoch(certificates_per_epoch)
226        .with_protocol_parameters(setup_protocol_parameters())
227        .build()
228}
mithril_common/test/crypto_helper/setup.rs

mithril_common/test/crypto_helper/
setup.rs
