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

mithril_dmq/model/
builder.rs
