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