Merge pull request #2161 from input-output-hk/sfa/2153/accelerate_signer_fixtures_creation_in_tests

Accelerate signer fixtures creation in tests

Author: sfauvel

Cargo.lock

@@ -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 }

mithril-aggregator/Cargo.toml

@@ -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),

mithril-aggregator/src/http_server/routes/status.rs

@@ -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())
         };
@@ -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,
             }
         );

mithril-aggregator/src/services/epoch_service.rs

@@ -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 }

mithril-common/Cargo.toml

@@ -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(),
@@ -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 =
@@ -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(),
@@ -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 {

mithril-common/src/entities/mithril_stake_distribution.rs

@@ -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};
 
@@ -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,
@@ -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()
                     }
@@ -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()
@@ -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();
@@ -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
+        );
+    }
 }

mithril-common/src/test_utils/fixture_builder.rs

@@ -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")]