Implement signable and artifact builders for Cardano Stake Distribution

CHANGELOG.md

@@ -9,6 +9,10 @@ As a minor extension, we have adopted a slightly different versioning convention
 
 ## Mithril Distribution [XXXX] - UNRELEASED
 
+- **UNSTABLE** Cardano stake distribution certification:
+
+  - Implement the signable and artifact builders for the signed entity type `CardanoStakeDistribution`
+
 - Crates versions:
 
 | Crate | Version |

internal/mithril-persistence/src/store/stake_store.rs

@@ -1,5 +1,6 @@
 use async_trait::async_trait;
 use mithril_common::entities::{Epoch, StakeDistribution};
+use mithril_common::signable_builder::StakeDistributionRetriever;
 use mithril_common::StdResult;
 use tokio::sync::RwLock;
 
@@ -79,6 +80,15 @@ impl StakeStorer for StakeStore {
     }
 }
 
+#[async_trait]
+impl StakeDistributionRetriever for StakeStore {
+    async fn retrieve(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>> {
+        let stake_distribution = self.get_stakes(epoch.offset_to_recording_epoch()).await?;
+
+        Ok(stake_distribution)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::super::adapter::MemoryAdapter;
@@ -167,4 +177,34 @@ mod tests {
             .unwrap();
         assert!(store.get_stakes(Epoch(1)).await.unwrap().is_none());
     }
+
+    #[tokio::test]
+    async fn retrieve_with_no_stakes_returns_none() {
+        let store = init_store(0, 0, None);
+
+        let result = store.retrieve(Epoch(1)).await.unwrap();
+
+        assert!(result.is_none());
+    }
+
+    #[tokio::test]
+    async fn retrieve_returns_cardano_stake_distribution_with_epoch_offset() {
+        let initial_stake_distribution = StakeDistribution::from([("pool-123".to_string(), 123)]);
+        let store = init_store(0, 0, None);
+        store
+            .save_stakes(Epoch(1), initial_stake_distribution.clone())
+            .await
+            .unwrap();
+
+        let mut next_stake_distribution = initial_stake_distribution;
+        next_stake_distribution.insert("pool-456".to_string(), 456);
+        store
+            .save_stakes(Epoch(2), next_stake_distribution.clone())
+            .await
+            .unwrap();
+
+        let stake_distribution = store.retrieve(Epoch(1)).await.unwrap();
+
+        assert_eq!(Some(next_stake_distribution), stake_distribution);
+    }
 }

mithril-aggregator/src/artifact_builder/cardano_stake_distribution.rs

@@ -0,0 +1,96 @@
+use anyhow::anyhow;
+use async_trait::async_trait;
+use std::sync::Arc;
+
+use mithril_common::{
+    entities::{CardanoStakeDistribution, Certificate, Epoch},
+    signable_builder::StakeDistributionRetriever,
+    StdResult,
+};
+
+use crate::ArtifactBuilder;
+
+/// A [CardanoStakeDistributionArtifact] builder
+pub struct CardanoStakeDistributionArtifactBuilder {
+    stake_distribution_retriever: Arc<dyn StakeDistributionRetriever>,
+}
+
+impl CardanoStakeDistributionArtifactBuilder {
+    /// CardanoStakeDistribution artifact builder factory
+    pub fn new(stake_distribution_retriever: Arc<dyn StakeDistributionRetriever>) -> Self {
+        Self {
+            stake_distribution_retriever,
+        }
+    }
+}
+
+#[async_trait]
+impl ArtifactBuilder<Epoch, CardanoStakeDistribution> for CardanoStakeDistributionArtifactBuilder {
+    async fn compute_artifact(
+        &self,
+        epoch: Epoch,
+        _certificate: &Certificate,
+    ) -> StdResult<CardanoStakeDistribution> {
+        let stake_distribution = self
+            .stake_distribution_retriever
+            .retrieve(epoch)
+            .await?
+            .ok_or_else(|| anyhow!("No stake distribution found for epoch '{}'", epoch))?;
+
+        Ok(CardanoStakeDistribution::new(epoch, stake_distribution))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use mithril_common::{entities::StakeDistribution, test_utils::fake_data};
+    use mockall::{mock, predicate::eq};
+
+    use super::*;
+
+    mock! {
+        pub StakeDistributionRetrieverImpl {}
+
+        #[async_trait]
+        impl StakeDistributionRetriever for StakeDistributionRetrieverImpl {
+            async fn retrieve(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>>;
+        }
+    }
+
+    #[tokio::test]
+    async fn compute_artifact_returns_valid_artifact_if_stakes_stored_for_epoch() {
+        let epoch = Epoch(1);
+        let certificate = fake_data::certificate("whatever".to_string());
+        let stake_distribution = StakeDistribution::from([("pool-123".to_string(), 123)]);
+        let stake_distribution_clone = stake_distribution.clone();
+        let mut mock_storer = MockStakeDistributionRetrieverImpl::new();
+        mock_storer
+            .expect_retrieve()
+            .with(eq(epoch))
+            .return_once(move |_| Ok(Some(stake_distribution_clone)));
+        let builder = CardanoStakeDistributionArtifactBuilder::new(Arc::new(mock_storer));
+
+        let cardano_stake_distribution =
+            builder.compute_artifact(epoch, &certificate).await.unwrap();
+
+        let expected = CardanoStakeDistribution::new(epoch, stake_distribution);
+        assert_eq!(cardano_stake_distribution, expected);
+    }
+
+    #[tokio::test]
+    async fn compute_artifact_returns_error_if_no_stakes_found_for_epoch() {
+        let epoch = Epoch(1);
+        let certificate = fake_data::certificate("whatever".to_string());
+        let mut mock_storer = MockStakeDistributionRetrieverImpl::new();
+        mock_storer
+            .expect_retrieve()
+            .with(eq(epoch))
+            .return_once(move |_| Ok(None));
+        let builder = CardanoStakeDistributionArtifactBuilder::new(Arc::new(mock_storer));
+
+        builder
+            .compute_artifact(epoch, &certificate)
+            .await
+            .expect_err("Should return error");
+    }
+}

mithril-aggregator/src/artifact_builder/mod.rs

@@ -1,10 +1,12 @@
 //! The module used for building artifact
 mod cardano_immutable_files_full;
+mod cardano_stake_distribution;
 mod cardano_transactions;
 mod interface;
 mod mithril_stake_distribution;
 
 pub use cardano_immutable_files_full::*;
+pub use cardano_stake_distribution::*;
 pub use cardano_transactions::*;
 pub use interface::*;
 pub use mithril_stake_distribution::*;

mithril-aggregator/src/database/repository/open_message_repository.rs

@@ -45,7 +45,7 @@ impl OpenMessageRepository {
     ) -> StdResult<Option<OpenMessageWithSingleSignaturesRecord>> {
         self.connection.fetch_first(
             GetOpenMessageWithSingleSignaturesQuery::by_epoch_and_signed_entity_type(
-                signed_entity_type.get_epoch(),
+                signed_entity_type.get_epoch_when_signed_entity_type_is_signed(),
                 signed_entity_type,
             )?,
         )

mithril-aggregator/src/database/repository/stake_pool_store.rs

@@ -5,6 +5,7 @@ use anyhow::Context;
 use async_trait::async_trait;
 
 use mithril_common::entities::{Epoch, StakeDistribution};
+use mithril_common::signable_builder::StakeDistributionRetriever;
 use mithril_common::StdResult;
 use mithril_persistence::sqlite::{ConnectionExtensions, SqliteConnection};
 use mithril_persistence::store::adapter::AdapterError;
@@ -87,6 +88,15 @@ impl StakeStorer for StakePoolStore {
     }
 }
 
+#[async_trait]
+impl StakeDistributionRetriever for StakePoolStore {
+    async fn retrieve(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>> {
+        let stake_distribution = self.get_stakes(epoch.offset_to_recording_epoch()).await?;
+
+        Ok(stake_distribution)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::database::test_helper::{insert_stake_pool, main_db_connection};
@@ -120,4 +130,36 @@ mod tests {
             "Stakes at epoch 2 should still exist",
         );
     }
+
+    #[tokio::test]
+    async fn retrieve_with_no_stakes_returns_none() {
+        let connection = main_db_connection().unwrap();
+        let store = StakePoolStore::new(Arc::new(connection), None);
+
+        let result = store.retrieve(Epoch(1)).await.unwrap();
+
+        assert!(result.is_none());
+    }
+
+    #[tokio::test]
+    async fn retrieve_returns_cardano_stake_distribution_with_epoch_offset() {
+        let initial_stake_distribution = StakeDistribution::from([("pool-123".to_string(), 123)]);
+        let connection = main_db_connection().unwrap();
+        let store = StakePoolStore::new(Arc::new(connection), None);
+        store
+            .save_stakes(Epoch(1), initial_stake_distribution.clone())
+            .await
+            .unwrap();
+
+        let mut next_stake_distribution = initial_stake_distribution;
+        next_stake_distribution.insert("pool-456".to_string(), 456);
+        store
+            .save_stakes(Epoch(2), next_stake_distribution.clone())
+            .await
+            .unwrap();
+
+        let stake_distribution = store.retrieve(Epoch(1)).await.unwrap();
+
+        assert_eq!(Some(next_stake_distribution), stake_distribution);
+    }
 }

mithril-aggregator/src/dependency_injection/builder.rs

@@ -37,9 +37,9 @@ use mithril_common::{
         EraChecker, EraMarker, EraReader, EraReaderAdapter, SupportedEra,
     },
     signable_builder::{
-        CardanoImmutableFilesFullSignableBuilder, CardanoTransactionsSignableBuilder,
-        MithrilSignableBuilderService, MithrilStakeDistributionSignableBuilder,
-        SignableBuilderService, TransactionsImporter,
+        CardanoImmutableFilesFullSignableBuilder, CardanoStakeDistributionSignableBuilder,
+        CardanoTransactionsSignableBuilder, MithrilSignableBuilderService,
+        MithrilStakeDistributionSignableBuilder, SignableBuilderService, TransactionsImporter,
     },
     signed_entity_type_lock::SignedEntityTypeLock,
     MithrilTickerService, TickerService,
@@ -52,8 +52,8 @@ use mithril_persistence::{
 
 use crate::{
     artifact_builder::{
-        CardanoImmutableFilesFullArtifactBuilder, CardanoTransactionsArtifactBuilder,
-        MithrilStakeDistributionArtifactBuilder,
+        CardanoImmutableFilesFullArtifactBuilder, CardanoStakeDistributionArtifactBuilder,
+        CardanoTransactionsArtifactBuilder, MithrilStakeDistributionArtifactBuilder,
     },
     configuration::ExecutionEnvironment,
     database::repository::{
@@ -1090,10 +1090,14 @@ impl DependenciesBuilder {
             block_range_root_retriever,
             self.get_logger()?,
         ));
+        let cardano_stake_distribution_builder = Arc::new(
+            CardanoStakeDistributionSignableBuilder::new(self.get_stake_store().await?),
+        );
         let signable_builder_service = Arc::new(MithrilSignableBuilderService::new(
             mithril_stake_distribution_builder,
             immutable_signable_builder,
             cardano_transactions_builder,
+            cardano_stake_distribution_builder,
         ));
 
         Ok(signable_builder_service)
@@ -1130,12 +1134,16 @@ impl DependenciesBuilder {
         let cardano_transactions_artifact_builder = Arc::new(
             CardanoTransactionsArtifactBuilder::new(prover_service.clone()),
         );
+        let stake_store = self.get_stake_store().await?;
+        let cardano_stake_distribution_artifact_builder =
+            Arc::new(CardanoStakeDistributionArtifactBuilder::new(stake_store));
         let signed_entity_service = Arc::new(MithrilSignedEntityService::new(
             signed_entity_storer,
             mithril_stake_distribution_artifact_builder,
             cardano_immutable_files_full_artifact_builder,
             cardano_transactions_artifact_builder,
             self.get_signed_entity_lock().await?,
+            cardano_stake_distribution_artifact_builder,
         ));
 
         // Compute the cache pool for prover service

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

@@ -52,11 +52,17 @@ mod handlers {
 
         let signed_entity_type = match message.signed_entity_type.clone() {
             Some(signed_entity_type) => Ok(signed_entity_type),
+            // The '.unwrap()' inside the 'None' branch is safe
+            // This case does not occur anymore, as there are no signers using the Mithril Signer node
+            // version that sends a 'RegisterSignatureMessage' without the 'signed_entity_type' field.
+            // Make the 'signed_entity_type' field mandatory in the 'RegisterSignatureMessage' will be done in a new PR.
             None => ticker_service.get_current_time_point().await.map(|t| {
-                signed_entity_config.time_point_to_signed_entity(
-                    SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
-                    &t,
-                )
+                signed_entity_config
+                    .time_point_to_signed_entity(
+                        SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
+                        &t,
+                    )
+                    .unwrap()
             }),
         };
 

mithril-aggregator/src/runtime/runner.rs

@@ -147,15 +147,18 @@ impl AggregatorRunner {
     async fn list_available_signed_entity_types(
         &self,
         time_point: &TimePoint,
-    ) -> Vec<SignedEntityType> {
+    ) -> StdResult<Vec<SignedEntityType>> {
         let signed_entity_types = self
             .dependencies
             .signed_entity_config
-            .list_allowed_signed_entity_types(time_point);
-        self.dependencies
+            .list_allowed_signed_entity_types(time_point)?;
+        let unlocked_signed_entities = self
+            .dependencies
             .signed_entity_type_lock
             .filter_unlocked_entries(signed_entity_types)
-            .await
+            .await;
+
+        Ok(unlocked_signed_entities)
     }
 }
 
@@ -197,7 +200,7 @@ impl AggregatorRunnerTrait for AggregatorRunner {
         debug!("RUNNER: get_current_non_certified_open_message"; "time_point" => #?current_time_point);
         let signed_entity_types = self
             .list_available_signed_entity_types(current_time_point)
-            .await;
+            .await?;
         for signed_entity_type in signed_entity_types {
             let current_open_message = self.get_current_open_message_for_signed_entity_type(&signed_entity_type)
                 .await
@@ -1192,6 +1195,7 @@ pub mod tests {
         let signed_entities: Vec<SignedEntityTypeDiscriminants> = runner
             .list_available_signed_entity_types(&time_point)
             .await
+            .unwrap()
             .into_iter()
             .map(Into::into)
             .collect();
@@ -1223,6 +1227,7 @@ pub mod tests {
         let signed_entities: Vec<SignedEntityTypeDiscriminants> = runner
             .list_available_signed_entity_types(&time_point)
             .await
+            .unwrap()
             .into_iter()
             .map(Into::into)
             .collect();

mithril-aggregator/src/runtime/state_machine.rs

@@ -240,7 +240,7 @@ impl AggregatorRuntime {
                         .time_point_to_signed_entity(
                             &state.open_message.signed_entity_type,
                             &last_time_point,
-                        );
+                        )?;
                     new_signed_entity_type != state.open_message.signed_entity_type
                 };