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