mithril_dmq/
message.rs

1use std::sync::Arc;
2
3use anyhow::{Context, anyhow};
4use blake2::{Blake2b, Digest, digest::consts::U64};
5use pallas_network::miniprotocols::localmsgsubmission::DmqMsg;
6
7use mithril_cardano_node_chain::chain_observer::ChainObserver;
8use mithril_common::{
9    StdResult,
10    crypto_helper::{KesSigner, TryToBytes},
11};
12
13/// The TTL (Time To Live) for DMQ messages in blocks.
14const DMQ_MESSAGE_TTL_IN_BLOCKS: u16 = 100;
15
16/// A builder for creating DMQ messages.
17pub struct DmqMessageBuilder {
18    kes_signer: Arc<dyn KesSigner>,
19    chain_observer: Arc<dyn ChainObserver>,
20    ttl_blocks: u16,
21}
22
23impl DmqMessageBuilder {
24    /// Creates a new instance of `DmqMessageBuilder`.
25    pub fn new(kes_signer: Arc<dyn KesSigner>, chain_observer: Arc<dyn ChainObserver>) -> Self {
26        Self {
27            kes_signer,
28            chain_observer,
29            ttl_blocks: DMQ_MESSAGE_TTL_IN_BLOCKS,
30        }
31    }
32
33    /// Set the TTL (Time To Live) for DMQ messages in blocks.
34    pub fn set_ttl(mut self, ttl_blocks: u16) -> Self {
35        self.ttl_blocks = ttl_blocks;
36
37        self
38    }
39
40    /// Builds a DMQ message from the provided message bytes.
41    pub async fn build(&self, message_bytes: &[u8]) -> StdResult<DmqMsg> {
42        fn compute_msg_id(dmq_message: &DmqMsg) -> Vec<u8> {
43            let mut hasher = Blake2b::<U64>::new();
44            hasher.update(&dmq_message.msg_body);
45            hasher.update(dmq_message.block_number.to_be_bytes());
46            hasher.update(dmq_message.ttl.to_be_bytes());
47            hasher.update(&dmq_message.kes_signature);
48            hasher.update(&dmq_message.operational_certificate);
49            hasher.update(dmq_message.kes_period.to_be_bytes());
50
51            hasher.finalize().to_vec()
52        }
53
54        let block_number = self
55            .chain_observer
56            .get_current_chain_point()
57            .await
58            .with_context(|| "Failed to get current chain point while building DMQ message")?
59            .ok_or(anyhow!(
60                "No current chain point available while building DMQ message"
61            ))?
62            .block_number;
63        let block_number = (*block_number)
64            .try_into()
65            .with_context(|| "Failed to convert block number to u32")?;
66        let (kes_signature, operational_certificate) = self
67            .kes_signer
68            .sign(message_bytes, block_number)
69            .with_context(|| "Failed to KES sign message while building DMQ message")?;
70        let kes_period = self
71            .chain_observer
72            .get_current_kes_period(&operational_certificate)
73            .await
74            .with_context(|| "Failed to get KES period while building DMQ message")?
75            .unwrap_or_default();
76        let mut dmq_message = DmqMsg {
77            msg_id: vec![],
78            msg_body: message_bytes.to_vec(),
79            block_number,
80            ttl: self.ttl_blocks,
81            kes_signature: kes_signature.to_bytes_vec()?,
82            operational_certificate: operational_certificate.to_bytes_vec()?,
83            kes_period,
84        };
85        dmq_message.msg_id = compute_msg_id(&dmq_message);
86
87        Ok(dmq_message)
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use mithril_cardano_node_chain::test::double::FakeChainObserver;
94    use mithril_common::{
95        crypto_helper::TryToBytes,
96        entities::{BlockNumber, ChainPoint, TimePoint},
97        test::{crypto_helper::KesSignerFake, double::Dummy},
98    };
99
100    use super::*;
101
102    mod test_utils {
103        use super::*;
104
105        pub(super) struct TestMessage {
106            pub(super) content: Vec<u8>,
107        }
108
109        impl TryToBytes for TestMessage {
110            fn to_bytes_vec(&self) -> StdResult<Vec<u8>> {
111                Ok(self.content.clone())
112            }
113        }
114    }
115
116    #[tokio::test]
117    async fn test_build_dmq_message() {
118        let (kes_signature, operational_certificate) = KesSignerFake::dummy_signature();
119        let kes_signer = Arc::new(KesSignerFake::new(vec![Ok((
120            kes_signature,
121            operational_certificate.clone(),
122        ))]));
123        let chain_observer = Arc::new(FakeChainObserver::new(Some(TimePoint {
124            chain_point: ChainPoint {
125                block_number: BlockNumber(123),
126                ..ChainPoint::dummy()
127            },
128            ..TimePoint::dummy()
129        })));
130        let builder = DmqMessageBuilder::new(kes_signer, chain_observer).set_ttl(100);
131        let message = test_utils::TestMessage {
132            content: b"test".to_vec(),
133        };
134
135        let dmq_message = builder.build(&message.to_bytes_vec().unwrap()).await.unwrap();
136
137        assert!(!dmq_message.msg_id.is_empty());
138        assert_eq!(
139            DmqMsg {
140                msg_id: vec![],
141                msg_body: b"test".to_vec(),
142                block_number: 123,
143                ttl: 100,
144                kes_signature: kes_signature.to_bytes_vec().unwrap(),
145                operational_certificate: operational_certificate.to_bytes_vec().unwrap(),
146                kes_period: 0,
147            },
148            DmqMsg {
149                msg_id: vec![],
150                ..dmq_message
151            }
152        );
153    }
154}