Update VerificationKeyStorer 'save_verification_key' signature

To prepare migration to 'signer_registration' table that also records stake.

Author: jpraynaud

mithril-aggregator/src/dependency.rs

@@ -4,9 +4,7 @@ use mithril_common::{
     chain_observer::ChainObserver,
     crypto_helper::ProtocolGenesisVerifier,
     digesters::{ImmutableDigester, ImmutableFileObserver},
-    entities::{
-        Certificate, Epoch, ProtocolParameters, Signer, SignerWithStake, StakeDistribution,
-    },
+    entities::{Certificate, Epoch, ProtocolParameters, SignerWithStake, StakeDistribution},
     era::{EraChecker, EraReader},
     store::StakeStorer,
     BeaconProvider,
@@ -253,11 +251,7 @@ impl DependencyManager {
     }
 
     async fn fill_verification_key_store(&self, target_epoch: Epoch, signers: &[SignerWithStake]) {
-        for signer in signers
-            .iter()
-            .map(|s| s.to_owned().into())
-            .collect::<Vec<Signer>>()
-        {
+        for signer in signers {
             self.verification_key_store
                 .save_verification_key(target_epoch, signer.clone())
                 .await

mithril-aggregator/src/multi_signer.rs

@@ -839,7 +839,7 @@ mod tests {
             verification_key_store
                 .save_verification_key(
                     start_epoch.offset_to_recording_epoch(),
-                    signer_with_stake.to_owned().into(),
+                    signer_with_stake.to_owned(),
                 )
                 .await
                 .expect("register should have succeeded");

mithril-aggregator/src/signer_registerer.rs

@@ -208,7 +208,7 @@ impl SignerRegisterer for MithrilSignerRegisterer {
 
         match self
             .verification_key_store
-            .save_verification_key(registration_round.epoch, signer_save.clone().into())
+            .save_verification_key(registration_round.epoch, signer_save.clone())
             .await?
         {
             Some(_) => Err(SignerRegistrationError::ExistingSigner(signer_save)),

mithril-aggregator/src/store/verification_key_store.rs

@@ -3,7 +3,7 @@ use mithril_common::store::StorePruner;
 use std::collections::HashMap;
 use tokio::sync::RwLock;
 
-use mithril_common::entities::{Epoch, PartyId, Signer};
+use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
 use mithril_common::store::{adapter::StoreAdapter, StoreError};
 
 type Adapter = Box<dyn StoreAdapter<Key = Epoch, Record = HashMap<PartyId, Signer>>>;
@@ -15,8 +15,8 @@ pub trait VerificationKeyStorer {
     async fn save_verification_key(
         &self,
         epoch: Epoch,
-        signer: Signer,
-    ) -> Result<Option<Signer>, StoreError>;
+        signer: SignerWithStake,
+    ) -> Result<Option<SignerWithStake>, StoreError>;
 
     /// Returns a HashMap of [Signer] indexed by [PartyId] for the given `Beacon`.
     async fn get_verification_keys(
@@ -61,21 +61,21 @@ impl VerificationKeyStorer for VerificationKeyStore {
     async fn save_verification_key(
         &self,
         epoch: Epoch,
-        signer: Signer,
-    ) -> Result<Option<Signer>, StoreError> {
+        signer: SignerWithStake,
+    ) -> Result<Option<SignerWithStake>, StoreError> {
         let mut signers = match self.adapter.read().await.get_record(&epoch).await? {
             Some(s) => s,
             None => HashMap::new(),
         };
-        let prev_signer = signers.insert(signer.party_id.to_owned(), signer);
+        let prev_signer = signers.insert(signer.party_id.to_owned(), signer.clone().into());
         self.adapter
             .write()
             .await
             .store_record(&epoch, &signers)
             .await?;
         self.prune().await?;
 
-        Ok(prev_signer)
+        Ok(prev_signer.map(|prev_signer| SignerWithStake::from_signer(prev_signer, signer.stake)))
     }
 
     async fn get_verification_keys(
@@ -135,12 +135,13 @@ mod tests {
         let res = store
             .save_verification_key(
                 Epoch(0),
-                Signer {
+                SignerWithStake {
                     party_id: "0".to_string(),
                     verification_key: "OK".to_string(),
                     verification_key_signature: None,
                     operational_certificate: None,
                     kes_period: None,
+                    stake: 10,
                 },
             )
             .await
@@ -155,25 +156,27 @@ mod tests {
         let res = store
             .save_verification_key(
                 Epoch(1),
-                Signer {
+                SignerWithStake {
                     party_id: "1".to_string(),
                     verification_key: "test".to_string(),
                     verification_key_signature: None,
                     operational_certificate: None,
                     kes_period: None,
+                    stake: 10,
                 },
             )
             .await
             .unwrap();
 
         assert!(res.is_some());
         assert_eq!(
-            Signer {
+            SignerWithStake {
                 party_id: "1".to_string(),
                 verification_key: "vkey 1".to_string(),
                 verification_key_signature: None,
                 operational_certificate: None,
                 kes_period: None,
+                stake: 10,
             },
             res.unwrap(),
         );
@@ -207,12 +210,13 @@ mod tests {
         let _ = store
             .save_verification_key(
                 Epoch(3),
-                Signer {
+                SignerWithStake {
                     party_id: "party_id".to_string(),
                     verification_key: "whatever".to_string(),
                     verification_key_signature: None,
                     operational_certificate: None,
                     kes_period: None,
+                    stake: 10,
                 },
             )
             .await