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
13const DMQ_MESSAGE_TTL_IN_BLOCKS: u16 = 100;
15
16pub struct DmqMessageBuilder {
18 kes_signer: Arc<dyn KesSigner>,
19 chain_observer: Arc<dyn ChainObserver>,
20 ttl_blocks: u16,
21}
22
23impl DmqMessageBuilder {
24 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 pub fn set_ttl(mut self, ttl_blocks: u16) -> Self {
35 self.ttl_blocks = ttl_blocks;
36
37 self
38 }
39
40 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}