Refactor: signable builder services compute full protocol message

CHANGELOG.md

@@ -19,6 +19,8 @@ As a minor extension, we have adopted a slightly different versioning convention
 
 - Support for signing the protocol parameters in the Genesis certificate.
 
+- Refactor the builder of the protocol messages to be signed.
+
 - Crates versions:
 
 | Crate | Version |

mithril-aggregator/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-aggregator"
-version = "0.5.67"
+version = "0.5.68"
 description = "A Mithril Aggregator server"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-aggregator/src/dependency_injection/builder.rs

@@ -40,7 +40,8 @@ use mithril_common::{
     signable_builder::{
         CardanoImmutableFilesFullSignableBuilder, CardanoStakeDistributionSignableBuilder,
         CardanoTransactionsSignableBuilder, MithrilSignableBuilderService,
-        MithrilStakeDistributionSignableBuilder, SignableBuilderService, TransactionsImporter,
+        MithrilStakeDistributionSignableBuilder, SignableBuilderService, SignableSeedBuilder,
+        TransactionsImporter,
     },
     signed_entity_type_lock::SignedEntityTypeLock,
     MithrilTickerService, TickerService,
@@ -66,11 +67,11 @@ use crate::{
     event_store::{EventMessage, EventStore, TransmitterService},
     http_server::routes::router,
     services::{
-        AggregatorUpkeepService, BufferedCertifierService, CardanoTransactionsImporter,
-        CertifierService, MessageService, MithrilCertifierService, MithrilEpochService,
-        MithrilMessageService, MithrilProverService, MithrilSignedEntityService,
-        MithrilStakeDistributionService, ProverService, SignedEntityService,
-        StakeDistributionService, UpkeepService,
+        AggregatorSignableSeedBuilder, AggregatorUpkeepService, BufferedCertifierService,
+        CardanoTransactionsImporter, CertifierService, MessageService, MithrilCertifierService,
+        MithrilEpochService, MithrilMessageService, MithrilProverService,
+        MithrilSignedEntityService, MithrilStakeDistributionService, ProverService,
+        SignedEntityService, StakeDistributionService, UpkeepService,
     },
     tools::{CExplorerSignerRetriever, GcpFileUploader, GenesisToolsDependency, SignersImporter},
     AggregatorConfig, AggregatorRunner, AggregatorRuntime, CertificatePendingStore,
@@ -202,6 +203,9 @@ pub struct DependenciesBuilder {
     /// Signer Store
     pub signer_store: Option<Arc<SignerStore>>,
 
+    /// Signable Seed Builder
+    pub signable_seed_builder: Option<Arc<dyn SignableSeedBuilder>>,
+
     /// Signable Builder Service
     pub signable_builder_service: Option<Arc<dyn SignableBuilderService>>,
 
@@ -274,6 +278,7 @@ impl DependenciesBuilder {
             stake_distribution_service: None,
             ticker_service: None,
             signer_store: None,
+            signable_seed_builder: None,
             signable_builder_service: None,
             signed_entity_service: None,
             certifier_service: None,
@@ -1079,6 +1084,7 @@ impl DependenciesBuilder {
     }
 
     async fn build_signable_builder_service(&mut self) -> Result<Arc<dyn SignableBuilderService>> {
+        let seed_signable_builder = self.get_signable_seed_builder().await?;
         let mithril_stake_distribution_builder =
             Arc::new(MithrilStakeDistributionSignableBuilder::default());
         let immutable_signable_builder = Arc::new(CardanoImmutableFilesFullSignableBuilder::new(
@@ -1099,6 +1105,7 @@ impl DependenciesBuilder {
             CardanoStakeDistributionSignableBuilder::new(self.get_stake_store().await?),
         );
         let signable_builder_service = Arc::new(MithrilSignableBuilderService::new(
+            seed_signable_builder,
             mithril_stake_distribution_builder,
             immutable_signable_builder,
             cardano_transactions_builder,
@@ -1119,11 +1126,29 @@ impl DependenciesBuilder {
         Ok(self.signable_builder_service.as_ref().cloned().unwrap())
     }
 
+    async fn build_signable_seed_builder(&mut self) -> Result<Arc<dyn SignableSeedBuilder>> {
+        let signable_seed_builder_service = Arc::new(AggregatorSignableSeedBuilder::new(
+            self.get_epoch_service().await?,
+        ));
+
+        Ok(signable_seed_builder_service)
+    }
+
+    /// [SignableSeedBuilder] service
+    pub async fn get_signable_seed_builder(&mut self) -> Result<Arc<dyn SignableSeedBuilder>> {
+        if self.signable_seed_builder.is_none() {
+            self.signable_seed_builder = Some(self.build_signable_seed_builder().await?);
+        }
+
+        Ok(self.signable_seed_builder.as_ref().cloned().unwrap())
+    }
+
     async fn build_signed_entity_service(&mut self) -> Result<Arc<dyn SignedEntityService>> {
         let signed_entity_storer = self.build_signed_entity_storer().await?;
         let epoch_service = self.get_epoch_service().await?;
-        let mithril_stake_distribution_artifact_builder =
-            Arc::new(MithrilStakeDistributionArtifactBuilder::new(epoch_service));
+        let mithril_stake_distribution_artifact_builder = Arc::new(
+            MithrilStakeDistributionArtifactBuilder::new(epoch_service.clone()),
+        );
         let snapshotter = self.build_snapshotter().await?;
         let snapshot_uploader = self.build_snapshot_uploader().await?;
         let cardano_node_version = Version::parse(&self.configuration.cardano_node_version)

mithril-aggregator/src/runtime/runner.rs

@@ -5,8 +5,8 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use mithril_common::entities::{
-    Certificate, CertificatePending, Epoch, ProtocolMessage, ProtocolMessagePartKey,
-    SignedEntityConfig, SignedEntityType, Signer, TimePoint,
+    Certificate, CertificatePending, Epoch, ProtocolMessage, SignedEntityConfig, SignedEntityType,
+    Signer, TimePoint,
 };
 use mithril_common::StdResult;
 use mithril_persistence::store::StakeStorer;
@@ -284,22 +284,13 @@ impl AggregatorRunnerTrait for AggregatorRunner {
         signed_entity_type: &SignedEntityType,
     ) -> StdResult<ProtocolMessage> {
         debug!("RUNNER: compute protocol message");
-        let mut protocol_message = self
+        let protocol_message = self
             .dependencies
             .signable_builder_service
             .compute_protocol_message(signed_entity_type.to_owned())
             .await
             .with_context(|| format!("Runner can not compute protocol message for signed entity type: '{signed_entity_type}'"))?;
 
-        let epoch_service = self.dependencies.epoch_service.read().await;
-        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")?,
-        );
-
         Ok(protocol_message)
     }
 

mithril-aggregator/src/services/mod.rs

@@ -14,6 +14,7 @@ mod certifier;
 mod epoch_service;
 mod message;
 mod prover;
+mod signable_builder;
 mod signed_entity;
 mod stake_distribution;
 mod upkeep;
@@ -23,6 +24,7 @@ pub use certifier::*;
 pub use epoch_service::*;
 pub use message::*;
 pub use prover::*;
+pub use signable_builder::*;
 pub use signed_entity::*;
 pub use stake_distribution::*;
 pub use upkeep::*;

mithril-aggregator/src/services/signable_builder/mod.rs

@@ -0,0 +1,3 @@
+mod signable_seed_builder;
+
+pub use signable_seed_builder::*;

mithril-aggregator/src/services/signable_builder/signable_seed_builder.rs

@@ -0,0 +1,79 @@
+//! ## AggregatorSignableSeedBuilder
+//!
+//! This service is responsible for computing the seed protocol message
+//! that is used by the [SignableBuilder] to compute the final protocol message.
+//!
+use anyhow::Context;
+use async_trait::async_trait;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+
+use mithril_common::{
+    entities::ProtocolMessagePartValue, signable_builder::SignableSeedBuilder, StdResult,
+};
+
+use crate::services::EpochService;
+
+/// SignableSeedBuilder aggregator implementation
+pub struct AggregatorSignableSeedBuilder {
+    epoch_service: Arc<RwLock<dyn EpochService>>,
+}
+
+impl AggregatorSignableSeedBuilder {
+    /// AggregatorSignableSeedBuilder factory
+    pub fn new(epoch_service: Arc<RwLock<dyn EpochService>>) -> Self {
+        Self { epoch_service }
+    }
+}
+
+#[async_trait]
+impl SignableSeedBuilder for AggregatorSignableSeedBuilder {
+    async fn compute_next_aggregate_verification_key_protocol_message_part_value(
+        &self,
+    ) -> StdResult<ProtocolMessagePartValue> {
+        let epoch_service = self.epoch_service.read().await;
+        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(next_aggregate_verification_key)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use mithril_common::{entities::Epoch, test_utils::MithrilFixtureBuilder};
+
+    use crate::services::FakeEpochService;
+
+    use super::*;
+
+    #[tokio::test]
+    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();
+        let expected_next_aggregate_verification_key = next_fixture.compute_and_encode_avk();
+        let epoch_service = Arc::new(RwLock::new(FakeEpochService::with_data(
+            epoch,
+            &fixture.protocol_parameters(),
+            &next_fixture.protocol_parameters(),
+            &next_fixture.protocol_parameters(),
+            &fixture.signers_with_stake(),
+            &next_fixture.signers_with_stake(),
+        )));
+        let signable_seed_builder = AggregatorSignableSeedBuilder::new(epoch_service);
+
+        let next_aggregate_verification_key = signable_seed_builder
+            .compute_next_aggregate_verification_key_protocol_message_part_value()
+            .await
+            .unwrap();
+
+        assert_eq!(
+            next_aggregate_verification_key,
+            expected_next_aggregate_verification_key
+        );
+    }
+}

mithril-aggregator/tests/test_extensions/runtime_tester.rs

@@ -16,8 +16,8 @@ use mithril_common::{
     digesters::{DumbImmutableDigester, DumbImmutableFileObserver},
     entities::{
         BlockNumber, Certificate, CertificateSignature, ChainPoint, Epoch, ImmutableFileNumber,
-        ProtocolMessagePartKey, SignedEntityType, SignedEntityTypeDiscriminants,
-        SingleSignatureAuthenticationStatus, SlotNumber, Snapshot, StakeDistribution, TimePoint,
+        SignedEntityType, SignedEntityTypeDiscriminants, SingleSignatureAuthenticationStatus,
+        SlotNumber, Snapshot, StakeDistribution, TimePoint,
     },
     era::{adapters::EraReaderDummyAdapter, EraMarker, EraReader, SupportedEra},
     test_utils::{
@@ -417,24 +417,12 @@ impl RuntimeTester {
             .build_current_signed_entity_type(discriminant)
             .await?;
 
-        // Code copied from `AggregatorRunner::compute_protocol_message`
-        // Todo: Refactor this code to avoid code duplication by making the signable_builder_service
-        // able to retrieve the next avk by itself.
-        let mut message = self
+        let message = self
             .dependencies
             .signable_builder_service
             .compute_protocol_message(signed_entity_type.clone())
             .await?;
 
-        let epoch_service = self.dependencies.epoch_service.read().await;
-        message.set_message_part(
-            ProtocolMessagePartKey::NextAggregateVerificationKey,
-            epoch_service
-                .next_aggregate_verification_key()?
-                .to_json_hex()
-                .with_context(|| "convert next avk to json hex failure")?,
-        );
-
         for signer_fixture in signers {
             if let Some(mut single_signatures) = signer_fixture.sign(&message) {
                 if authentication_status == SingleSignatureAuthenticationStatus::Authenticated {

mithril-common/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-common"
-version = "0.4.57"
+version = "0.4.58"
 description = "Common types, interfaces, and utilities for Mithril nodes."
 authors = { workspace = true }
 edition = { workspace = true }

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;
@@ -16,7 +19,7 @@ pub trait Artifact: Debug + Send + Sync {
     fn get_id(&self) -> String;
 }
 
-/// SignableBuilder is trait for building a protocol message for a beacon
+/// SignableBuilder is a trait for building a protocol message for a beacon
 #[cfg_attr(test, automock)]
 #[async_trait]
 pub trait SignableBuilder<U>: Send + Sync
@@ -26,3 +29,13 @@ where
     /// Compute a protocol message
     async fn compute_protocol_message(&self, beacon: U) -> StdResult<ProtocolMessage>;
 }
+
+/// SignableSeedBuilder is a trait for building seed protocol message part values
+#[cfg_attr(test, automock)]
+#[async_trait]
+pub trait SignableSeedBuilder: Send + Sync {
+    /// Compute next aggregate verification key protocol message part value
+    async fn compute_next_aggregate_verification_key_protocol_message_part_value(
+        &self,
+    ) -> StdResult<ProtocolMessagePartValue>;
+}

mithril-common/src/signable_builder/mithril_stake_distribution.rs

@@ -12,7 +12,6 @@ pub struct MithrilStakeDistributionSignableBuilder {}
 
 #[async_trait]
 impl SignableBuilder<Epoch> for MithrilStakeDistributionSignableBuilder {
-    // We just need to return an empty protocol message as the next AVK will be appended by the signing engine automatically
     async fn compute_protocol_message(&self, _beacon: Epoch) -> StdResult<ProtocolMessage> {
         Ok(ProtocolMessage::new())
     }