1use serde::{Deserialize, Serialize};
2use sha2::{Digest, Sha256};
3use std::{collections::BTreeMap, fmt::Display};
4
5#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
7pub enum ProtocolMessagePartKey {
8 #[serde(rename = "snapshot_digest")]
10 SnapshotDigest,
11
12 #[serde(rename = "cardano_transactions_merkle_root")]
14 CardanoTransactionsMerkleRoot,
15
16 #[serde(rename = "next_aggregate_verification_key")]
21 NextAggregateVerificationKey,
22
23 #[serde(rename = "next_protocol_parameters")]
28 NextProtocolParameters,
29
30 #[serde(rename = "current_epoch")]
34 CurrentEpoch,
35
36 #[serde(rename = "latest_block_number")]
38 LatestBlockNumber,
39
40 #[serde(rename = "cardano_stake_distribution_epoch")]
42 CardanoStakeDistributionEpoch,
43
44 #[serde(rename = "cardano_stake_distribution_merkle_root")]
46 CardanoStakeDistributionMerkleRoot,
47
48 #[serde(rename = "cardano_database_merkle_root")]
50 CardanoDatabaseMerkleRoot,
51}
52
53impl Display for ProtocolMessagePartKey {
54 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55 match *self {
56 Self::SnapshotDigest => write!(f, "snapshot_digest"),
57 Self::NextAggregateVerificationKey => write!(f, "next_aggregate_verification_key"),
58 Self::NextProtocolParameters => write!(f, "next_protocol_parameters"),
59 Self::CurrentEpoch => write!(f, "current_epoch"),
60 Self::CardanoTransactionsMerkleRoot => write!(f, "cardano_transactions_merkle_root"),
61 Self::LatestBlockNumber => write!(f, "latest_block_number"),
62 Self::CardanoStakeDistributionEpoch => write!(f, "cardano_stake_distribution_epoch"),
63 Self::CardanoStakeDistributionMerkleRoot => {
64 write!(f, "cardano_stake_distribution_merkle_root")
65 }
66 Self::CardanoDatabaseMerkleRoot => write!(f, "cardano_database_merkle_root"),
67 }
68 }
69}
70
71pub type ProtocolMessagePartValue = String;
73
74#[derive(Clone, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
76pub struct ProtocolMessage {
77 pub message_parts: BTreeMap<ProtocolMessagePartKey, ProtocolMessagePartValue>,
80}
81
82impl ProtocolMessage {
83 pub fn new() -> ProtocolMessage {
85 ProtocolMessage {
86 message_parts: BTreeMap::new(),
87 }
88 }
89
90 pub fn set_message_part(
93 &mut self,
94 key: ProtocolMessagePartKey,
95 value: ProtocolMessagePartValue,
96 ) -> Option<ProtocolMessagePartValue> {
97 self.message_parts.insert(key, value)
98 }
99
100 pub fn get_message_part(
102 &self,
103 key: &ProtocolMessagePartKey,
104 ) -> Option<&ProtocolMessagePartValue> {
105 self.message_parts.get(key)
106 }
107
108 pub fn compute_hash(&self) -> String {
110 let mut hasher = Sha256::new();
111 self.message_parts.iter().for_each(|(k, v)| {
112 hasher.update(k.to_string().as_bytes());
113 hasher.update(v.as_bytes());
114 });
115 hex::encode(hasher.finalize())
116 }
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122
123 #[test]
124 fn test_protocol_message_compute_hash_include_next_aggregate_verification_key() {
125 let protocol_message = build_protocol_message_reference();
126 let hash_expected = protocol_message.compute_hash();
127
128 let mut protocol_message_modified = protocol_message.clone();
129 protocol_message_modified.set_message_part(
130 ProtocolMessagePartKey::NextAggregateVerificationKey,
131 "next-avk-456".to_string(),
132 );
133
134 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
135 }
136
137 #[test]
138 fn test_protocol_message_compute_hash_include_snapshot_digest() {
139 let protocol_message = build_protocol_message_reference();
140 let hash_expected = protocol_message.compute_hash();
141
142 let mut protocol_message_modified = protocol_message.clone();
143 protocol_message_modified.set_message_part(
144 ProtocolMessagePartKey::SnapshotDigest,
145 "snapshot-digest-456".to_string(),
146 );
147
148 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
149 }
150
151 #[test]
152 fn test_protocol_message_compute_hash_include_cardano_transactions_merkle_root() {
153 let protocol_message = build_protocol_message_reference();
154 let hash_expected = protocol_message.compute_hash();
155
156 let mut protocol_message_modified = protocol_message.clone();
157 protocol_message_modified.set_message_part(
158 ProtocolMessagePartKey::CardanoTransactionsMerkleRoot,
159 "ctx-merke-root-456".to_string(),
160 );
161
162 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
163 }
164
165 #[test]
166 fn test_protocol_message_compute_hash_include_cardano_stake_distribution_epoch() {
167 let protocol_message = build_protocol_message_reference();
168 let hash_expected = protocol_message.compute_hash();
169
170 let mut protocol_message_modified = protocol_message.clone();
171 protocol_message_modified.set_message_part(
172 ProtocolMessagePartKey::CardanoStakeDistributionEpoch,
173 "cardano-stake-distribution-epoch-456".to_string(),
174 );
175
176 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
177 }
178
179 #[test]
180 fn test_protocol_message_compute_hash_include_cardano_stake_distribution_merkle_root() {
181 let protocol_message = build_protocol_message_reference();
182 let hash_expected = protocol_message.compute_hash();
183
184 let mut protocol_message_modified = protocol_message.clone();
185 protocol_message_modified.set_message_part(
186 ProtocolMessagePartKey::CardanoStakeDistributionMerkleRoot,
187 "cardano-stake-distribution-merkle-root-456".to_string(),
188 );
189
190 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
191 }
192
193 #[test]
194 fn test_protocol_message_compute_hash_include_lastest_immutable_file_number() {
195 let protocol_message = build_protocol_message_reference();
196 let hash_expected = protocol_message.compute_hash();
197
198 let mut protocol_message_modified = protocol_message.clone();
199 protocol_message_modified.set_message_part(
200 ProtocolMessagePartKey::LatestBlockNumber,
201 "latest-immutable-file-number-456".to_string(),
202 );
203
204 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
205 }
206
207 #[test]
208 fn test_protocol_message_compute_hash_include_cardano_database_merkle_root() {
209 let protocol_message = build_protocol_message_reference();
210 let hash_expected = protocol_message.compute_hash();
211
212 let mut protocol_message_modified = protocol_message.clone();
213 protocol_message_modified.set_message_part(
214 ProtocolMessagePartKey::CardanoDatabaseMerkleRoot,
215 "cardano-database-merkle-root-456".to_string(),
216 );
217
218 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
219 }
220
221 #[test]
222 fn test_protocol_message_compute_hash_include_next_protocol_parameters() {
223 let protocol_message = build_protocol_message_reference();
224 let hash_expected = protocol_message.compute_hash();
225
226 let mut protocol_message_modified = protocol_message.clone();
227 protocol_message_modified.set_message_part(
228 ProtocolMessagePartKey::NextProtocolParameters,
229 "latest-protocol-parameters-456".to_string(),
230 );
231
232 assert_ne!(hash_expected, protocol_message_modified.compute_hash());
233 }
234
235 #[test]
236 fn test_protocol_message_compute_hash_the_same_hash_with_same_protocol_message() {
237 assert_eq!(
238 build_protocol_message_reference().compute_hash(),
239 build_protocol_message_reference().compute_hash()
240 );
241 }
242
243 fn build_protocol_message_reference() -> ProtocolMessage {
244 let mut protocol_message = ProtocolMessage::new();
245 protocol_message.set_message_part(
246 ProtocolMessagePartKey::SnapshotDigest,
247 "snapshot-digest-123".to_string(),
248 );
249 protocol_message.set_message_part(
250 ProtocolMessagePartKey::NextAggregateVerificationKey,
251 "next-avk-123".to_string(),
252 );
253 protocol_message.set_message_part(
254 ProtocolMessagePartKey::NextProtocolParameters,
255 "next-protocol-parameters-123".to_string(),
256 );
257 protocol_message.set_message_part(
258 ProtocolMessagePartKey::CardanoTransactionsMerkleRoot,
259 "ctx-merkle-root-123".to_string(),
260 );
261 protocol_message.set_message_part(
262 ProtocolMessagePartKey::LatestBlockNumber,
263 "latest-immutable-file-number-123".to_string(),
264 );
265 protocol_message.set_message_part(
266 ProtocolMessagePartKey::CardanoStakeDistributionEpoch,
267 "cardano-stake-distribution-epoch-123".to_string(),
268 );
269 protocol_message.set_message_part(
270 ProtocolMessagePartKey::CardanoStakeDistributionMerkleRoot,
271 "cardano-stake-distribution-merkle-root-123".to_string(),
272 );
273 protocol_message.set_message_part(
274 ProtocolMessagePartKey::CardanoDatabaseMerkleRoot,
275 "cardano-database-merkle-root-123".to_string(),
276 );
277
278 protocol_message
279 }
280}