refactor(dmq): update the message builder and its KES signature

Author: jpraynaud

internal/mithril-dmq/src/model/builder.rs

@@ -1,25 +1,26 @@
 use std::sync::Arc;
 
-use anyhow::{Context, anyhow};
+use anyhow::Context;
 use blake2::{Blake2b, Digest, digest::consts::U64};
-use pallas_network::miniprotocols::localmsgsubmission::DmqMsg;
+use pallas_network::miniprotocols::localmsgsubmission::{DmqMsg, DmqMsgPayload};
 
 use mithril_cardano_node_chain::chain_observer::ChainObserver;
 use mithril_common::{
     StdResult,
     crypto_helper::{KesSigner, TryToBytes},
 };
 
-use crate::model::DmqMessage;
+use crate::model::{DmqMessage, SystemUnixTimestampProvider, UnixTimestampProvider};
 
-/// The TTL (Time To Live) for DMQ messages in blocks.
-const DMQ_MESSAGE_TTL_IN_BLOCKS: u16 = 100;
+/// The TTL (Time To Live) for DMQ messages in seconds (default is 30 minutes).
+const DMQ_MESSAGE_TTL_IN_SECONDS: u16 = 1800;
 
 /// A builder for creating DMQ messages.
 pub struct DmqMessageBuilder {
     kes_signer: Arc<dyn KesSigner>,
     chain_observer: Arc<dyn ChainObserver>,
-    ttl_blocks: u16,
+    timestamp_provider: Arc<dyn UnixTimestampProvider>,
+    ttl_seconds: u16,
 }
 
 impl DmqMessageBuilder {
@@ -28,43 +29,45 @@ impl DmqMessageBuilder {
         Self {
             kes_signer,
             chain_observer,
-            ttl_blocks: DMQ_MESSAGE_TTL_IN_BLOCKS,
+            timestamp_provider: Arc::new(SystemUnixTimestampProvider),
+            ttl_seconds: DMQ_MESSAGE_TTL_IN_SECONDS,
         }
     }
 
-    /// Set the TTL (Time To Live) for DMQ messages in blocks.
-    pub fn set_ttl(mut self, ttl_blocks: u16) -> Self {
-        self.ttl_blocks = ttl_blocks;
+    /// Sets the timestamp provider for the DMQ message builder.
+    pub fn set_timestamp_provider(
+        mut self,
+        timestamp_provider: Arc<dyn UnixTimestampProvider>,
+    ) -> Self {
+        self.timestamp_provider = timestamp_provider;
+        self
+    }
+
+    /// Sets the TTL (Time To Live) for DMQ messages in seconds.
+    pub fn set_ttl(mut self, ttl_seconds: u16) -> Self {
+        self.ttl_seconds = ttl_seconds;
 
         self
     }
 
+    /// Computes a message id for a DMQ message payload.
+    fn compute_msg_id(dmq_message_payload: &DmqMsgPayload) -> Vec<u8> {
+        let mut hasher = Blake2b::<U64>::new();
+        hasher.update(&dmq_message_payload.msg_body);
+        hasher.update(dmq_message_payload.kes_period.to_be_bytes());
+        hasher.update(&dmq_message_payload.operational_certificate);
+        hasher.update(&dmq_message_payload.cold_verification_key);
+        hasher.update(dmq_message_payload.expires_at.to_be_bytes());
+
+        hasher.finalize().to_vec()
+    }
+
     /// Builds a DMQ message from the provided message bytes.
     pub async fn build(&self, message_bytes: &[u8]) -> StdResult<DmqMessage> {
-        fn compute_msg_id(dmq_message: &DmqMsg) -> Vec<u8> {
-            let mut hasher = Blake2b::<U64>::new();
-            hasher.update(&dmq_message.msg_body);
-            hasher.update(dmq_message.block_number.to_be_bytes());
-            hasher.update(dmq_message.ttl.to_be_bytes());
-            hasher.update(&dmq_message.kes_signature);
-            hasher.update(&dmq_message.operational_certificate);
-            hasher.update(dmq_message.kes_period.to_be_bytes());
-
-            hasher.finalize().to_vec()
-        }
-
-        let block_number = self
-            .chain_observer
-            .get_current_chain_point()
-            .await
-            .with_context(|| "Failed to get current chain point while building DMQ message")?
-            .ok_or(anyhow!(
-                "No current chain point available while building DMQ message"
-            ))?
-            .block_number;
-        let block_number = (*block_number)
+        let expires_at: u32 = (self.timestamp_provider.current_timestamp()?
+            + self.ttl_seconds as u64)
             .try_into()
-            .with_context(|| "Failed to convert block number to u32")?;
+            .with_context(|| "Failed to compute expires_at while building DMQ message")?;
         let kes_period = self
             .chain_observer
             .get_current_kes_period()
@@ -75,16 +78,23 @@ impl DmqMessageBuilder {
             .kes_signer
             .sign(message_bytes, kes_period)
             .with_context(|| "Failed to KES sign message while building DMQ message")?;
-        let mut dmq_message = DmqMsg {
-            msg_id: vec![],
-            msg_body: message_bytes.to_vec(),
-            block_number,
-            ttl: self.ttl_blocks,
+
+        let dmq_message = DmqMsg {
+            msg_payload: {
+                let mut dmq_message_payload = DmqMsgPayload {
+                    msg_id: vec![],
+                    msg_body: message_bytes.to_vec(),
+                    kes_period: kes_period as u64,
+                    operational_certificate: operational_certificate.to_bytes_vec()?, // TODO: remove the cold verification key in the op cert
+                    cold_verification_key: vec![],
+                    expires_at,
+                };
+                dmq_message_payload.msg_id = Self::compute_msg_id(&dmq_message_payload);
+
+                dmq_message_payload
+            },
             kes_signature: kes_signature.to_bytes_vec()?,
-            operational_certificate: operational_certificate.to_bytes_vec()?,
-            kes_period,
         };
-        dmq_message.msg_id = compute_msg_id(&dmq_message);
 
         Ok(dmq_message.into())
     }
@@ -100,6 +110,8 @@ mod tests {
         test::{crypto_helper::KesSignerFake, double::Dummy},
     };
 
+    use crate::model::MockUnixTimestampProvider;
+
     use super::*;
 
     mod test_utils {
@@ -131,28 +143,39 @@ mod tests {
             },
             ..TimePoint::dummy()
         })));
-        let builder = DmqMessageBuilder::new(kes_signer, chain_observer).set_ttl(100);
+        let builder = DmqMessageBuilder::new(kes_signer, chain_observer)
+            .set_ttl(1000)
+            .set_timestamp_provider(Arc::new({
+                let mut mock_timestamp_provider = MockUnixTimestampProvider::new();
+                mock_timestamp_provider
+                    .expect_current_timestamp()
+                    .returning(|| Ok(234));
+
+                mock_timestamp_provider
+            }));
         let message = test_utils::TestMessage {
             content: b"test".to_vec(),
         };
 
         let dmq_message = builder.build(&message.to_bytes_vec().unwrap()).await.unwrap();
 
-        assert!(!dmq_message.msg_id.is_empty());
+        let DmqMsg {
+            msg_payload,
+            kes_signature: _,
+        } = &*dmq_message;
         assert_eq!(
             DmqMsg {
-                msg_id: vec![],
-                msg_body: b"test".to_vec(),
-                block_number: 123,
-                ttl: 100,
+                msg_payload: DmqMsgPayload {
+                    msg_id: DmqMessageBuilder::compute_msg_id(msg_payload),
+                    msg_body: b"test".to_vec(),
+                    kes_period: 0,
+                    operational_certificate: operational_certificate.to_bytes_vec().unwrap(),
+                    cold_verification_key: vec![], // TODO: fix
+                    expires_at: 1234,
+                },
                 kes_signature: kes_signature.to_bytes_vec().unwrap(),
-                operational_certificate: operational_certificate.to_bytes_vec().unwrap(),
-                kes_period: 0,
             },
-            DmqMsg {
-                msg_id: vec![],
-                ..dmq_message.into()
-            }
+            dmq_message.into()
         );
     }
 }

internal/mithril-dmq/src/model/message.rs

@@ -72,18 +72,23 @@ impl<'de> Deserialize<'de> for DmqMessage {
 
 #[cfg(test)]
 mod tests {
+    use pallas_network::miniprotocols::localmsgsubmission::DmqMsgPayload;
+
     use super::*;
 
     #[test]
     fn test_dmq_message_serialize_deserialize() {
         let dmq_msg = DmqMsg {
-            msg_id: vec![0, 1],
-            msg_body: vec![0, 1, 2],
+            msg_payload: DmqMsgPayload {
+                msg_id: vec![0, 1],
+                msg_body: vec![0, 1, 2],
+
+                kes_period: 10,
+                operational_certificate: vec![0, 1, 2, 3, 4],
+                cold_verification_key: vec![0, 1, 2, 3, 4, 5],
+                expires_at: 100,
+            },
             kes_signature: vec![0, 1, 2, 3],
-            kes_period: 10,
-            operational_certificate: vec![0, 1, 2, 3, 4],
-            cold_verification_key: vec![0, 1, 2, 3, 4, 5],
-            expires_at: 100,
         };
 
         let dmq_message = DmqMessage::from(dmq_msg.clone());