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        KesSigner, KesSignerStandard, OpCert, ProtocolInitializer, ProtocolKeyRegistration,
10        ProtocolOpCert, ProtocolParameters, ProtocolPartyId, ProtocolStakeDistribution,
11        SerDeShelleyFileFormat,
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(|_| 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_period: u32,
95) -> SignerWithStake {
96    let kes_period = operational_certificate.as_ref().and(Some(kes_period));
97
98    SignerWithStake::new(
99        party_id.to_owned(),
100        protocol_initializer.verification_key().into(),
101        protocol_initializer.verification_key_signature(),
102        operational_certificate,
103        kes_period,
104        stake,
105    )
106}
107
108fn decode_op_cert_in_dir(dir: Option<PathBuf>) -> Option<ProtocolOpCert> {
109    dir.as_ref().map(|dir| {
110        OpCert::from_file(dir.join("opcert.cert"))
111            .expect("operational certificate decoding should not fail")
112            .into()
113    })
114}
115
116/// Instantiate a list of protocol signers based on the given [ProtocolStakeDistribution] and [ProtocolParameters], use this for tests only.
117pub fn setup_signers_from_stake_distribution(
118    stake_distribution: &ProtocolStakeDistribution,
119    protocol_parameters: &ProtocolParameters,
120) -> Vec<SignerFixture> {
121    let mut key_registration = ProtocolKeyRegistration::init(stake_distribution);
122    let mut signers: Vec<(
123        SignerWithStake,
124        ProtocolInitializer,
125        Option<PathBuf>,
126        Option<PathBuf>,
127    )> = vec![];
128
129    for (party_id, stake) in stake_distribution {
130        let kes_period = 0;
131        let temp_dir = setup_temp_directory_for_signer(party_id, false);
132        let kes_secret_key_path: Option<PathBuf> = temp_dir.as_ref().map(|dir| dir.join("kes.sk"));
133        let operational_certificate_path = temp_dir.as_ref().map(|dir| dir.join("opcert.cert"));
134        let protocol_initializer = setup_protocol_initializer(
135            party_id,
136            kes_secret_key_path.clone(),
137            operational_certificate_path.clone(),
138            *stake,
139            protocol_parameters,
140        );
141        let operational_certificate = decode_op_cert_in_dir(temp_dir);
142        let signer_with_stake = setup_signer_with_stake(
143            party_id,
144            *stake,
145            &protocol_initializer,
146            operational_certificate.clone(),
147            kes_period,
148        );
149
150        key_registration
151            .register(
152                Some(signer_with_stake.party_id.to_owned()),
153                operational_certificate,
154                protocol_initializer.verification_key_signature(),
155                Some(kes_period),
156                protocol_initializer.verification_key().into(),
157            )
158            .expect("key registration should have succeeded");
159
160        signers.push((
161            signer_with_stake,
162            protocol_initializer,
163            kes_secret_key_path,
164            operational_certificate_path,
165        ));
166    }
167
168    let closed_key_registration = key_registration.close();
169
170    signers
171        .into_iter()
172        .map(
173            |(
174                signer_with_stake,
175                protocol_initializer,
176                kes_secret_key_path,
177                operational_certificate_path,
178            )| {
179                let protocol_closed_key_registration = closed_key_registration.clone();
180                let protocol_signer = protocol_initializer
181                    .clone()
182                    .new_signer(protocol_closed_key_registration.clone())
183                    .expect("creating a new protocol signer should not fail");
184
185                SignerFixture {
186                    signer_with_stake,
187                    protocol_signer,
188                    protocol_initializer,
189                    protocol_closed_key_registration,
190                    kes_secret_key_path,
191                    operational_certificate_path,
192                }
193            },
194        )
195        .collect::<_>()
196}
197
198/// Instantiate a certificate chain, use this for tests only.
199pub fn setup_certificate_chain(
200    total_certificates: u64,
201    certificates_per_epoch: u64,
202) -> CertificateChainFixture {
203    CertificateChainBuilder::new()
204        .with_total_certificates(total_certificates)
205        .with_certificates_per_epoch(certificates_per_epoch)
206        .with_protocol_parameters(setup_protocol_parameters())
207        .build()
208}
mithril_common/test/crypto_helper/setup.rs

mithril_common/test/crypto_helper/
setup.rs
