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.into(),
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                    .into(),
113                cert_sig: operational_certificate_without_cold_verification_key
114                    .cert_sig()
115                    .to_bytes()
116                    .to_vec(),
117            },
118            cold_verification_key: cold_verification_key.to_bytes().to_vec(),
119        };
120
121        Ok(dmq_message.into())
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use mithril_cardano_node_chain::test::double::FakeChainObserver;
128    use mithril_common::{
129        crypto_helper::TryToBytes,
130        current_function,
131        entities::{BlockNumber, ChainPoint, TimePoint},
132        test::{crypto_helper::KesSignerFake, double::Dummy},
133    };
134
135    use crate::model::MockUnixTimestampProvider;
136
137    use super::*;
138
139    mod test_utils {
140        use super::*;
141
142        pub(super) struct TestMessage {
143            pub(super) content: Vec<u8>,
144        }
145
146        impl TryToBytes for TestMessage {
147            fn to_bytes_vec(&self) -> StdResult<Vec<u8>> {
148                Ok(self.content.clone())
149            }
150        }
151    }
152
153    #[tokio::test]
154    async fn test_build_dmq_message() {
155        let (kes_signature, operational_certificate) =
156            KesSignerFake::dummy_signature(current_function!());
157        let kes_signer = Arc::new(KesSignerFake::new(vec![Ok((
158            kes_signature,
159            operational_certificate.clone(),
160        ))]));
161        let chain_observer = Arc::new(FakeChainObserver::new(Some(TimePoint {
162            chain_point: ChainPoint {
163                block_number: BlockNumber(123),
164                ..ChainPoint::dummy()
165            },
166            ..TimePoint::dummy()
167        })));
168        let builder = DmqMessageBuilder::new(kes_signer.clone(), chain_observer)
169            .set_ttl(1000)
170            .set_timestamp_provider(Arc::new({
171                let mut mock_timestamp_provider = MockUnixTimestampProvider::new();
172                mock_timestamp_provider
173                    .expect_current_timestamp()
174                    .returning(|| Ok(234));
175
176                mock_timestamp_provider
177            }));
178        let message = test_utils::TestMessage {
179            content: b"test".to_vec(),
180        };
181
182        let dmq_message = builder.build(&message.to_bytes_vec().unwrap()).await.unwrap();
183
184        let expected_msg_payload = DmqMessageBuilder::enrich_msg_payload_with_id(DmqMsgPayload {
185            msg_id: vec![],
186            msg_body: b"test".to_vec(),
187            kes_period: 0,
188            expires_at: 1234,
189        });
190        assert_eq!(
191            DmqMsg {
192                msg_payload: expected_msg_payload.clone(),
193                kes_signature: kes_signature.to_bytes_vec().unwrap(),
194                operational_certificate: DmqMsgOperationalCertificate {
195                    kes_vk: operational_certificate
196                        .get_opcert_without_cold_verification_key()
197                        .kes_vk()
198                        .as_bytes()
199                        .to_vec(),
200                    issue_number: operational_certificate
201                        .get_opcert_without_cold_verification_key()
202                        .issue_number(),
203                    start_kes_period: operational_certificate
204                        .get_opcert_without_cold_verification_key()
205                        .start_kes_period()
206                        .into(),
207                    cert_sig: operational_certificate
208                        .get_opcert_without_cold_verification_key()
209                        .cert_sig()
210                        .to_bytes()
211                        .to_vec(),
212                },
213                cold_verification_key: operational_certificate
214                    .get_cold_verification_key()
215                    .to_bytes()
216                    .to_vec(),
217            },
218            dmq_message.clone().into()
219        );
220
221        let signed_messages = kes_signer.get_signed_messages();
222        assert_eq!(
223            vec![expected_msg_payload.bytes_to_sign().unwrap()],
224            signed_messages
225        );
226    }
227}
mithril_dmq/model/builder.rs

mithril_dmq/model/
builder.rs
