Cleanup protocol message from multi-signer

Author: jpraynaud

mithril-aggregator/src/multi_signer.rs

@@ -1,5 +1,4 @@
 use async_trait::async_trait;
-use hex::ToHex;
 use slog_scope::{debug, warn};
 use std::sync::Arc;
 use thiserror::Error;
@@ -78,15 +77,6 @@ pub enum ProtocolError {
 #[cfg_attr(test, automock)]
 #[async_trait]
 pub trait MultiSigner: Sync + Send {
-    /// Get current message
-    async fn get_current_message(&self) -> Option<entities::ProtocolMessage>;
-
-    /// Update current message
-    async fn update_current_message(
-        &mut self,
-        message: entities::ProtocolMessage,
-    ) -> Result<(), ProtocolError>;
-
     /// Get current beacon
     async fn get_current_beacon(&self) -> Option<entities::Beacon>;
 
@@ -200,9 +190,6 @@ pub trait MultiSigner: Sync + Send {
 
 /// MultiSignerImpl is an implementation of the MultiSigner
 pub struct MultiSignerImpl {
-    /// Message that is currently signed
-    current_message: Option<entities::ProtocolMessage>,
-
     /// Beacon that is currently used
     current_beacon: Option<entities::Beacon>,
 
@@ -225,7 +212,6 @@ impl MultiSignerImpl {
     ) -> Self {
         debug!("New MultiSignerImpl created");
         Self {
-            current_message: None,
             current_beacon: None,
             verification_key_store,
             stake_store,
@@ -321,23 +307,6 @@ impl MultiSignerImpl {
 
 #[async_trait]
 impl MultiSigner for MultiSignerImpl {
-    /// Get current message
-    async fn get_current_message(&self) -> Option<entities::ProtocolMessage> {
-        debug!("Get current message");
-        self.current_message.clone()
-    }
-
-    /// Update current message
-    async fn update_current_message(
-        &mut self,
-        message: entities::ProtocolMessage,
-    ) -> Result<(), ProtocolError> {
-        debug!("Update current_message"; "protocol_message" =>  #?message, "signed message" => message.compute_hash().encode_hex::<String>());
-
-        self.current_message = Some(message);
-        Ok(())
-    }
-
     async fn get_current_beacon(&self) -> Option<entities::Beacon> {
         self.current_beacon.clone()
     }
@@ -698,23 +667,6 @@ mod tests {
         result
     }
 
-    #[tokio::test]
-    async fn test_multi_signer_current_message_ok() {
-        let mut multi_signer = setup_multi_signer().await;
-
-        let current_message_expected = setup_message();
-        multi_signer
-            .update_current_message(current_message_expected.clone())
-            .await
-            .expect("update current message failed");
-
-        let current_message = multi_signer
-            .get_current_message()
-            .await
-            .expect("current message should have been retrieved");
-        assert_eq!(current_message_expected, current_message)
-    }
-
     #[tokio::test]
     async fn test_multi_signer_protocol_parameters_ok() {
         let mut multi_signer = setup_multi_signer().await;
@@ -816,14 +768,7 @@ mod tests {
             Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
         )
         .await;
-        // We have to update the current message AFTER we reached the epoch for
-        // which the signers registered for the stake distribution to be valid
-        // hence the multisigner be able to create a Clerk and being able to
-        // register single signatures
-        multi_signer
-            .update_current_message(message.clone())
-            .await
-            .expect("update current message failed");
+
         let mut signatures = Vec::new();
 
         let mut expected_certificate_signers: Vec<SignerWithStake> = Vec::new();

mithril-aggregator/src/runtime/runner.rs

@@ -100,10 +100,8 @@ pub trait AggregatorRunnerTrait: Sync + Send {
         new_beacon: &Beacon,
     ) -> Result<(), Box<dyn StdError + Sync + Send>>;
 
-    /// Set the message to sign in the multisigner. The digest is only one part
-    /// of the message, the next signing stake distribution must also be signed
-    /// as part of the message.
-    async fn update_message_in_multisigner(
+    /// Compute the protocol message
+    async fn compute_protocol_message(
         &self,
         digest: String,
     ) -> Result<ProtocolMessage, Box<dyn StdError + Sync + Send>>;
@@ -138,6 +136,7 @@ pub trait AggregatorRunnerTrait: Sync + Send {
     async fn create_snapshot_archive(
         &self,
         beacon: &Beacon,
+        protocol_message: &ProtocolMessage,
     ) -> Result<OngoingSnapshot, Box<dyn StdError + Sync + Send>>;
 
     /// Upload the snapshot at the given location using the configured uploader(s).
@@ -378,12 +377,12 @@ impl AggregatorRunnerTrait for AggregatorRunner {
             .map_err(|e| e.into())
     }
 
-    async fn update_message_in_multisigner(
+    async fn compute_protocol_message(
         &self,
         digest: String,
     ) -> Result<ProtocolMessage, Box<dyn StdError + Sync + Send>> {
-        debug!("RUNNER: update message in multisigner");
-        let mut multi_signer = self.dependencies.multi_signer.write().await;
+        debug!("RUNNER: compute protocol message");
+        let multi_signer = self.dependencies.multi_signer.write().await;
         let mut protocol_message = ProtocolMessage::new();
         protocol_message.set_message_part(ProtocolMessagePartKey::SnapshotDigest, digest);
         protocol_message.set_message_part(
@@ -393,9 +392,6 @@ impl AggregatorRunnerTrait for AggregatorRunner {
                 .await?
                 .unwrap_or_default(),
         );
-        multi_signer
-            .update_current_message(protocol_message.clone())
-            .await?;
 
         Ok(protocol_message)
     }
@@ -491,22 +487,11 @@ impl AggregatorRunnerTrait for AggregatorRunner {
     async fn create_snapshot_archive(
         &self,
         beacon: &Beacon,
+        protocol_message: &ProtocolMessage,
     ) -> Result<OngoingSnapshot, Box<dyn StdError + Sync + Send>> {
         debug!("RUNNER: create snapshot archive");
 
         let snapshotter = self.dependencies.snapshotter.clone();
-        let protocol_message = self
-            .dependencies
-            .multi_signer
-            .read()
-            .await
-            .get_current_message()
-            .await
-            .ok_or_else(|| {
-                RunnerError::MissingProtocolMessage(format!(
-                    "no message found for beacon '{beacon:?}'."
-                ))
-            })?;
         let snapshot_digest = protocol_message
             .get_message_part(&ProtocolMessagePartKey::SnapshotDigest)
             .ok_or_else(|| {
@@ -899,41 +884,6 @@ pub mod tests {
         assert_eq!(expected, certificate);
     }
 
-    #[tokio::test]
-    async fn test_update_message_in_multisigner() {
-        let (deps, config) = initialize_dependencies().await;
-        let deps = Arc::new(deps);
-        let runner = AggregatorRunner::new(config, deps.clone());
-        let digest = "1+2+3+4=10".to_string();
-        runner.update_beacon(&fake_data::beacon()).await.unwrap();
-        let fixture = MithrilFixtureBuilder::default().build();
-        deps.prepare_for_genesis(
-            fixture.signers_with_stake(),
-            fixture.signers_with_stake(),
-            &fixture.protocol_parameters(),
-        )
-        .await;
-
-        runner
-            .update_message_in_multisigner(digest)
-            .await
-            .expect("update_message_in_multisigner should not fail");
-        let message = deps
-            .multi_signer
-            .read()
-            .await
-            .get_current_message()
-            .await
-            .unwrap();
-
-        assert_eq!(
-            "1+2+3+4=10",
-            message
-                .get_message_part(&ProtocolMessagePartKey::SnapshotDigest)
-                .unwrap()
-        );
-    }
-
     #[tokio::test]
     async fn test_save_pending_certificate() {
         let (deps, config) = initialize_dependencies().await;
@@ -1012,16 +962,13 @@ pub mod tests {
             ProtocolMessagePartKey::SnapshotDigest,
             "test+digest".to_string(),
         );
-        let mut mock_multi_signer = MockMultiSigner::new();
-        mock_multi_signer
-            .expect_get_current_message()
-            .return_once(move || Some(message));
+        let mock_multi_signer = MockMultiSigner::new();
         let (mut deps, config) = initialize_dependencies().await;
         deps.multi_signer = Arc::new(RwLock::new(mock_multi_signer));
         let runner = AggregatorRunner::new(config, Arc::new(deps));
 
         let ongoing_snapshot = runner
-            .create_snapshot_archive(&beacon)
+            .create_snapshot_archive(&beacon, &message)
             .await
             .expect("create_snapshot_archive should not fail");
 

mithril-aggregator/src/runtime/state_machine.rs

@@ -1,4 +1,7 @@
-use crate::runtime::{AggregatorRunnerTrait, RuntimeError};
+use crate::{
+    database::provider::OpenMessage,
+    runtime::{AggregatorRunnerTrait, RuntimeError},
+};
 
 use mithril_common::entities::{Beacon, SignedEntityType};
 use slog_scope::{crit, info, trace, warn};
@@ -19,6 +22,7 @@ pub struct ReadyState {
 #[derive(Clone, Debug, PartialEq)]
 pub struct SigningState {
     current_beacon: Beacon,
+    open_message: OpenMessage,
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -270,12 +274,9 @@ impl AggregatorRuntime {
         state: SigningState,
     ) -> Result<IdleState, RuntimeError> {
         trace!("launching transition from SIGNING to IDLE state");
-        // TODO: Temporary, we need to compute the signed entity type from the current open message
-        let signed_entity_type =
-            SignedEntityType::CardanoImmutableFilesFull(state.current_beacon.clone());
         let certificate = self
             .runner
-            .create_certificate(&signed_entity_type)
+            .create_certificate(&state.open_message.signed_entity_type)
             .await?
             .ok_or_else(|| RuntimeError::KeepState {
                 message: "not enough signature yet to create a certificate, waiting…".to_string(),
@@ -285,7 +286,7 @@ impl AggregatorRuntime {
         self.runner.drop_pending_certificate().await?;
         let ongoing_snapshot = self
             .runner
-            .create_snapshot_archive(&state.current_beacon)
+            .create_snapshot_archive(&state.current_beacon, &state.open_message.protocol_message)
             .await?;
         let locations = self
             .runner
@@ -323,16 +324,16 @@ impl AggregatorRuntime {
         new_beacon: Beacon,
     ) -> Result<SigningState, RuntimeError> {
         trace!("launching transition from READY to SIGNING state");
-        // TODO: Temporary, we need to compute the signed entity type from the current open message
+        // TODO: Temporary, we need to compute the signed entity type for other types than Cardano immutable files
         let signed_entity_type = SignedEntityType::CardanoImmutableFilesFull(new_beacon.clone());
         self.runner.update_beacon(&new_beacon).await?;
 
         let digester_result = self.runner.compute_digest(&new_beacon).await?;
         let protocol_message = self
             .runner
-            .update_message_in_multisigner(digester_result)
+            .compute_protocol_message(digester_result)
             .await?;
-        let _open_message = self
+        let open_message = self
             .runner
             .create_open_message(&signed_entity_type, &protocol_message)
             .await?;
@@ -348,6 +349,7 @@ impl AggregatorRuntime {
             .await?;
         let state = SigningState {
             current_beacon: new_beacon,
+            open_message,
         };
 
         Ok(state)
@@ -589,7 +591,7 @@ mod tests {
             .once()
             .returning(|_| Ok(()));
         runner
-            .expect_update_message_in_multisigner()
+            .expect_compute_protocol_message()
             .with(predicate::eq("whatever".to_string()))
             .once()
             .returning(|_| Ok(ProtocolMessage::new()));
@@ -640,6 +642,7 @@ mod tests {
 
                 beacon
             },
+            open_message: OpenMessage::dummy(),
         };
         let mut runtime = init_runtime(Some(AggregatorState::Signing(state)), runner).await;
         runtime.cycle().await.unwrap();
@@ -660,6 +663,7 @@ mod tests {
             .returning(|_| Ok(None));
         let state = SigningState {
             current_beacon: fake_data::beacon(),
+            open_message: OpenMessage::dummy(),
         };
         let mut runtime = init_runtime(Some(AggregatorState::Signing(state)), runner).await;
         runtime

mithril-aggregator/tests/test_extensions/runtime_tester.rs

@@ -15,11 +15,11 @@ use mithril_aggregator::{
     AggregatorRuntime, Configuration, DumbSnapshotUploader, DumbSnapshotter,
     ProtocolParametersStorer, SignerRegisterer,
 };
-use mithril_common::crypto_helper::{key_encode_hex, ProtocolClerk, ProtocolGenesisSigner};
+use mithril_common::crypto_helper::{ProtocolClerk, ProtocolGenesisSigner};
 use mithril_common::digesters::DumbImmutableFileObserver;
 use mithril_common::entities::{
-    Certificate, Epoch, ImmutableFileNumber, SignedEntityType, SignerWithStake, SingleSignatures,
-    Snapshot, StakeDistribution,
+    Certificate, Epoch, ImmutableFileNumber, SignedEntityType, SignerWithStake, Snapshot,
+    StakeDistribution,
 };
 use mithril_common::{chain_observer::FakeObserver, digesters::DumbImmutableDigester};
 
@@ -268,24 +268,15 @@ impl RuntimeTester {
         signers: &[SignerFixture],
     ) -> Result<(), String> {
         let certifier_service = self.deps_builder.get_certifier_service().await.unwrap();
-        let lock = self.deps_builder.get_multi_signer().await.unwrap();
-        let multisigner = lock.read().await;
-        let message = multisigner
-            .get_current_message()
+        let message = certifier_service
+            .get_open_message(signed_entity_type)
             .await
-            .ok_or("There should be a message to be signed.")?;
+            .unwrap()
+            .ok_or("There should be a message to be signed.")?
+            .protocol_message;
 
         for signer_fixture in signers {
-            if let Some(signature) = signer_fixture
-                .protocol_signer
-                .sign(message.compute_hash().as_bytes())
-            {
-                let single_signatures = SingleSignatures::new(
-                    signer_fixture.signer_with_stake.party_id.to_owned(),
-                    key_encode_hex(&signature).expect("hex encoding should not fail"),
-                    signature.indexes,
-                );
-
+            if let Some(single_signatures) = signer_fixture.sign(&message) {
                 certifier_service
                     .register_single_signature(signed_entity_type, &single_signatures)
                     .await