feat: cleanup unexpected files in immutable folder after download

mithril-client/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-client"
-version = "0.12.5"
+version = "0.12.6"
 description = "Mithril client library"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-client/src/cardano_database_client/download_unpack/internal_downloader.rs

@@ -15,8 +15,8 @@ use crate::cardano_database_client::ImmutableFileRange;
 use crate::feedback::{FeedbackSender, MithrilEvent, MithrilEventCardanoDatabase};
 use crate::file_downloader::{DownloadEvent, FileDownloader, FileDownloaderUri};
 use crate::utils::{
-    create_bootstrap_node_files, AncillaryVerifier, VecDequeExtensions,
-    ANCILLARIES_NOT_SIGNED_BY_MITHRIL,
+    create_bootstrap_node_files, AncillaryVerifier, UnexpectedDownloadedFileVerifier,
+    VecDequeExtensions, ANCILLARIES_NOT_SIGNED_BY_MITHRIL,
 };
 use crate::MithrilResult;
 
@@ -73,6 +73,16 @@ impl InternalArtifactDownloader {
             .verify_compatibility(&immutable_file_number_range, last_immutable_file_number)?;
         download_unpack_options.verify_can_write_to_target_directory(target_dir)?;
 
+        let expected_files_after_download = UnexpectedDownloadedFileVerifier::new(
+            target_dir,
+            &cardano_database_snapshot.network,
+            download_unpack_options.include_ancillary,
+            last_immutable_file_number,
+            &self.logger,
+        )
+        .compute_expected_state_after_download()
+        .await?;
+
         let mut tasks = VecDeque::from(self.build_download_tasks_for_immutables(
             &cardano_database_snapshot.immutables,
             immutable_file_number_range,
@@ -89,8 +99,15 @@ impl InternalArtifactDownloader {
         } else {
             slog::warn!(self.logger, "The fast bootstrap of the Cardano node is not available with the current parameters used in this command: the ledger state will be recomputed from genesis at startup of the Cardano node. Set the include_ancillary entry to true in the DownloadUnpackOptions.");
         }
-        self.batch_download_unpack(tasks, download_unpack_options.max_parallel_downloads)
+
+        // Return the result later so unexpected file removal is always run
+        let download_result = self
+            .batch_download_unpack(tasks, download_unpack_options.max_parallel_downloads)
+            .await;
+        expected_files_after_download
+            .remove_unexpected_files()
             .await?;
+        download_result?;
 
         create_bootstrap_node_files(&self.logger, target_dir, &cardano_database_snapshot.network)?;
 
@@ -268,11 +285,13 @@ mod tests {
     use super::*;
 
     mod download_unpack {
+        use mithril_common::assert_dir_eq;
         use mithril_common::crypto_helper::ManifestSigner;
+        use mithril_common::digesters::IMMUTABLE_DIR;
         use mithril_common::entities::CompressionAlgorithm;
         use mithril_common::messages::DigestsMessagePart;
 
-        use crate::file_downloader::FakeAncillaryFileBuilder;
+        use crate::file_downloader::{FakeAncillaryFileBuilder, MockFileDownloader};
 
         use super::*;
 
@@ -349,7 +368,7 @@ mod tests {
                 ..CardanoDatabaseSnapshot::dummy()
             };
             let target_dir = temp_dir_create!();
-            fs::create_dir_all(target_dir.join("immutable")).unwrap();
+            fs::create_dir_all(target_dir.join(IMMUTABLE_DIR)).unwrap();
             let client =
                 CardanoDatabaseClientDependencyInjector::new().build_cardano_database_client();
 
@@ -436,6 +455,50 @@ mod tests {
                 .unwrap();
         }
 
+        #[tokio::test]
+        async fn download_unpack_remove_unexpected_downloaded_files() {
+            let target_dir = temp_dir_create!();
+            let immutable_dir = target_dir.join(IMMUTABLE_DIR);
+
+            let immutable_file_range = ImmutableFileRange::UpTo(1);
+            let download_unpack_options = DownloadUnpackOptions {
+                include_ancillary: false,
+                ..DownloadUnpackOptions::default()
+            };
+            let cardano_db_snapshot = CardanoDatabaseSnapshot::dummy();
+            let client = CardanoDatabaseClientDependencyInjector::new()
+                .with_http_file_downloader(Arc::new({
+                    let mut mock_downloader = MockFileDownloader::new();
+                    mock_downloader
+                        .expect_download_unpack()
+                        .returning(move |_, _, _, _, _| {
+                            // Simulate an additional file written mid-download
+                            if !immutable_dir.exists() {
+                                fs::create_dir(&immutable_dir).unwrap();
+                            }
+                            fs::File::create(immutable_dir.join("unexpected.md")).unwrap();
+                            Ok(())
+                        });
+                    mock_downloader
+                }))
+                .build_cardano_database_client();
+
+            client
+                .download_unpack(
+                    &cardano_db_snapshot,
+                    &immutable_file_range,
+                    &target_dir,
+                    download_unpack_options,
+                )
+                .await
+                .unwrap();
+
+            assert_dir_eq!(
+                &target_dir,
+                format!("* {IMMUTABLE_DIR}/\n* clean\n * protocolMagicId")
+            );
+        }
+
         #[tokio::test]
         async fn fail_if_include_ancillary_is_true_and_ancillary_verifier_is_not_set() {
             let download_unpack_options = DownloadUnpackOptions {

mithril-client/src/snapshot_client.rs

@@ -120,8 +120,9 @@ use crate::file_downloader::{DownloadEvent, FileDownloader};
 #[cfg(feature = "fs")]
 use crate::utils::create_bootstrap_node_files;
 #[cfg(feature = "fs")]
-use crate::utils::AncillaryVerifier;
-use crate::utils::ANCILLARIES_NOT_SIGNED_BY_MITHRIL;
+use crate::utils::{
+    AncillaryVerifier, UnexpectedDownloadedFileVerifier, ANCILLARIES_NOT_SIGNED_BY_MITHRIL,
+};
 use crate::{MithrilResult, Snapshot, SnapshotListItem};
 
 /// Error for the Snapshot client
@@ -225,23 +226,29 @@ impl SnapshotClient {
             snapshot: &Snapshot,
             target_dir: &Path,
         ) -> MithrilResult<()> {
-            use crate::feedback::MithrilEvent;
             if self.ancillary_verifier.is_none() {
                 return Err(SnapshotClientError::MissingAncillaryVerifier.into());
             }
 
-            let download_id = MithrilEvent::new_snapshot_download_id();
-            self.download_unpack_immutables_files(snapshot, target_dir, &download_id)
-                .await?;
-            self.download_unpack_ancillary(snapshot, target_dir, &download_id)
-                .await?;
-            create_bootstrap_node_files(
-                &self.logger,
+            let include_ancillary = true;
+            let expected_files_after_download = UnexpectedDownloadedFileVerifier::new(
                 target_dir,
                 &snapshot.network,
-            )?;
+                include_ancillary,
+                snapshot.beacon.immutable_file_number,
+                &self.logger
+            )
+            .compute_expected_state_after_download()
+            .await?;
 
-            Ok(())
+            // Return the result later so unexpected file removal is always run
+            let result = self.run_download_unpack(snapshot, target_dir, include_ancillary).await;
+
+            expected_files_after_download
+                .remove_unexpected_files()
+                .await?;
+
+            result
         }
 
         /// Download and unpack the given immutable files of the snapshot to the given directory
@@ -260,16 +267,48 @@ impl SnapshotClient {
                 self.logger,
                 "The fast bootstrap of the Cardano node is not available with the current parameters used in this command: the ledger state will be recomputed from genesis at startup of the Cardano node. Use the extra function download_unpack_full to allow it."
             );
+
+            let include_ancillary = false;
+            let expected_files_after_download = UnexpectedDownloadedFileVerifier::new(
+                target_dir,
+                &snapshot.network,
+                include_ancillary,
+                snapshot.beacon.immutable_file_number,
+                &self.logger
+            )
+            .compute_expected_state_after_download()
+            .await?;
+
+            // Return the result later so unexpected file removal is always run
+            let result = self.run_download_unpack(snapshot, target_dir, include_ancillary).await;
+
+            expected_files_after_download
+                .remove_unexpected_files()
+                .await?;
+
+            result
+        }
+
+        async fn run_download_unpack(
+            &self,
+            snapshot: &Snapshot,
+            target_dir: &Path,
+            include_ancillary: bool,
+        ) -> MithrilResult<()> {
             use crate::feedback::MithrilEvent;
+
             let download_id = MithrilEvent::new_snapshot_download_id();
             self.download_unpack_immutables_files(snapshot, target_dir, &download_id)
                 .await?;
+            if include_ancillary {
+                self.download_unpack_ancillary(snapshot, target_dir, &download_id)
+                    .await?;
+            }
             create_bootstrap_node_files(
                 &self.logger,
                 target_dir,
                 &snapshot.network,
             )?;
-
             Ok(())
         }
 
@@ -439,9 +478,12 @@ mod tests {
         test_utils::TestLogger,
     };
 
-    use super::*;
+    use mithril_common::{
+        assert_dir_eq, crypto_helper::ManifestSigner, digesters::IMMUTABLE_DIR, temp_dir_create,
+        test_utils::fake_keys,
+    };
 
-    use mithril_common::{assert_dir_eq, temp_dir_create};
+    use super::*;
 
     fn dummy_download_event() -> DownloadEvent {
         DownloadEvent::Full {
@@ -622,6 +664,38 @@ mod tests {
                 "Expected SnapshotClientError::MissingAncillaryVerifier, but got: {error:#?}"
             );
         }
+
+        #[tokio::test]
+        async fn remove_unexpected_downloaded_files_even_if_failing_to_verify_ancillary() {
+            let test_dir = temp_dir_create!();
+            let immutable_dir = test_dir.join(IMMUTABLE_DIR);
+            std::fs::create_dir(&immutable_dir).unwrap();
+            let snapshot = Snapshot::dummy();
+
+            let mut mock_downloader = MockFileDownloader::new();
+            mock_downloader
+                .expect_download_unpack()
+                .returning(move |_, _, _, _, _| {
+                    // Simulate an additional file written mid-download
+                    std::fs::File::create(immutable_dir.join("unexpected.md")).unwrap();
+                    Ok(())
+                });
+
+            let client = setup_snapshot_client(
+                Arc::new(mock_downloader),
+                Some(Arc::new(AncillaryVerifier::new(
+                    fake_keys::manifest_verification_key()[0]
+                        .try_into()
+                        .unwrap(),
+                ))),
+            );
+
+            client
+                .download_unpack_full(&snapshot, &test_dir)
+                .await
+                .unwrap_err();
+            assert_dir_eq!(&test_dir, format!("* {IMMUTABLE_DIR}/"));
+        }
     }
 
     mod download_unpack {
@@ -651,12 +725,34 @@ mod tests {
                 "Expected log message not found, logs: {log_inspector}"
             );
         }
+
+        #[tokio::test]
+        async fn remove_unexpected_downloaded_files() {
+            let test_dir = temp_dir_create!();
+            let immutable_dir = test_dir.join(IMMUTABLE_DIR);
+            std::fs::create_dir(&immutable_dir).unwrap();
+            let snapshot = Snapshot::dummy();
+
+            let mut mock_downloader = MockFileDownloader::new();
+            mock_downloader
+                .expect_download_unpack()
+                .returning(move |_, _, _, _, _| {
+                    // Simulate an additional file written mid-download
+                    std::fs::File::create(immutable_dir.join("unexpected.md")).unwrap();
+                    Ok(())
+                });
+
+            let client = setup_snapshot_client(Arc::new(mock_downloader), None);
+
+            client.download_unpack(&snapshot, &test_dir).await.unwrap();
+            assert_dir_eq!(
+                &test_dir,
+                format!("* {IMMUTABLE_DIR}/\n* clean\n * protocolMagicId")
+            );
+        }
     }
 
     mod download_unpack_ancillary {
-        use mithril_common::crypto_helper::ManifestSigner;
-        use mithril_common::test_utils::fake_keys;
-
         use crate::file_downloader::FakeAncillaryFileBuilder;
 
         use super::*;

mithril-client/src/utils/mod.rs

@@ -1,14 +1,16 @@
 //! Utilities module
 //! This module contains tools needed mostly for the snapshot download and unpack.
 
-pub const ANCILLARIES_NOT_SIGNED_BY_MITHRIL:&str = "Ancillary verification does not use the Mithril certification: as a mitigation, IOG owned keys are used to sign these files.";
-
 cfg_fs! {
+    pub const ANCILLARIES_NOT_SIGNED_BY_MITHRIL:&str = "Ancillary verification does not use the Mithril certification: as a mitigation, IOG owned keys are used to sign these files.";
+
     mod ancillary_verifier;
     mod stream_reader;
     mod bootstrap_files;
+    mod unexpected_downloaded_file_verifier;
 
     pub use ancillary_verifier::AncillaryVerifier;
+    pub(crate) use unexpected_downloaded_file_verifier::*;
     pub use stream_reader::*;
     pub use bootstrap_files::*;
 }