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::{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}