mithril_common/digesters/
dumb_immutable_digester.rs

1use std::{collections::BTreeMap, ops::RangeInclusive, path::Path};
2
3use crate::{
4    crypto_helper::{MKTree, MKTreeStoreInMemory},
5    digesters::{ImmutableDigester, ImmutableDigesterError},
6    entities::{CardanoDbBeacon, ImmutableFileNumber},
7};
8use async_trait::async_trait;
9use tokio::sync::RwLock;
10
11use super::{immutable_digester::ComputedImmutablesDigests, ImmutableFile};
12
13/// A [ImmutableDigester] returning configurable result for testing purpose.
14pub struct DumbImmutableDigester {
15    digest: RwLock<String>,
16    mktree_leaves: RwLock<Vec<String>>,
17    is_success: bool,
18}
19
20impl DumbImmutableDigester {
21    /// Set the digest returned by [compute_digest][DumbImmutableDigester::compute_digest]
22    pub fn with_digest(mut self, new_digest: &str) -> Self {
23        self.digest = RwLock::new(new_digest.to_string());
24        self
25    }
26
27    /// Set the leaves used to construct the merkle tree returned by [compute_merkle_tree][DumbImmutableDigester::compute_merkle_tree]
28    pub fn with_merkle_tree(mut self, leaves: Vec<String>) -> Self {
29        self.mktree_leaves = RwLock::new(leaves);
30        self
31    }
32
33    /// Update digest returned by [compute_digest][DumbImmutableDigester::compute_digest]
34    pub async fn update_digest(&self, new_digest: String) {
35        let mut digest = self.digest.write().await;
36        *digest = new_digest;
37    }
38
39    /// Update the leaves used to construct the merkle tree returned by [compute_merkle_tree][DumbImmutableDigester::compute_merkle_tree]
40    pub async fn update_merkle_tree(&self, leaves: Vec<String>) {
41        let mut mktree_leaves = self.mktree_leaves.write().await;
42        *mktree_leaves = leaves;
43    }
44}
45
46impl Default for DumbImmutableDigester {
47    fn default() -> Self {
48        Self {
49            digest: RwLock::new(String::from("1234")),
50            mktree_leaves: RwLock::new(vec!["1".to_string(), "2".to_string(), "3".to_string()]),
51            is_success: true,
52        }
53    }
54}
55
56#[async_trait]
57impl ImmutableDigester for DumbImmutableDigester {
58    async fn compute_digest(
59        &self,
60        dirpath: &Path,
61        beacon: &CardanoDbBeacon,
62    ) -> Result<String, ImmutableDigesterError> {
63        if self.is_success {
64            Ok(self.digest.read().await.clone())
65        } else {
66            Err(ImmutableDigesterError::NotEnoughImmutable {
67                expected_number: beacon.immutable_file_number,
68                found_number: None,
69                db_dir: dirpath.to_owned(),
70            })
71        }
72    }
73
74    async fn compute_digests_for_range(
75        &self,
76        dirpath: &Path,
77        range: &RangeInclusive<ImmutableFileNumber>,
78    ) -> Result<ComputedImmutablesDigests, ImmutableDigesterError> {
79        if self.is_success {
80            let immutable_file_paths = range
81                .clone()
82                .flat_map(|immutable_file_number| {
83                    vec![
84                        Path::new(&format!("{immutable_file_number:0>5}.chunk")).to_path_buf(),
85                        Path::new(&format!("{immutable_file_number:0>5}.primary")).to_path_buf(),
86                        Path::new(&format!("{immutable_file_number:0>5}.secondary")).to_path_buf(),
87                    ]
88                })
89                .collect::<Vec<_>>();
90            let digest = self.digest.read().await.clone();
91
92            Ok(ComputedImmutablesDigests::compute_immutables_digests(
93                BTreeMap::from_iter(immutable_file_paths.into_iter().map(|immutable_file_path| {
94                    (
95                        ImmutableFile::new(immutable_file_path).unwrap(),
96                        Some(digest.clone()),
97                    )
98                })),
99                slog::Logger::root(slog::Discard, slog::o!()),
100            )?)
101        } else {
102            Err(ImmutableDigesterError::NotEnoughImmutable {
103                expected_number: *range.end(),
104                found_number: None,
105                db_dir: dirpath.to_owned(),
106            })
107        }
108    }
109
110    async fn compute_merkle_tree(
111        &self,
112        dirpath: &Path,
113        beacon: &CardanoDbBeacon,
114    ) -> Result<MKTree<MKTreeStoreInMemory>, ImmutableDigesterError> {
115        if self.is_success {
116            let leaves = self.mktree_leaves.read().await;
117            Ok(MKTree::new(&leaves).unwrap())
118        } else {
119            Err(ImmutableDigesterError::NotEnoughImmutable {
120                expected_number: beacon.immutable_file_number,
121                found_number: None,
122                db_dir: dirpath.to_owned(),
123            })
124        }
125    }
126}