refactor: make signable seed builder compute protocol message part values

jpraynaud · jpraynaud · commit b322df8ba41c · 2024-09-24T16:04:48.000+02:00
And let the responsibility of using it in the signable builder service.
diff --git a/mithril-aggregator/src/services/signable_seed_builder.rs b/mithril-aggregator/src/services/signable_seed_builder.rs
@@ -9,9 +9,7 @@ use std::sync::Arc;
 use tokio::sync::RwLock;
 
 use mithril_common::{
-    entities::{ProtocolMessage, ProtocolMessagePartKey},
-    signable_builder::SignableSeedBuilder,
-    StdResult,
+    entities::ProtocolMessagePartValue, signable_builder::SignableSeedBuilder, StdResult,
 };
 
 use crate::services::EpochService;
@@ -30,22 +28,17 @@ impl SignableSeedBuilderService {
 
 #[async_trait]
 impl SignableSeedBuilder for SignableSeedBuilderService {
-    async fn compute_seeded_protocol_message(
+    async fn compute_next_aggregate_verification_key_protocol_message_part_value(
         &self,
-        protocol_message: ProtocolMessage,
-    ) -> StdResult<ProtocolMessage> {
+    ) -> StdResult<ProtocolMessagePartValue> {
         let epoch_service = self.epoch_service.read().await;
-        let mut protocol_message = protocol_message;
-        protocol_message.set_message_part(
-            ProtocolMessagePartKey::NextAggregateVerificationKey,
-            (*epoch_service)
-                .next_aggregate_verification_key()?
-                .to_json_hex()
-                .with_context(|| "convert next avk to json hex failure")?
-                .to_string(),
-        );
+        let next_aggregate_verification_key = (*epoch_service)
+            .next_aggregate_verification_key()?
+            .to_json_hex()
+            .with_context(|| "convert next avk to json hex failure")?
+            .to_string();
 
-        Ok(protocol_message)
+        Ok(next_aggregate_verification_key)
     }
 }
 
@@ -58,7 +51,7 @@ mod tests {
     use super::*;
 
     #[tokio::test]
-    async fn test_compute_seeded_protocol_message() {
+    async fn test_compute_next_aggregate_verification_key_protocol_message_value() {
         let epoch = Epoch(5);
         let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
         let next_fixture = MithrilFixtureBuilder::default().with_signers(4).build();
@@ -73,15 +66,14 @@ mod tests {
         )));
         let signable_seed_builder_service = SignableSeedBuilderService::new(epoch_service);
 
-        let protocol_message = signable_seed_builder_service
-            .compute_seeded_protocol_message(ProtocolMessage::new())
+        let next_aggregate_verification_key = signable_seed_builder_service
+            .compute_next_aggregate_verification_key_protocol_message_part_value()
             .await
             .unwrap();
 
         assert_eq!(
-            protocol_message
-                .get_message_part(&ProtocolMessagePartKey::NextAggregateVerificationKey),
-            Some(&expected_next_aggregate_verification_key)
+            next_aggregate_verification_key,
+            expected_next_aggregate_verification_key
         );
     }
 }
diff --git a/mithril-common/src/signable_builder/interface.rs b/mithril-common/src/signable_builder/interface.rs
@@ -1,7 +1,10 @@
 use async_trait::async_trait;
 use std::fmt::Debug;
 
-use crate::{entities::ProtocolMessage, StdResult};
+use crate::{
+    entities::{ProtocolMessage, ProtocolMessagePartValue},
+    StdResult,
+};
 
 #[cfg(test)]
 use mockall::automock;
@@ -27,13 +30,12 @@ where
     async fn compute_protocol_message(&self, beacon: U) -> StdResult<ProtocolMessage>;
 }
 
-/// SignableSeedBuilder is a trait for building a seed protocol message
+/// SignableSeedBuilder is a trait for building seed protocol message part values
 #[cfg_attr(test, automock)]
 #[async_trait]
 pub trait SignableSeedBuilder: Send + Sync {
-    /// Compute seeded protocol message
-    async fn compute_seeded_protocol_message(
+    /// Compute next aggregate verification key protocol message part value
+    async fn compute_next_aggregate_verification_key_protocol_message_part_value(
         &self,
-        protocol_message: ProtocolMessage,
-    ) -> StdResult<ProtocolMessage>;
+    ) -> StdResult<ProtocolMessagePartValue>;
 }
diff --git a/mithril-common/src/signable_builder/signable_builder_service.rs b/mithril-common/src/signable_builder/signable_builder_service.rs
@@ -3,16 +3,16 @@ use async_trait::async_trait;
 use std::sync::Arc;
 
 use crate::{
-    entities::{BlockNumber, CardanoDbBeacon, Epoch, ProtocolMessage, SignedEntityType},
+    entities::{
+        BlockNumber, CardanoDbBeacon, Epoch, ProtocolMessage, ProtocolMessagePartKey,
+        SignedEntityType,
+    },
     signable_builder::{SignableBuilder, SignableSeedBuilder},
     StdResult,
 };
 
-#[cfg(test)]
-use mockall::automock;
-
 /// ArtifactBuilder Service trait
-#[cfg_attr(test, automock)]
+#[cfg_attr(test, mockall::automock)]
 #[async_trait]
 pub trait SignableBuilderService: Send + Sync {
     /// Compute signable from signed entity type
@@ -48,11 +48,8 @@ impl MithrilSignableBuilderService {
             cardano_stake_distribution_builder,
         }
     }
-}
 
-#[async_trait]
-impl SignableBuilderService for MithrilSignableBuilderService {
-    async fn compute_protocol_message(
+    async fn compute_signed_entity_protocol_message(
         &self,
         signed_entity_type: SignedEntityType,
     ) -> StdResult<ProtocolMessage> {
@@ -84,8 +81,38 @@ impl SignableBuilderService for MithrilSignableBuilderService {
                     "Signable builder service can not compute protocol message with block_number: '{block_number}'"
                 ))?,
         };
-        let protocol_message = self
+
+        Ok(protocol_message)
+    }
+
+    async fn compute_seeded_protocol_message(
+        &self,
+        protocol_message: ProtocolMessage,
+    ) -> StdResult<ProtocolMessage> {
+        let next_aggregate_verification_key = self
             .seed_signable_builder
+            .compute_next_aggregate_verification_key_protocol_message_part_value()
+            .await?;
+        let mut protocol_message = protocol_message;
+        protocol_message.set_message_part(
+            ProtocolMessagePartKey::NextAggregateVerificationKey,
+            next_aggregate_verification_key,
+        );
+
+        Ok(protocol_message)
+    }
+}
+
+#[async_trait]
+impl SignableBuilderService for MithrilSignableBuilderService {
+    async fn compute_protocol_message(
+        &self,
+        signed_entity_type: SignedEntityType,
+    ) -> StdResult<ProtocolMessage> {
+        let protocol_message = self
+            .compute_signed_entity_protocol_message(signed_entity_type)
+            .await?;
+        let protocol_message = self
             .compute_seeded_protocol_message(protocol_message)
             .await?;
 
@@ -98,7 +125,7 @@ mod tests {
     use super::*;
 
     use crate::{
-        entities::{BlockNumber, Epoch, ProtocolMessage},
+        entities::{BlockNumber, Epoch, ProtocolMessage, ProtocolMessagePartValue},
         signable_builder::{Beacon as Beaconnable, SignableBuilder},
         StdResult,
     };
@@ -121,7 +148,7 @@ mod tests {
 
         #[async_trait]
         impl SignableSeedBuilder for SignableSeedBuilderImpl {
-            async fn compute_seeded_protocol_message(&self, protocol_message: ProtocolMessage) -> StdResult<ProtocolMessage>;
+            async fn compute_next_aggregate_verification_key_protocol_message_part_value(&self) -> StdResult<ProtocolMessagePartValue>;
         }
     }
 
@@ -164,9 +191,9 @@ mod tests {
         let mut mock_container = MockDependencyInjector::new();
         mock_container
             .mock_signable_seed_builder
-            .expect_compute_seeded_protocol_message()
+            .expect_compute_next_aggregate_verification_key_protocol_message_value()
             .once()
-            .return_once(move |_| Ok(ProtocolMessage::new()));
+            .return_once(move || Ok("next-avk-123".to_string()));
         mock_container
             .mock_mithril_stake_distribution_signable_builder
             .expect_compute_protocol_message()
@@ -188,9 +215,9 @@ mod tests {
         let mut mock_container = MockDependencyInjector::new();
         mock_container
             .mock_signable_seed_builder
-            .expect_compute_seeded_protocol_message()
+            .expect_compute_next_aggregate_verification_key_protocol_message_value()
             .once()
-            .return_once(move |_| Ok(ProtocolMessage::new()));
+            .return_once(move || Ok("next-avk-123".to_string()));
         mock_container
             .mock_cardano_immutable_files_full_signable_builder
             .expect_compute_protocol_message()
@@ -213,9 +240,9 @@ mod tests {
         let mut mock_container = MockDependencyInjector::new();
         mock_container
             .mock_signable_seed_builder
-            .expect_compute_seeded_protocol_message()
+            .expect_compute_next_aggregate_verification_key_protocol_message_value()
             .once()
-            .return_once(move |_| Ok(ProtocolMessage::new()));
+            .return_once(move || Ok("next-avk-123".to_string()));
         mock_container
             .mock_cardano_transactions_signable_builder
             .expect_compute_protocol_message()
@@ -238,9 +265,9 @@ mod tests {
         let mut mock_container = MockDependencyInjector::new();
         mock_container
             .mock_signable_seed_builder
-            .expect_compute_seeded_protocol_message()
+            .expect_compute_next_aggregate_verification_key_protocol_message_value()
             .once()
-            .return_once(move |_| Ok(ProtocolMessage::new()));
+            .return_once(move || Ok("next-avk-123".to_string()));
         mock_container
             .mock_cardano_stake_distribution_signable_builder
             .expect_compute_protocol_message()
diff --git a/mithril-signer/src/services/signable_seed_builder.rs b/mithril-signer/src/services/signable_seed_builder.rs
@@ -9,9 +9,7 @@ use std::sync::Arc;
 use tokio::sync::RwLock;
 
 use mithril_common::{
-    entities::{ProtocolMessage, ProtocolMessagePartKey},
-    signable_builder::SignableSeedBuilder,
-    StdResult,
+    entities::ProtocolMessagePartValue, signable_builder::SignableSeedBuilder, StdResult,
 };
 
 use crate::{
@@ -43,11 +41,9 @@ impl SignableSeedBuilderService {
 
 #[async_trait]
 impl SignableSeedBuilder for SignableSeedBuilderService {
-    async fn compute_seeded_protocol_message(
+    async fn compute_next_aggregate_verification_key_protocol_message_part_value(
         &self,
-        protocol_message: ProtocolMessage,
-    ) -> StdResult<ProtocolMessage> {
-        let mut protocol_message = protocol_message;
+    ) -> StdResult<ProtocolMessagePartValue> {
         let epoch_service = self.epoch_service.read().await;
         let epoch = (*epoch_service).epoch_of_current_data()?;
         let next_signer_retrieval_epoch = epoch.offset_to_next_signer_retrieval_epoch();
@@ -60,17 +56,15 @@ impl SignableSeedBuilder for SignableSeedBuilderService {
             })?;
 
         let next_signers_with_stake = epoch_service.next_signers_with_stake().await?;
-        let avk = self
+        let next_aggregate_verification_key = self
             .single_signer
             .compute_aggregate_verification_key(
                 &next_signers_with_stake,
                 &next_protocol_initializer,
             )?
             .ok_or_else(|| anyhow!("next_signers avk".to_string()))?;
-        protocol_message
-            .set_message_part(ProtocolMessagePartKey::NextAggregateVerificationKey, avk);
 
-        Ok(protocol_message)
+        Ok(next_aggregate_verification_key)
     }
 }
 
@@ -87,7 +81,7 @@ mod tests {
     use super::*;
 
     #[tokio::test]
-    async fn test_compute_seeded_protocol_message() {
+    async fn test_compute_next_aggregate_verification_key_protocol_message_value() {
         let epoch = Epoch(5);
         let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
         let next_fixture = MithrilFixtureBuilder::default().with_signers(4).build();
@@ -124,15 +118,14 @@ mod tests {
             protocol_initializer_store,
         );
 
-        let protocol_message = signable_seed_builder_service
-            .compute_seeded_protocol_message(ProtocolMessage::new())
+        let next_aggregate_verification_key = signable_seed_builder_service
+            .compute_next_aggregate_verification_key_protocol_message_part_value()
             .await
             .unwrap();
 
         assert_eq!(
-            protocol_message
-                .get_message_part(&ProtocolMessagePartKey::NextAggregateVerificationKey),
-            Some(&expected_next_aggregate_verification_key)
+            next_aggregate_verification_key,
+            expected_next_aggregate_verification_key
         );
     }
 }
