1use std::{
2 collections::BTreeMap,
3 fs,
4 path::{Path, PathBuf},
5 sync::Arc,
6};
7
8use anyhow::{anyhow, Context};
9
10use mithril_common::{
11 crypto_helper::{MKProof, MKTree, MKTreeNode, MKTreeStoreInMemory},
12 digesters::{CardanoImmutableDigester, ImmutableDigester, ImmutableFile},
13 entities::{DigestLocation, HexEncodedDigest, ImmutableFileName},
14 messages::{
15 CardanoDatabaseDigestListItemMessage, CardanoDatabaseSnapshotMessage, CertificateMessage,
16 DigestsMessagePart,
17 },
18};
19
20use crate::{
21 feedback::MithrilEvent,
22 file_downloader::{DownloadEvent, FileDownloader, FileDownloaderUri},
23 utils::{create_directory_if_not_exists, delete_directory, read_files_in_directory},
24 MithrilResult,
25};
26
27use super::immutable_file_range::ImmutableFileRange;
28
29pub struct InternalArtifactProver {
30 http_file_downloader: Arc<dyn FileDownloader>,
31 logger: slog::Logger,
32}
33
34impl InternalArtifactProver {
35 pub fn new(http_file_downloader: Arc<dyn FileDownloader>, logger: slog::Logger) -> Self {
37 Self {
38 http_file_downloader,
39 logger,
40 }
41 }
42
43 pub async fn compute_merkle_proof(
45 &self,
46 certificate: &CertificateMessage,
47 cardano_database_snapshot: &CardanoDatabaseSnapshotMessage,
48 immutable_file_range: &ImmutableFileRange,
49 database_dir: &Path,
50 ) -> MithrilResult<MKProof> {
51 self.download_unpack_digest_file(
52 &cardano_database_snapshot.digests,
53 &Self::digest_target_dir(database_dir),
54 )
55 .await?;
56 let network = certificate.metadata.network.clone();
57 let last_immutable_file_number = cardano_database_snapshot.beacon.immutable_file_number;
58 let immutable_file_number_range =
59 immutable_file_range.to_range_inclusive(last_immutable_file_number)?;
60 let downloaded_digests = self.read_digest_file(&Self::digest_target_dir(database_dir))?;
61 let downloaded_digests_values = downloaded_digests
62 .into_iter()
63 .filter(|(immutable_file_name, _)| {
64 match ImmutableFile::new(Path::new(immutable_file_name).to_path_buf()) {
65 Ok(immutable_file) => immutable_file.number <= last_immutable_file_number,
66 Err(_) => false,
67 }
68 })
69 .map(|(_immutable_file_name, digest)| digest)
70 .collect::<Vec<_>>();
71 let merkle_tree: MKTree<MKTreeStoreInMemory> = MKTree::new(&downloaded_digests_values)?;
72 let immutable_digester = CardanoImmutableDigester::new(network, None, self.logger.clone());
73 let computed_digests = immutable_digester
74 .compute_digests_for_range(database_dir, &immutable_file_number_range)
75 .await?
76 .entries
77 .values()
78 .map(MKTreeNode::from)
79 .collect::<Vec<_>>();
80 delete_directory(&Self::digest_target_dir(database_dir))?;
81
82 merkle_tree.compute_proof(&computed_digests)
83 }
84
85 async fn download_unpack_digest_file(
86 &self,
87 digests_locations: &DigestsMessagePart,
88 digests_file_target_dir: &Path,
89 ) -> MithrilResult<()> {
90 create_directory_if_not_exists(digests_file_target_dir)?;
91 let mut locations_sorted = digests_locations.sanitized_locations()?;
92 locations_sorted.sort();
93 for location in locations_sorted {
94 let download_id = MithrilEvent::new_cardano_database_download_id();
95 let (file_downloader, compression_algorithm) = match &location {
96 DigestLocation::CloudStorage {
97 uri: _,
98 compression_algorithm,
99 } => (self.http_file_downloader.clone(), *compression_algorithm),
100 DigestLocation::Aggregator { .. } => (self.http_file_downloader.clone(), None),
101 DigestLocation::Unknown => unreachable!(),
103 };
104 let file_downloader_uri: FileDownloaderUri = location.try_into()?;
105 let downloaded = file_downloader
106 .download_unpack(
107 &file_downloader_uri,
108 digests_locations.size_uncompressed,
109 digests_file_target_dir,
110 compression_algorithm,
111 DownloadEvent::Digest {
112 download_id: download_id.clone(),
113 },
114 )
115 .await;
116 match downloaded {
117 Ok(_) => {
118 return Ok(());
119 }
120 Err(e) => {
121 slog::error!(
122 self.logger,
123 "Failed downloading and unpacking digest for location {file_downloader_uri:?}"; "error" => ?e
124 );
125 }
126 }
127 }
128
129 Err(anyhow!(
130 "Failed downloading and unpacking digests for all locations"
131 ))
132 }
133
134 fn read_digest_file(
135 &self,
136 digest_file_target_dir: &Path,
137 ) -> MithrilResult<BTreeMap<ImmutableFileName, HexEncodedDigest>> {
138 let digest_files = read_files_in_directory(digest_file_target_dir)?;
139 if digest_files.len() > 1 {
140 return Err(anyhow!(
141 "Multiple digest files found in directory: {digest_file_target_dir:?}"
142 ));
143 }
144 if digest_files.is_empty() {
145 return Err(anyhow!(
146 "No digest file found in directory: {digest_file_target_dir:?}"
147 ));
148 }
149
150 let digest_file = &digest_files[0];
151 let content = fs::read_to_string(digest_file)
152 .with_context(|| format!("Failed reading digest file: {digest_file:?}"))?;
153 let digest_messages: Vec<CardanoDatabaseDigestListItemMessage> =
154 serde_json::from_str(&content)
155 .with_context(|| format!("Failed deserializing digest file: {digest_file:?}"))?;
156 let digest_map = digest_messages
157 .into_iter()
158 .map(|message| (message.immutable_file_name, message.digest))
159 .collect::<BTreeMap<_, _>>();
160
161 Ok(digest_map)
162 }
163
164 fn digest_target_dir(target_dir: &Path) -> PathBuf {
165 target_dir.join("digest")
166 }
167}
168
169#[cfg(test)]
170mod tests {
171 use std::collections::BTreeMap;
172 use std::fs;
173 use std::io::Write;
174 use std::path::Path;
175 use std::sync::Arc;
176
177 use mithril_common::{
178 digesters::{DummyCardanoDbBuilder, ImmutableDigester, ImmutableFile},
179 entities::{CardanoDbBeacon, Epoch, HexEncodedDigest},
180 messages::CardanoDatabaseDigestListItemMessage,
181 test_utils::TempDir,
182 };
183
184 use crate::{
185 cardano_database_client::CardanoDatabaseClientDependencyInjector,
186 file_downloader::MockFileDownloaderBuilder, test_utils,
187 };
188
189 use super::*;
190
191 mod compute_merkle_proof {
192
193 use std::ops::RangeInclusive;
194
195 use mithril_common::{entities::ImmutableFileNumber, messages::DigestsMessagePart};
196
197 use super::*;
198
199 async fn create_fake_digest_artifact(
200 dir_name: &str,
201 beacon: &CardanoDbBeacon,
202 immutable_file_range: &RangeInclusive<ImmutableFileNumber>,
203 digests_offset: usize,
204 ) -> (
205 PathBuf,
206 CardanoDatabaseSnapshotMessage,
207 CertificateMessage,
208 MKTree<MKTreeStoreInMemory>,
209 ) {
210 let cardano_database_snapshot = CardanoDatabaseSnapshotMessage {
211 hash: "hash-123".to_string(),
212 beacon: beacon.clone(),
213 digests: DigestsMessagePart {
214 size_uncompressed: 1024,
215 locations: vec![DigestLocation::CloudStorage {
216 uri: "http://whatever/digests.json".to_string(),
217 compression_algorithm: None,
218 }],
219 },
220 ..CardanoDatabaseSnapshotMessage::dummy()
221 };
222 let certificate = CertificateMessage {
223 hash: "cert-hash-123".to_string(),
224 ..CertificateMessage::dummy()
225 };
226 let cardano_db = DummyCardanoDbBuilder::new(dir_name)
227 .with_immutables(&immutable_file_range.clone().collect::<Vec<_>>())
228 .append_immutable_trio()
229 .build();
230 let database_dir = cardano_db.get_dir();
231 let immutable_digester = CardanoImmutableDigester::new(
232 certificate.metadata.network.to_string(),
233 None,
234 test_utils::test_logger(),
235 );
236 let computed_digests = immutable_digester
237 .compute_digests_for_range(database_dir, immutable_file_range)
238 .await
239 .unwrap();
240 write_digest_file(&database_dir.join("digest"), &computed_digests.entries).await;
241
242 for (immutable_file, _digest) in
245 computed_digests.entries.iter().rev().take(digests_offset)
246 {
247 fs::remove_file(
248 database_dir.join(
249 database_dir
250 .join("immutable")
251 .join(immutable_file.filename.clone()),
252 ),
253 )
254 .unwrap();
255 }
256
257 let merkle_tree = immutable_digester
258 .compute_merkle_tree(database_dir, beacon)
259 .await
260 .unwrap();
261
262 (
263 database_dir.to_owned(),
264 cardano_database_snapshot,
265 certificate,
266 merkle_tree,
267 )
268 }
269
270 async fn write_digest_file(
271 digest_dir: &Path,
272 digests: &BTreeMap<ImmutableFile, HexEncodedDigest>,
273 ) {
274 let digest_file_path = digest_dir.join("digests.json");
275 if !digest_dir.exists() {
276 fs::create_dir_all(digest_dir).unwrap();
277 }
278
279 let immutable_digest_messages = digests
280 .iter()
281 .map(
282 |(immutable_file, digest)| CardanoDatabaseDigestListItemMessage {
283 immutable_file_name: immutable_file.filename.clone(),
284 digest: digest.to_string(),
285 },
286 )
287 .collect::<Vec<_>>();
288 serde_json::to_writer(
289 fs::File::create(digest_file_path).unwrap(),
290 &immutable_digest_messages,
291 )
292 .unwrap();
293 }
294
295 #[tokio::test]
296 async fn compute_merkle_proof_succeeds() {
297 let beacon = CardanoDbBeacon {
298 epoch: Epoch(123),
299 immutable_file_number: 10,
300 };
301 let immutable_file_range = 1..=15;
302 let immutable_file_range_to_prove = ImmutableFileRange::Range(2, 4);
303 let digests_offset = 3;
304 let (database_dir, cardano_database_snapshot, certificate, merkle_tree) =
305 create_fake_digest_artifact(
306 "compute_merkle_proof_succeeds",
307 &beacon,
308 &immutable_file_range,
309 digests_offset,
310 )
311 .await;
312 let expected_merkle_root = merkle_tree.compute_root().unwrap();
313 let client = CardanoDatabaseClientDependencyInjector::new()
314 .with_http_file_downloader(Arc::new(
315 MockFileDownloaderBuilder::default()
316 .with_file_uri("http://whatever/digests.json")
317 .with_target_dir(database_dir.join("digest"))
318 .with_compression(None)
319 .with_success()
320 .build(),
321 ))
322 .build_cardano_database_client();
323
324 let merkle_proof = client
325 .compute_merkle_proof(
326 &certificate,
327 &cardano_database_snapshot,
328 &immutable_file_range_to_prove,
329 &database_dir,
330 )
331 .await
332 .unwrap();
333 merkle_proof.verify().unwrap();
334
335 let merkle_proof_root = merkle_proof.root().to_owned();
336 assert_eq!(expected_merkle_root, merkle_proof_root);
337
338 assert!(!database_dir.join("digest").exists());
339 }
340 }
341
342 mod download_unpack_digest_file {
343
344 use mithril_common::entities::CompressionAlgorithm;
345
346 use crate::file_downloader::MockFileDownloader;
347
348 use super::*;
349
350 #[tokio::test]
351 async fn fails_if_no_location_is_retrieved() {
352 let target_dir = Path::new(".");
353 let artifact_prover = InternalArtifactProver::new(
354 Arc::new(
355 MockFileDownloaderBuilder::default()
356 .with_compression(None)
357 .with_failure()
358 .with_times(2)
359 .build(),
360 ),
361 test_utils::test_logger(),
362 );
363
364 artifact_prover
365 .download_unpack_digest_file(
366 &DigestsMessagePart {
367 locations: vec![
368 DigestLocation::CloudStorage {
369 uri: "http://whatever-1/digests.json".to_string(),
370 compression_algorithm: None,
371 },
372 DigestLocation::Aggregator {
373 uri: "http://whatever-2/digest".to_string(),
374 },
375 ],
376 size_uncompressed: 0,
377 },
378 target_dir,
379 )
380 .await
381 .expect_err("download_unpack_digest_file should fail");
382 }
383
384 #[tokio::test]
385 async fn fails_if_all_locations_are_unknown() {
386 let target_dir = Path::new(".");
387 let artifact_prover = InternalArtifactProver::new(
388 Arc::new(MockFileDownloader::new()),
389 test_utils::test_logger(),
390 );
391
392 artifact_prover
393 .download_unpack_digest_file(
394 &DigestsMessagePart {
395 locations: vec![DigestLocation::Unknown],
396 size_uncompressed: 0,
397 },
398 target_dir,
399 )
400 .await
401 .expect_err("download_unpack_digest_file should fail");
402 }
403
404 #[tokio::test]
405 async fn succeeds_if_at_least_one_location_is_retrieved() {
406 let target_dir = Path::new(".");
407 let artifact_prover = InternalArtifactProver::new(
408 Arc::new(
409 MockFileDownloaderBuilder::default()
410 .with_compression(None)
411 .with_failure()
412 .next_call()
413 .with_compression(None)
414 .with_success()
415 .build(),
416 ),
417 test_utils::test_logger(),
418 );
419
420 artifact_prover
421 .download_unpack_digest_file(
422 &DigestsMessagePart {
423 locations: vec![
424 DigestLocation::CloudStorage {
425 uri: "http://whatever-1/digests.json".to_string(),
426 compression_algorithm: None,
427 },
428 DigestLocation::Aggregator {
429 uri: "http://whatever-2/digest".to_string(),
430 },
431 ],
432 size_uncompressed: 0,
433 },
434 target_dir,
435 )
436 .await
437 .unwrap();
438 }
439
440 #[tokio::test]
441 async fn succeeds_when_first_location_is_retrieved() {
442 let target_dir = Path::new(".");
443 let artifact_prover = InternalArtifactProver::new(
444 Arc::new(
445 MockFileDownloaderBuilder::default()
446 .with_compression(None)
447 .with_times(1)
448 .with_success()
449 .build(),
450 ),
451 test_utils::test_logger(),
452 );
453
454 artifact_prover
455 .download_unpack_digest_file(
456 &DigestsMessagePart {
457 locations: vec![
458 DigestLocation::CloudStorage {
459 uri: "http://whatever-1/digests.json".to_string(),
460 compression_algorithm: None,
461 },
462 DigestLocation::Aggregator {
463 uri: "http://whatever-2/digest".to_string(),
464 },
465 ],
466 size_uncompressed: 0,
467 },
468 target_dir,
469 )
470 .await
471 .unwrap();
472 }
473
474 #[tokio::test]
475 async fn should_call_download_with_compression_algorithm() {
476 let target_dir = Path::new(".");
477 let artifact_prover = InternalArtifactProver::new(
478 Arc::new(
479 MockFileDownloaderBuilder::default()
480 .with_compression(Some(CompressionAlgorithm::Gzip))
481 .with_times(1)
482 .with_success()
483 .build(),
484 ),
485 test_utils::test_logger(),
486 );
487
488 artifact_prover
489 .download_unpack_digest_file(
490 &DigestsMessagePart {
491 locations: vec![
492 DigestLocation::CloudStorage {
493 uri: "http://whatever-1/digests.tar.gz".to_string(),
494 compression_algorithm: Some(CompressionAlgorithm::Gzip),
495 },
496 DigestLocation::Aggregator {
497 uri: "http://whatever-2/digest".to_string(),
498 },
499 ],
500 size_uncompressed: 0,
501 },
502 target_dir,
503 )
504 .await
505 .unwrap();
506 }
507 }
508
509 mod read_digest_file {
510
511 use super::*;
512
513 fn create_valid_fake_digest_file(
514 file_path: &Path,
515 digest_messages: &[CardanoDatabaseDigestListItemMessage],
516 ) {
517 let mut file = fs::File::create(file_path).unwrap();
518 let digest_json = serde_json::to_string(&digest_messages).unwrap();
519 file.write_all(digest_json.as_bytes()).unwrap();
520 }
521
522 fn create_invalid_fake_digest_file(file_path: &Path) {
523 let mut file = fs::File::create(file_path).unwrap();
524 file.write_all(b"incorrect-digest").unwrap();
525 }
526
527 #[test]
528 fn read_digest_file_fails_when_no_digest_file() {
529 let target_dir = TempDir::new(
530 "cardano_database_client",
531 "read_digest_file_fails_when_no_digest_file",
532 )
533 .build();
534 let artifact_prover = InternalArtifactProver::new(
535 Arc::new(
536 MockFileDownloaderBuilder::default()
537 .with_times(0)
538 .with_success()
539 .build(),
540 ),
541 test_utils::test_logger(),
542 );
543 artifact_prover
544 .read_digest_file(&target_dir)
545 .expect_err("read_digest_file should fail");
546 }
547
548 #[test]
549 fn read_digest_file_fails_when_multiple_digest_files() {
550 let target_dir = TempDir::new(
551 "cardano_database_client",
552 "read_digest_file_fails_when_multiple_digest_files",
553 )
554 .build();
555 create_valid_fake_digest_file(&target_dir.join("digests.json"), &[]);
556 create_valid_fake_digest_file(&target_dir.join("digests-2.json"), &[]);
557 let artifact_prover = InternalArtifactProver::new(
558 Arc::new(
559 MockFileDownloaderBuilder::default()
560 .with_times(0)
561 .with_success()
562 .build(),
563 ),
564 test_utils::test_logger(),
565 );
566 artifact_prover
567 .read_digest_file(&target_dir)
568 .expect_err("read_digest_file should fail");
569 }
570
571 #[test]
572 fn read_digest_file_fails_when_invalid_unique_digest_file() {
573 let target_dir = TempDir::new(
574 "cardano_database_client",
575 "read_digest_file_fails_when_invalid_unique_digest_file",
576 )
577 .build();
578 create_invalid_fake_digest_file(&target_dir.join("digests.json"));
579 let artifact_prover = InternalArtifactProver::new(
580 Arc::new(
581 MockFileDownloaderBuilder::default()
582 .with_times(0)
583 .with_success()
584 .build(),
585 ),
586 test_utils::test_logger(),
587 );
588 artifact_prover
589 .read_digest_file(&target_dir)
590 .expect_err("read_digest_file should fail");
591 }
592
593 #[test]
594 fn read_digest_file_succeeds_when_valid_unique_digest_file() {
595 let target_dir = TempDir::new(
596 "cardano_database_client",
597 "read_digest_file_succeeds_when_valid_unique_digest_file",
598 )
599 .build();
600 let digest_messages = vec![
601 CardanoDatabaseDigestListItemMessage {
602 immutable_file_name: "00001.chunk".to_string(),
603 digest: "digest-1".to_string(),
604 },
605 CardanoDatabaseDigestListItemMessage {
606 immutable_file_name: "00002.chunk".to_string(),
607 digest: "digest-2".to_string(),
608 },
609 ];
610 create_valid_fake_digest_file(&target_dir.join("digests.json"), &digest_messages);
611 let artifact_prover = InternalArtifactProver::new(
612 Arc::new(
613 MockFileDownloaderBuilder::default()
614 .with_times(0)
615 .with_success()
616 .build(),
617 ),
618 test_utils::test_logger(),
619 );
620
621 let digests = artifact_prover.read_digest_file(&target_dir).unwrap();
622 assert_eq!(
623 BTreeMap::from([
624 ("00001.chunk".to_string(), "digest-1".to_string()),
625 ("00002.chunk".to_string(), "digest-2".to_string())
626 ]),
627 digests
628 )
629 }
630 }
631}