1use std::collections::BTreeSet;
2use std::str::FromStr;
3use std::time::Duration;
4
5use anyhow::anyhow;
6use digest::Update;
7use serde::{Deserialize, Serialize};
8use sha2::Sha256;
9use strum::{AsRefStr, Display, EnumDiscriminants, EnumIter, EnumString, IntoEnumIterator};
10
11use crate::{
12 StdResult,
13 crypto_helper::{TryFromBytes, TryToBytes},
14};
15
16use super::{BlockNumber, CardanoDbBeacon, Epoch};
17
18const ENTITY_TYPE_MITHRIL_STAKE_DISTRIBUTION: usize = 0;
20
21const ENTITY_TYPE_CARDANO_STAKE_DISTRIBUTION: usize = 1;
23
24const ENTITY_TYPE_CARDANO_IMMUTABLE_FILES_FULL: usize = 2;
26
27const ENTITY_TYPE_CARDANO_TRANSACTIONS: usize = 3;
29
30const ENTITY_TYPE_CARDANO_DATABASE: usize = 4;
32
33#[derive(Display, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, EnumDiscriminants)]
41#[strum(serialize_all = "PascalCase")]
42#[strum_discriminants(doc = "The discriminants of the SignedEntityType enum.")]
43#[strum_discriminants(derive(
44 Display,
45 EnumString,
46 AsRefStr,
47 Serialize,
48 Deserialize,
49 PartialOrd,
50 Ord,
51 EnumIter,
52))]
53pub enum SignedEntityType {
54 MithrilStakeDistribution(Epoch),
56
57 CardanoStakeDistribution(Epoch),
59
60 CardanoImmutableFilesFull(CardanoDbBeacon),
62
63 CardanoDatabase(CardanoDbBeacon),
65
66 CardanoTransactions(Epoch, BlockNumber),
68}
69
70impl SignedEntityType {
71 pub fn genesis(epoch: Epoch) -> Self {
73 Self::MithrilStakeDistribution(epoch)
74 }
75
76 pub fn get_epoch(&self) -> Epoch {
78 match self {
79 Self::CardanoImmutableFilesFull(b) | Self::CardanoDatabase(b) => b.epoch,
80 Self::CardanoStakeDistribution(e)
81 | Self::MithrilStakeDistribution(e)
82 | Self::CardanoTransactions(e, _) => *e,
83 }
84 }
85
86 pub fn get_epoch_when_signed_entity_type_is_signed(&self) -> Epoch {
88 match self {
89 Self::CardanoImmutableFilesFull(beacon) | Self::CardanoDatabase(beacon) => beacon.epoch,
90 Self::CardanoStakeDistribution(epoch) => epoch.next(),
91 Self::MithrilStakeDistribution(epoch) | Self::CardanoTransactions(epoch, _) => *epoch,
92 }
93 }
94
95 pub fn index(&self) -> usize {
97 match self {
98 Self::MithrilStakeDistribution(_) => ENTITY_TYPE_MITHRIL_STAKE_DISTRIBUTION,
99 Self::CardanoStakeDistribution(_) => ENTITY_TYPE_CARDANO_STAKE_DISTRIBUTION,
100 Self::CardanoImmutableFilesFull(_) => ENTITY_TYPE_CARDANO_IMMUTABLE_FILES_FULL,
101 Self::CardanoTransactions(_, _) => ENTITY_TYPE_CARDANO_TRANSACTIONS,
102 Self::CardanoDatabase(_) => ENTITY_TYPE_CARDANO_DATABASE,
103 }
104 }
105
106 pub fn get_json_beacon(&self) -> StdResult<String> {
108 let value = match self {
109 Self::CardanoImmutableFilesFull(value) | Self::CardanoDatabase(value) => {
110 serde_json::to_string(value)?
111 }
112 Self::CardanoStakeDistribution(value) | Self::MithrilStakeDistribution(value) => {
113 serde_json::to_string(value)?
114 }
115 Self::CardanoTransactions(epoch, block_number) => {
116 let json = serde_json::json!({
117 "epoch": epoch,
118 "block_number": block_number,
119 });
120 serde_json::to_string(&json)?
121 }
122 };
123
124 Ok(value)
125 }
126
127 pub fn get_open_message_timeout(&self) -> Option<Duration> {
129 match self {
130 Self::MithrilStakeDistribution(_) | Self::CardanoImmutableFilesFull(_) => None,
131 Self::CardanoStakeDistribution(_) => Some(Duration::from_secs(600)),
132 Self::CardanoTransactions(_, _) => Some(Duration::from_secs(1800)),
133 Self::CardanoDatabase(_) => Some(Duration::from_secs(1800)),
134 }
135 }
136
137 pub(crate) fn feed_hash(&self, hasher: &mut Sha256) {
138 match self {
139 SignedEntityType::MithrilStakeDistribution(epoch)
140 | SignedEntityType::CardanoStakeDistribution(epoch) => {
141 hasher.update(&epoch.to_be_bytes())
142 }
143 SignedEntityType::CardanoImmutableFilesFull(db_beacon)
144 | SignedEntityType::CardanoDatabase(db_beacon) => {
145 hasher.update(&db_beacon.epoch.to_be_bytes());
146 hasher.update(&db_beacon.immutable_file_number.to_be_bytes());
147 }
148 SignedEntityType::CardanoTransactions(epoch, block_number) => {
149 hasher.update(&epoch.to_be_bytes());
150 hasher.update(&block_number.to_be_bytes())
151 }
152 }
153 }
154}
155
156impl TryFromBytes for SignedEntityType {
157 fn try_from_bytes(bytes: &[u8]) -> StdResult<Self> {
158 let (res, _) =
159 bincode::serde::decode_from_slice::<Self, _>(bytes, bincode::config::standard())?;
160
161 Ok(res)
162 }
163}
164
165impl TryToBytes for SignedEntityType {
166 fn to_bytes_vec(&self) -> StdResult<Vec<u8>> {
167 bincode::serde::encode_to_vec(self, bincode::config::standard()).map_err(|e| e.into())
168 }
169}
170
171impl SignedEntityTypeDiscriminants {
172 pub fn all() -> BTreeSet<Self> {
174 SignedEntityTypeDiscriminants::iter().collect()
175 }
176
177 pub fn index(&self) -> usize {
179 match self {
180 Self::MithrilStakeDistribution => ENTITY_TYPE_MITHRIL_STAKE_DISTRIBUTION,
181 Self::CardanoStakeDistribution => ENTITY_TYPE_CARDANO_STAKE_DISTRIBUTION,
182 Self::CardanoImmutableFilesFull => ENTITY_TYPE_CARDANO_IMMUTABLE_FILES_FULL,
183 Self::CardanoTransactions => ENTITY_TYPE_CARDANO_TRANSACTIONS,
184 Self::CardanoDatabase => ENTITY_TYPE_CARDANO_DATABASE,
185 }
186 }
187
188 pub fn from_id(signed_entity_type_id: usize) -> StdResult<SignedEntityTypeDiscriminants> {
190 match signed_entity_type_id {
191 ENTITY_TYPE_MITHRIL_STAKE_DISTRIBUTION => Ok(Self::MithrilStakeDistribution),
192 ENTITY_TYPE_CARDANO_STAKE_DISTRIBUTION => Ok(Self::CardanoStakeDistribution),
193 ENTITY_TYPE_CARDANO_IMMUTABLE_FILES_FULL => Ok(Self::CardanoImmutableFilesFull),
194 ENTITY_TYPE_CARDANO_TRANSACTIONS => Ok(Self::CardanoTransactions),
195 ENTITY_TYPE_CARDANO_DATABASE => Ok(Self::CardanoDatabase),
196 index => Err(anyhow!("Invalid entity_type_id {index}.")),
197 }
198 }
199
200 pub fn parse_list<T: AsRef<str>>(discriminants_string: T) -> StdResult<BTreeSet<Self>> {
205 let mut discriminants = BTreeSet::new();
206 let mut invalid_discriminants = Vec::new();
207
208 for name in discriminants_string
209 .as_ref()
210 .split(',')
211 .map(str::trim)
212 .filter(|s| !s.is_empty())
213 {
214 match Self::from_str(name) {
215 Ok(discriminant) => {
216 discriminants.insert(discriminant);
217 }
218 Err(_) => {
219 invalid_discriminants.push(name);
220 }
221 }
222 }
223
224 if invalid_discriminants.is_empty() {
225 Ok(discriminants)
226 } else {
227 Err(anyhow!(Self::format_parse_list_error(
228 invalid_discriminants
229 )))
230 }
231 }
232
233 fn format_parse_list_error(invalid_discriminants: Vec<&str>) -> String {
234 format!(
235 r#"Invalid signed entity types discriminants: {}.
236
237Accepted values are (case-sensitive): {}."#,
238 invalid_discriminants.join(", "),
239 Self::accepted_discriminants()
240 )
241 }
242
243 fn accepted_discriminants() -> String {
244 Self::iter().map(|d| d.to_string()).collect::<Vec<_>>().join(", ")
245 }
246}
247
248#[cfg(test)]
249mod tests {
250 use digest::Digest;
251
252 use crate::test::assert_same_json;
253
254 use super::*;
255
256 #[test]
257 fn get_epoch_when_signed_entity_type_is_signed_for_cardano_stake_distribution_return_epoch_with_offset()
258 {
259 let signed_entity_type = SignedEntityType::CardanoStakeDistribution(Epoch(3));
260
261 assert_eq!(
262 signed_entity_type.get_epoch_when_signed_entity_type_is_signed(),
263 Epoch(4)
264 );
265 }
266
267 #[test]
268 fn get_epoch_when_signed_entity_type_is_signed_for_mithril_stake_distribution_return_epoch_stored_in_signed_entity_type()
269 {
270 let signed_entity_type = SignedEntityType::MithrilStakeDistribution(Epoch(3));
271 assert_eq!(
272 signed_entity_type.get_epoch_when_signed_entity_type_is_signed(),
273 Epoch(3)
274 );
275 }
276
277 #[test]
278 fn get_epoch_when_signed_entity_type_is_signed_for_cardano_immutable_files_full_return_epoch_stored_in_signed_entity_type()
279 {
280 let signed_entity_type =
281 SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::new(3, 100));
282 assert_eq!(
283 signed_entity_type.get_epoch_when_signed_entity_type_is_signed(),
284 Epoch(3)
285 );
286 }
287
288 #[test]
289 fn get_epoch_when_signed_entity_type_is_signed_for_cardano_transactions_return_epoch_stored_in_signed_entity_type()
290 {
291 let signed_entity_type = SignedEntityType::CardanoTransactions(Epoch(3), BlockNumber(77));
292 assert_eq!(
293 signed_entity_type.get_epoch_when_signed_entity_type_is_signed(),
294 Epoch(3)
295 );
296 }
297
298 #[test]
299 fn get_epoch_when_signed_entity_type_is_signed_for_cardano_database_return_epoch_stored_in_signed_entity_type()
300 {
301 let signed_entity_type = SignedEntityType::CardanoDatabase(CardanoDbBeacon::new(12, 987));
302 assert_eq!(
303 signed_entity_type.get_epoch_when_signed_entity_type_is_signed(),
304 Epoch(12)
305 );
306 }
307
308 #[test]
309 fn verify_signed_entity_type_properties_are_included_in_computed_hash() {
310 fn hash(signed_entity_type: SignedEntityType) -> String {
311 let mut hasher = Sha256::new();
312 signed_entity_type.feed_hash(&mut hasher);
313 hex::encode(hasher.finalize())
314 }
315
316 let reference_hash = hash(SignedEntityType::MithrilStakeDistribution(Epoch(5)));
317 assert_ne!(
318 reference_hash,
319 hash(SignedEntityType::MithrilStakeDistribution(Epoch(15)))
320 );
321
322 let reference_hash = hash(SignedEntityType::CardanoStakeDistribution(Epoch(5)));
323 assert_ne!(
324 reference_hash,
325 hash(SignedEntityType::CardanoStakeDistribution(Epoch(15)))
326 );
327
328 let reference_hash = hash(SignedEntityType::CardanoImmutableFilesFull(
329 CardanoDbBeacon::new(5, 100),
330 ));
331 assert_ne!(
332 reference_hash,
333 hash(SignedEntityType::CardanoImmutableFilesFull(
334 CardanoDbBeacon::new(20, 100)
335 ))
336 );
337 assert_ne!(
338 reference_hash,
339 hash(SignedEntityType::CardanoImmutableFilesFull(
340 CardanoDbBeacon::new(5, 507)
341 ))
342 );
343
344 let reference_hash = hash(SignedEntityType::CardanoTransactions(
345 Epoch(35),
346 BlockNumber(77),
347 ));
348 assert_ne!(
349 reference_hash,
350 hash(SignedEntityType::CardanoTransactions(
351 Epoch(3),
352 BlockNumber(77)
353 ))
354 );
355 assert_ne!(
356 reference_hash,
357 hash(SignedEntityType::CardanoTransactions(
358 Epoch(35),
359 BlockNumber(98765)
360 ))
361 );
362
363 let reference_hash = hash(SignedEntityType::CardanoDatabase(CardanoDbBeacon::new(
364 12, 987,
365 )));
366 assert_ne!(
367 reference_hash,
368 hash(SignedEntityType::CardanoDatabase(CardanoDbBeacon::new(
369 98, 987
370 )))
371 );
372 assert_ne!(
373 reference_hash,
374 hash(SignedEntityType::CardanoDatabase(CardanoDbBeacon::new(
375 12, 123
376 )))
377 );
378 }
379
380 #[test]
381 fn serialize_beacon_to_json() {
382 let cardano_stake_distribution_json = SignedEntityType::CardanoStakeDistribution(Epoch(25))
383 .get_json_beacon()
384 .unwrap();
385 assert_same_json!("25", &cardano_stake_distribution_json);
386
387 let cardano_transactions_json =
388 SignedEntityType::CardanoTransactions(Epoch(35), BlockNumber(77))
389 .get_json_beacon()
390 .unwrap();
391 assert_same_json!(
392 r#"{"epoch":35,"block_number":77}"#,
393 &cardano_transactions_json
394 );
395
396 let cardano_immutable_files_full_json =
397 SignedEntityType::CardanoImmutableFilesFull(CardanoDbBeacon::new(5, 100))
398 .get_json_beacon()
399 .unwrap();
400 assert_same_json!(
401 r#"{"epoch":5,"immutable_file_number":100}"#,
402 &cardano_immutable_files_full_json
403 );
404
405 let msd_json = SignedEntityType::MithrilStakeDistribution(Epoch(15))
406 .get_json_beacon()
407 .unwrap();
408 assert_same_json!("15", &msd_json);
409
410 let cardano_database_full_json =
411 SignedEntityType::CardanoDatabase(CardanoDbBeacon::new(12, 987))
412 .get_json_beacon()
413 .unwrap();
414 assert_same_json!(
415 r#"{"epoch":12,"immutable_file_number":987}"#,
416 &cardano_database_full_json
417 );
418 }
419
420 #[test]
421 fn bytes_encoding() {
422 let cardano_stake_distribution = SignedEntityType::CardanoStakeDistribution(Epoch(25));
423 let cardano_stake_distribution_bytes = cardano_stake_distribution.to_bytes_vec().unwrap();
424 let cardano_stake_distribution_from_bytes =
425 SignedEntityType::try_from_bytes(&cardano_stake_distribution_bytes).unwrap();
426
427 assert_eq!(
428 cardano_stake_distribution,
429 cardano_stake_distribution_from_bytes
430 );
431 }
432
433 #[test]
436 fn ordering_discriminant() {
437 let mut list = vec![
438 SignedEntityTypeDiscriminants::CardanoStakeDistribution,
439 SignedEntityTypeDiscriminants::CardanoDatabase,
440 SignedEntityTypeDiscriminants::CardanoTransactions,
441 SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
442 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
443 ];
444 list.sort();
445
446 assert_eq!(
447 list,
448 vec![
449 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
450 SignedEntityTypeDiscriminants::CardanoStakeDistribution,
451 SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
452 SignedEntityTypeDiscriminants::CardanoDatabase,
453 SignedEntityTypeDiscriminants::CardanoTransactions,
454 ]
455 );
456 }
457
458 #[test]
459 fn ordering_discriminant_with_duplicate() {
460 let mut list = vec![
461 SignedEntityTypeDiscriminants::CardanoDatabase,
462 SignedEntityTypeDiscriminants::CardanoStakeDistribution,
463 SignedEntityTypeDiscriminants::CardanoDatabase,
464 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
465 SignedEntityTypeDiscriminants::CardanoTransactions,
466 SignedEntityTypeDiscriminants::CardanoStakeDistribution,
467 SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
468 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
469 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
470 ];
471 list.sort();
472
473 assert_eq!(
474 list,
475 vec![
476 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
477 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
478 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
479 SignedEntityTypeDiscriminants::CardanoStakeDistribution,
480 SignedEntityTypeDiscriminants::CardanoStakeDistribution,
481 SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
482 SignedEntityTypeDiscriminants::CardanoDatabase,
483 SignedEntityTypeDiscriminants::CardanoDatabase,
484 SignedEntityTypeDiscriminants::CardanoTransactions,
485 ]
486 );
487 }
488
489 #[test]
490 fn parse_signed_entity_types_discriminants_discriminant_without_values() {
491 let discriminants_str = "";
492 let discriminants = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap();
493
494 assert_eq!(BTreeSet::new(), discriminants);
495
496 let discriminants_str = " ";
497 let discriminants = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap();
498
499 assert_eq!(BTreeSet::new(), discriminants);
500 }
501
502 #[test]
503 fn parse_signed_entity_types_discriminants_with_correctly_formed_values() {
504 let discriminants_str = "MithrilStakeDistribution,CardanoImmutableFilesFull";
505 let discriminants = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap();
506
507 assert_eq!(
508 BTreeSet::from([
509 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
510 SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
511 ]),
512 discriminants
513 );
514 }
515
516 #[test]
517 fn parse_signed_entity_types_discriminants_should_trim_values() {
518 let discriminants_str =
519 "MithrilStakeDistribution , CardanoImmutableFilesFull , CardanoTransactions ";
520 let discriminants = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap();
521
522 assert_eq!(
523 BTreeSet::from([
524 SignedEntityTypeDiscriminants::MithrilStakeDistribution,
525 SignedEntityTypeDiscriminants::CardanoTransactions,
526 SignedEntityTypeDiscriminants::CardanoImmutableFilesFull,
527 ]),
528 discriminants
529 );
530 }
531
532 #[test]
533 fn parse_signed_entity_types_discriminants_should_remove_duplicates() {
534 let discriminants_str =
535 "CardanoTransactions,CardanoTransactions,CardanoTransactions,CardanoTransactions";
536 let discriminant = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap();
537
538 assert_eq!(
539 BTreeSet::from([SignedEntityTypeDiscriminants::CardanoTransactions]),
540 discriminant
541 );
542 }
543
544 #[test]
545 fn parse_signed_entity_types_discriminants_should_be_case_sensitive() {
546 let discriminants_str = "mithrilstakedistribution,CARDANOIMMUTABLEFILESFULL";
547 let error = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap_err();
548
549 assert_eq!(
550 SignedEntityTypeDiscriminants::format_parse_list_error(vec![
551 "mithrilstakedistribution",
552 "CARDANOIMMUTABLEFILESFULL"
553 ]),
554 error.to_string()
555 );
556 }
557
558 #[test]
559 fn parse_signed_entity_types_discriminants_should_not_return_unknown_signed_entity_types() {
560 let discriminants_str = "Unknown";
561 let error = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap_err();
562
563 assert_eq!(
564 SignedEntityTypeDiscriminants::format_parse_list_error(vec!["Unknown"]),
565 error.to_string()
566 );
567 }
568
569 #[test]
570 fn parse_signed_entity_types_discriminants_should_fail_if_there_is_at_least_one_invalid_value()
571 {
572 let discriminants_str = "CardanoTransactions,Invalid,MithrilStakeDistribution";
573 let error = SignedEntityTypeDiscriminants::parse_list(discriminants_str).unwrap_err();
574
575 assert_eq!(
576 SignedEntityTypeDiscriminants::format_parse_list_error(vec!["Invalid"]),
577 error.to_string()
578 );
579 }
580
581 #[test]
582 fn parse_list_error_format_to_an_useful_message() {
583 let invalid_discriminants = vec!["Unknown", "Invalid"];
584 let error = SignedEntityTypeDiscriminants::format_parse_list_error(invalid_discriminants);
585
586 assert_eq!(
587 format!(
588 r#"Invalid signed entity types discriminants: Unknown, Invalid.
589
590Accepted values are (case-sensitive): {}."#,
591 SignedEntityTypeDiscriminants::accepted_discriminants()
592 ),
593 error
594 );
595 }
596}