Skip to content

Commit

Permalink
Merge pull request #2161 from input-output-hk/sfa/2153/accelerate_sig…
Browse files Browse the repository at this point in the history
…ner_fixtures_creation_in_tests

Accelerate signer fixtures creation in tests
  • Loading branch information
sfauvel authored Dec 9, 2024
2 parents f1441bc + 38314fe commit 84530e0
Show file tree
Hide file tree
Showing 9 changed files with 596 additions and 33 deletions.
4 changes: 2 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion mithril-aggregator/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "mithril-aggregator"
version = "0.5.121"
version = "0.5.122"
description = "A Mithril Aggregator server"
authors = { workspace = true }
edition = { workspace = true }
Expand Down
4 changes: 2 additions & 2 deletions mithril-aggregator/src/http_server/routes/status.rs
Original file line number Diff line number Diff line change
Expand Up @@ -237,8 +237,8 @@ mod tests {

#[tokio::test]
async fn retrieves_correct_total_signers_from_epoch_service() {
let total_signers = 12;
let total_next_signers = 345;
let total_signers = 5;
let total_next_signers = 4;
let epoch_service = FakeEpochServiceBuilder {
current_signers_with_stake: fake_data::signers_with_stakes(total_signers),
next_signers_with_stake: fake_data::signers_with_stakes(total_next_signers),
Expand Down
4 changes: 2 additions & 2 deletions mithril-aggregator/src/services/epoch_service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1057,7 +1057,7 @@ mod tests {
allowed_discriminants: SignedEntityConfig::dummy().allowed_discriminants,
cardano_era: "CardanoEra".to_string(),
mithril_era: SupportedEra::eras()[1],
total_spo: 40,
total_spo: 10,
total_stake: 20_000_000,
..EpochServiceBuilder::new(epoch, current_epoch_fixture.clone())
};
Expand Down Expand Up @@ -1095,7 +1095,7 @@ mod tests {
current_signers: current_epoch_fixture.signers().into_iter().collect(),
next_signers: next_epoch_fixture.signers().into_iter().collect(),
signed_entity_config: SignedEntityConfig::dummy(),
total_spo: 40,
total_spo: 10,
total_stake: 20_000_000,
}
);
Expand Down
2 changes: 1 addition & 1 deletion mithril-common/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "mithril-common"
version = "0.4.94"
version = "0.4.95"
description = "Common types, interfaces, and utilities for Mithril nodes."
authors = { workspace = true }
edition = { workspace = true }
Expand Down
10 changes: 5 additions & 5 deletions mithril-common/src/entities/mithril_stake_distribution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,11 +70,11 @@ mod tests {

use super::*;

const EXPECTED_HASH: &str = "47675a6fad57040b194655b103bc0439ff9d6920a7ba86bd53756db7ce2952f5";
const EXPECTED_HASH: &str = "c5c1ff02e37c751329e3db7625c77fa2a24e86b2a75422c54f1b9f9232374d6f";

#[test]
fn test_compute_hash() {
let fixtures = MithrilFixtureBuilder::default().with_signers(100).build();
let fixtures = MithrilFixtureBuilder::default().with_signers(10).build();
let stake_distribution = MithrilStakeDistribution::new(
Epoch(1),
fixtures.signers_with_stake(),
Expand All @@ -86,7 +86,7 @@ mod tests {

#[test]
fn test_hash_fail_for_different_stake() {
let fixtures = MithrilFixtureBuilder::default().with_signers(100).build();
let fixtures = MithrilFixtureBuilder::default().with_signers(10).build();
let mut signers = fixtures.signers_with_stake();
signers[0].stake += 1;
let stake_distribution =
Expand All @@ -97,7 +97,7 @@ mod tests {

#[test]
fn test_hash_fail_for_different_epoch() {
let fixtures = MithrilFixtureBuilder::default().with_signers(100).build();
let fixtures = MithrilFixtureBuilder::default().with_signers(10).build();
let stake_distribution = MithrilStakeDistribution::new(
Epoch(2),
fixtures.signers_with_stake(),
Expand All @@ -110,7 +110,7 @@ mod tests {
#[test]
fn test_independence_protocol_parameters() {
let signers = MithrilFixtureBuilder::default()
.with_signers(100)
.with_signers(10)
.build()
.signers_with_stake();
let protocol_parameters = ProtocolParameters {
Expand Down
133 changes: 114 additions & 19 deletions mithril-common/src/test_utils/fixture_builder.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use kes_summed_ed25519::{kes::Sum6Kes, traits::KesSk};
use kes_summed_ed25519::{kes::Sum6Kes, traits::KesSk, PublicKey as KesPublicKey};
use rand_chacha::ChaCha20Rng;
use rand_core::{RngCore, SeedableRng};

Expand All @@ -11,6 +11,8 @@ use crate::{
test_utils::{fake_data, mithril_fixture::MithrilFixture},
};

use super::precomputed_kes_key;

/// A builder of mithril types.
pub struct MithrilFixtureBuilder {
protocol_parameters: ProtocolParameters,
Expand Down Expand Up @@ -130,13 +132,9 @@ impl MithrilFixtureBuilder {
match self.stake_distribution_generation_method {
StakeDistributionGenerationMethod::Custom(_) => vec![],
_ => {
let mut kes_keys_seed = [0u8; 32];
let signers_party_ids = (0..self.number_of_signers).map(|party_index| {
if self.enable_signers_certification {
self.build_party_with_operational_certificate(
party_index,
&mut kes_keys_seed,
)
self.build_party_with_operational_certificate(party_index)
} else {
party_index.to_string()
}
Expand All @@ -146,25 +144,61 @@ impl MithrilFixtureBuilder {
}
}

fn build_party_with_operational_certificate(
&self,
party_index: usize,
kes_key_seed: &mut [u8],
) -> PartyId {
let mut cold_key_seed: Vec<u8> = (party_index)
fn provide_kes_key(kes_key_seed: &mut [u8]) -> (Sum6KesBytes, KesPublicKey) {
if let Some((kes_bytes, kes_verification_key)) =
MithrilFixtureBuilder::cached_kes_key(kes_key_seed)
{
(kes_bytes, kes_verification_key)
} else {
println!(
"KES key not found in test cache, generating a new one for the seed {:?}.",
kes_key_seed
);
MithrilFixtureBuilder::generate_kes_key(kes_key_seed)
}
}

fn cached_kes_key(kes_key_seed: &[u8]) -> Option<(Sum6KesBytes, KesPublicKey)> {
precomputed_kes_key::cached_kes_key(kes_key_seed).map(
|(kes_bytes, kes_verification_key)| {
let kes_verification_key = KesPublicKey::from_bytes(&kes_verification_key).unwrap();
let kes_bytes = Sum6KesBytes(kes_bytes);

(kes_bytes, kes_verification_key)
},
)
}

fn generate_kes_key(kes_key_seed: &mut [u8]) -> (Sum6KesBytes, KesPublicKey) {
let mut key_buffer = [0u8; Sum6Kes::SIZE + 4];

let (kes_secret_key, kes_verification_key) = Sum6Kes::keygen(&mut key_buffer, kes_key_seed);
let mut kes_bytes = Sum6KesBytes([0u8; Sum6Kes::SIZE + 4]);
kes_bytes.0.copy_from_slice(&kes_secret_key.clone_sk());

(kes_bytes, kes_verification_key)
}

fn generate_cold_key_seed(&self, party_index: usize) -> Vec<u8> {
let mut cold_key_seed: Vec<_> = (party_index)
.to_le_bytes()
.iter()
.zip(self.party_id_seed)
.map(|(v1, v2)| v1 + v2)
.collect();
cold_key_seed.resize(32, 0);

cold_key_seed
}

fn build_party_with_operational_certificate(&self, party_index: usize) -> PartyId {
let cold_key_seed = self.generate_cold_key_seed(party_index).to_vec();
let mut kes_key_seed = cold_key_seed.clone();

let keypair =
ColdKeyGenerator::create_deterministic_keypair(cold_key_seed.try_into().unwrap());
let mut dummy_buffer = [0u8; Sum6Kes::SIZE + 4];
let (kes_secret_key, kes_verification_key) =
Sum6Kes::keygen(&mut dummy_buffer, kes_key_seed);
let mut kes_bytes = Sum6KesBytes([0u8; Sum6Kes::SIZE + 4]);
kes_bytes.0.copy_from_slice(&kes_secret_key.clone_sk());
let (kes_bytes, kes_verification_key) =
MithrilFixtureBuilder::provide_kes_key(&mut kes_key_seed);
let operational_certificate = OpCert::new(kes_verification_key, 0, 0, keypair);
let party_id = operational_certificate
.compute_protocol_party_id()
Expand Down Expand Up @@ -265,11 +299,11 @@ mod tests {
#[test]
fn changing_party_id_seed_change_all_builded_party_ids() {
let first_signers = MithrilFixtureBuilder::default()
.with_signers(20)
.with_signers(10)
.build()
.signers_with_stake();
let different_party_id_seed_signers = MithrilFixtureBuilder::default()
.with_signers(20)
.with_signers(10)
.with_party_id_seed([1u8; 32])
.build()
.signers_with_stake();
Expand All @@ -279,4 +313,65 @@ mod tests {
assert!(!first_party_ids.contains(&party_id));
}
}

/// Verify that there is a cached kes key for a number of party id.
/// If the cache is not up to date, the test will generate the code that can be copied/pasted into the [precomputed_kes_key] module.
/// The number of party id that should be in cache is defined with `precomputed_number`
#[test]
fn verify_kes_key_cache_content() {
// Generate code that should be in the `cached_kes_key` function of the `precomputed_kes_key.rs` file.
// It can be copied and pasted to update the cache.
fn generate_code(party_ids: &Vec<(&[u8], [u8; 612], KesPublicKey)>) -> String {
party_ids
.iter()
.map(|(key, i, p)| format!("{:?} => ({:?}, {:?}),", key, i, p.as_bytes()))
.collect::<Vec<_>>()
.join("\n")
}

let precomputed_number = 10;

let fixture = MithrilFixtureBuilder::default();
let cold_keys: Vec<_> = (0..precomputed_number)
.map(|party_index| fixture.generate_cold_key_seed(party_index))
.collect();

let computed_keys_key: Vec<_> = cold_keys
.iter()
.map(|cold_key| {
let mut kes_key_seed: Vec<u8> = cold_key.clone();
let (kes_bytes, kes_verification_key) =
MithrilFixtureBuilder::generate_kes_key(&mut kes_key_seed);

(cold_key.as_slice(), kes_bytes.0, kes_verification_key)
})
.collect();

let cached_kes_key: Vec<_> = cold_keys
.iter()
.filter_map(|cold_key| {
MithrilFixtureBuilder::cached_kes_key(cold_key).map(
|(kes_bytes, kes_verification_key)| {
(cold_key.as_slice(), kes_bytes.0, kes_verification_key)
},
)
})
.collect();

let expected_code = generate_code(&computed_keys_key);
let actual_code = generate_code(&cached_kes_key);

assert_eq!(
computed_keys_key, cached_kes_key,
"Precomputed KES keys should be:\n{}\nbut seems to be:\n{}",
expected_code, actual_code
);

let kes_key_seed = fixture.generate_cold_key_seed(precomputed_number);
assert!(
MithrilFixtureBuilder::cached_kes_key(kes_key_seed.as_slice()).is_none(),
"We checked precomputed KES keys up to {} but it seems to be more.",
precomputed_number
);
}
}
2 changes: 1 addition & 1 deletion mithril-common/src/test_utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ mod cardano_transactions_builder;
mod certificate_chain_builder;
mod fixture_builder;
mod mithril_fixture;

mod precomputed_kes_key;
mod temp_dir;

#[cfg(feature = "test_http_server")]
Expand Down
Loading

0 comments on commit 84530e0

Please sign in to comment.