mithril_dmq/model/
builder.rs

1use std::sync::Arc;
2
3use anyhow::Context;
4use blake2::{Blake2b, Digest, digest::consts::U64};
5use pallas_network::miniprotocols::localmsgsubmission::{
6    DmqMsg, DmqMsgOperationalCertificate, DmqMsgPayload,
7};
8
9use mithril_cardano_node_chain::chain_observer::ChainObserver;
10use mithril_common::{
11    StdResult,
12    crypto_helper::{KesSigner, TryToBytes},
13};
14
15use crate::model::{DmqMessage, SystemUnixTimestampProvider, UnixTimestampProvider};
16
17/// The TTL (Time To Live) for DMQ messages in seconds (default is 30 minutes).
18const DMQ_MESSAGE_TTL_IN_SECONDS: u16 = 1800;
19
20/// A builder for creating DMQ messages.
21pub struct DmqMessageBuilder {
22    kes_signer: Arc<dyn KesSigner>,
23    chain_observer: Arc<dyn ChainObserver>,
24    timestamp_provider: Arc<dyn UnixTimestampProvider>,
25    ttl_seconds: u16,
26}
27
28impl DmqMessageBuilder {
29    /// Creates a new instance of `DmqMessageBuilder`.
30    pub fn new(kes_signer: Arc<dyn KesSigner>, chain_observer: Arc<dyn ChainObserver>) -> Self {
31        Self {
32            kes_signer,
33            chain_observer,
34            timestamp_provider: Arc::new(SystemUnixTimestampProvider),
35            ttl_seconds: DMQ_MESSAGE_TTL_IN_SECONDS,
36        }
37    }
38
39    /// Sets the timestamp provider for the DMQ message builder.
40    pub fn set_timestamp_provider(
41        mut self,
42        timestamp_provider: Arc<dyn UnixTimestampProvider>,
43    ) -> Self {
44        self.timestamp_provider = timestamp_provider;
45        self
46    }
47
48    /// Sets the TTL (Time To Live) for DMQ messages in seconds.
49    pub fn set_ttl(mut self, ttl_seconds: u16) -> Self {
50        self.ttl_seconds = ttl_seconds;
51
52        self
53    }
54
55    /// Computes a message id for a DMQ message payload.
56    fn compute_msg_id(dmq_message_payload: &DmqMsgPayload) -> Vec<u8> {
57        let mut hasher = Blake2b::<U64>::new();
58        hasher.update(&dmq_message_payload.msg_body);
59        hasher.update(dmq_message_payload.kes_period.to_be_bytes());
60        hasher.update(dmq_message_payload.expires_at.to_be_bytes());
61
62        hasher.finalize().to_vec()
63    }
64
65    /// Enriches a DMQ message payload with a computed message ID.
66    fn enrich_msg_payload_with_id(dmq_message_payload: DmqMsgPayload) -> DmqMsgPayload {
67        let msg_id = Self::compute_msg_id(&dmq_message_payload);
68        let mut dmq_message_payload_with_id = dmq_message_payload;
69        dmq_message_payload_with_id.msg_id = msg_id;
70
71        dmq_message_payload_with_id
72    }
73
74    /// Builds a DMQ message from the provided message bytes.
75    pub async fn build(&self, message_bytes: &[u8]) -> StdResult<DmqMessage> {
76        let expires_at: u32 = (self.timestamp_provider.current_timestamp()?
77            + self.ttl_seconds as u64)
78            .try_into()
79            .with_context(|| "Failed to compute expires_at while building DMQ message")?;
80        let kes_period = self
81            .chain_observer
82            .get_current_kes_period()
83            .await
84            .with_context(|| "Failed to get KES period while building DMQ message")?
85            .unwrap_or_default();
86        let dmq_message_payload = Self::enrich_msg_payload_with_id(DmqMsgPayload {
87            msg_id: vec![],
88            msg_body: message_bytes.to_vec(),
89            kes_period: kes_period as u64,
90            expires_at,
91        });
92
93        let (kes_signature, operational_certificate) = self
94            .kes_signer
95            .sign(&dmq_message_payload.bytes_to_sign()?, kes_period)
96            .with_context(|| "Failed to KES sign message while building DMQ message")?;
97        let operational_certificate_without_cold_verification_key =
98            operational_certificate.get_opcert_without_cold_verification_key();
99        let cold_verification_key = operational_certificate.get_cold_verification_key();
100
101        let dmq_message = DmqMsg {
102            msg_payload: dmq_message_payload,
103            kes_signature: kes_signature.to_bytes_vec()?,
104            operational_certificate: DmqMsgOperationalCertificate {
105                kes_vk: operational_certificate_without_cold_verification_key
106                    .kes_vk()
107                    .as_bytes()
108                    .to_vec(),
109                issue_number: operational_certificate_without_cold_verification_key.issue_number(),
110                start_kes_period: operational_certificate_without_cold_verification_key
111                    .start_kes_period(),
112                cert_sig: operational_certificate_without_cold_verification_key
113                    .cert_sig()
114                    .to_bytes()
115                    .to_vec(),
116            },
117            cold_verification_key: cold_verification_key.to_bytes().to_vec(),
118        };
119
120        Ok(dmq_message.into())
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use mithril_cardano_node_chain::test::double::FakeChainObserver;
127    use mithril_common::{
128        crypto_helper::TryToBytes,
129        current_function,
130        entities::{BlockNumber, ChainPoint, TimePoint},
131        test::{crypto_helper::KesSignerFake, double::Dummy},
132    };
133
134    use crate::model::MockUnixTimestampProvider;
135
136    use super::*;
137
138    mod test_utils {
139        use super::*;
140
141        pub(super) struct TestMessage {
142            pub(super) content: Vec<u8>,
143        }
144
145        impl TryToBytes for TestMessage {
146            fn to_bytes_vec(&self) -> StdResult<Vec<u8>> {
147                Ok(self.content.clone())
148            }
149        }
150    }
151
152    #[tokio::test]
153    async fn test_build_dmq_message() {
154        let (kes_signature, operational_certificate) =
155            KesSignerFake::dummy_signature(current_function!());
156        let kes_signer = Arc::new(KesSignerFake::new(vec![Ok((
157            kes_signature,
158            operational_certificate.clone(),
159        ))]));
160        let chain_observer = Arc::new(FakeChainObserver::new(Some(TimePoint {
161            chain_point: ChainPoint {
162                block_number: BlockNumber(123),
163                ..ChainPoint::dummy()
164            },
165            ..TimePoint::dummy()
166        })));
167        let builder = DmqMessageBuilder::new(kes_signer.clone(), chain_observer)
168            .set_ttl(1000)
169            .set_timestamp_provider(Arc::new({
170                let mut mock_timestamp_provider = MockUnixTimestampProvider::new();
171                mock_timestamp_provider
172                    .expect_current_timestamp()
173                    .returning(|| Ok(234));
174
175                mock_timestamp_provider
176            }));
177        let message = test_utils::TestMessage {
178            content: b"test".to_vec(),
179        };
180
181        let dmq_message = builder.build(&message.to_bytes_vec().unwrap()).await.unwrap();
182
183        let expected_msg_payload = DmqMessageBuilder::enrich_msg_payload_with_id(DmqMsgPayload {
184            msg_id: vec![],
185            msg_body: b"test".to_vec(),
186            kes_period: 0,
187            expires_at: 1234,
188        });
189        assert_eq!(
190            DmqMsg {
191                msg_payload: expected_msg_payload.clone(),
192                kes_signature: kes_signature.to_bytes_vec().unwrap(),
193                operational_certificate: DmqMsgOperationalCertificate {
194                    kes_vk: operational_certificate
195                        .get_opcert_without_cold_verification_key()
196                        .kes_vk()
197                        .as_bytes()
198                        .to_vec(),
199                    issue_number: operational_certificate
200                        .get_opcert_without_cold_verification_key()
201                        .issue_number(),
202                    start_kes_period: operational_certificate
203                        .get_opcert_without_cold_verification_key()
204                        .start_kes_period(),
205                    cert_sig: operational_certificate
206                        .get_opcert_without_cold_verification_key()
207                        .cert_sig()
208                        .to_bytes()
209                        .to_vec(),
210                },
211                cold_verification_key: operational_certificate
212                    .get_cold_verification_key()
213                    .to_bytes()
214                    .to_vec(),
215            },
216            dmq_message.clone().into()
217        );
218
219        let signed_messages = kes_signer.get_signed_messages();
220        assert_eq!(
221            vec![expected_msg_payload.bytes_to_sign().unwrap()],
222            signed_messages
223        );
224    }
225}
mithril_dmq/model/builder.rs

mithril_dmq/model/
builder.rs
