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::{KesSignerFake, TryToBytes},
96        entities::{BlockNumber, ChainPoint, TimePoint},
97    };
98
99    use super::*;
100
101    mod test_utils {
102        use super::*;
103
104        pub(super) struct TestMessage {
105            pub(super) content: Vec<u8>,
106        }
107
108        impl TryToBytes for TestMessage {
109            fn to_bytes_vec(&self) -> StdResult<Vec<u8>> {
110                Ok(self.content.clone())
111            }
112        }
113    }
114
115    #[tokio::test]
116    async fn test_build_dmq_message() {
117        let (kes_signature, operational_certificate) = KesSignerFake::dummy_signature();
118        let kes_signer = Arc::new(KesSignerFake::new(vec![Ok((
119            kes_signature,
120            operational_certificate.clone(),
121        ))]));
122        let chain_observer = Arc::new(FakeChainObserver::new(Some(TimePoint {
123            chain_point: ChainPoint {
124                block_number: BlockNumber(123),
125                ..ChainPoint::dummy()
126            },
127            ..TimePoint::dummy()
128        })));
129        let builder = DmqMessageBuilder::new(kes_signer, chain_observer).set_ttl(100);
130        let message = test_utils::TestMessage {
131            content: b"test".to_vec(),
132        };
133
134        let dmq_message = builder.build(&message.to_bytes_vec().unwrap()).await.unwrap();
135
136        assert!(!dmq_message.msg_id.is_empty());
137        assert_eq!(
138            DmqMsg {
139                msg_id: vec![],
140                msg_body: b"test".to_vec(),
141                block_number: 123,
142                ttl: 100,
143                kes_signature: kes_signature.to_bytes_vec().unwrap(),
144                operational_certificate: operational_certificate.to_bytes_vec().unwrap(),
145                kes_period: 0,
146            },
147            DmqMsg {
148                msg_id: vec![],
149                ..dmq_message
150            }
151        );
152    }
153}