From 3fa93c6d9c09b4623687bea46845921807260084 Mon Sep 17 00:00:00 2001
From: Wolf Vollprecht <w.vollprecht@gmail.com>
Date: Thu, 11 Jan 2024 16:02:33 +0100
Subject: [PATCH] add drop bomb

---
 crates/rattler/Cargo.toml                     |  1 +
 .../rattler/src/install/clobber_registry.rs   | 20 +++++++++++++++++--
 crates/rattler/src/install/mod.rs             | 14 ++++++++++++-
 crates/rattler/src/install/unlink.rs          |  7 ++++++-
 4 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/crates/rattler/Cargo.toml b/crates/rattler/Cargo.toml
index 5a38167df..4aa39446c 100644
--- a/crates/rattler/Cargo.toml
+++ b/crates/rattler/Cargo.toml
@@ -22,6 +22,7 @@ bytes = "1.5.0"
 chrono = { version = "0.4.31", default-features = false, features = ["std", "serde", "alloc"] }
 digest = "0.10.7"
 dirs = "5.0.1"
+drop_bomb = "0.1.5"
 futures = "0.3.28"
 fxhash = "0.2.1"
 hex = "0.4.3"
diff --git a/crates/rattler/src/install/clobber_registry.rs b/crates/rattler/src/install/clobber_registry.rs
index 4fa4144f4..d037b2968 100644
--- a/crates/rattler/src/install/clobber_registry.rs
+++ b/crates/rattler/src/install/clobber_registry.rs
@@ -6,16 +6,31 @@ use std::{
     path::{Path, PathBuf},
 };
 
+use drop_bomb::DropBomb;
 use rattler_conda_types::{package::PathsJson, PackageName, PrefixRecord};
 
 /// A registry for clobbering files
 /// The registry keeps track of all files that are installed by a package and
 /// can be used to rename files that are already installed by another package.
-#[derive(Default, Debug)]
+#[derive(Debug)]
 pub struct ClobberRegistry {
     paths_registry: HashMap<PathBuf, usize>,
     clobbers: HashMap<PathBuf, Vec<usize>>,
     package_names: Vec<PackageName>,
+    drop_bomb: DropBomb,
+}
+
+impl Default for ClobberRegistry {
+    fn default() -> Self {
+        Self {
+            paths_registry: HashMap::new(),
+            clobbers: HashMap::new(),
+            package_names: Vec::new(),
+            drop_bomb: DropBomb::new(
+                "did not call post_process on InstallDriver / ClobberRegistry",
+            ),
+        }
+    }
 }
 
 static CLOBBER_TEMPLATE: &str = "__clobber-from-";
@@ -131,10 +146,11 @@ impl ClobberRegistry {
 
     /// Unclobber the paths after all installation steps have been completed.
     pub fn post_process(
-        &self,
+        &mut self,
         sorted_prefix_records: &[&PrefixRecord],
         target_prefix: &Path,
     ) -> Result<(), std::io::Error> {
+        self.drop_bomb.defuse();
         let sorted_names = sorted_prefix_records
             .iter()
             .map(|p| p.repodata_record.package_record.name.clone())
diff --git a/crates/rattler/src/install/mod.rs b/crates/rattler/src/install/mod.rs
index fcf7b72c9..b74dff30c 100644
--- a/crates/rattler/src/install/mod.rs
+++ b/crates/rattler/src/install/mod.rs
@@ -714,6 +714,12 @@ mod test {
             })
             .await;
 
+        // test doesn't write conda-meta, so we ignore the post processing
+        let prefix_records = vec![];
+        install_driver
+            .post_process(&prefix_records, target_dir.path())
+            .unwrap();
+
         // Run the python command and validate the version it outputs
         let python_path = if Platform::current().is_windows() {
             "python.exe"
@@ -745,16 +751,22 @@ mod test {
         )
         .unwrap();
 
+        let install_driver = InstallDriver::default();
+
         // Link the package
         let paths = link_package(
             package_dir.path(),
             environment_dir.path(),
-            &InstallDriver::default(),
+            &install_driver,
             InstallOptions::default(),
         )
         .await
         .unwrap();
 
+        install_driver
+            .post_process(&vec![], environment_dir.path())
+            .unwrap();
+
         insta::assert_yaml_snapshot!(paths);
     }
 }
diff --git a/crates/rattler/src/install/unlink.rs b/crates/rattler/src/install/unlink.rs
index 3eb84f61a..45acda8ba 100644
--- a/crates/rattler/src/install/unlink.rs
+++ b/crates/rattler/src/install/unlink.rs
@@ -184,11 +184,12 @@ mod tests {
             ..InstallOptions::default()
         };
 
+        let install_driver = InstallDriver::default();
         // Link the package
         let paths = link_package(
             package_dir.path(),
             target_prefix,
-            &InstallDriver::default(),
+            &install_driver,
             install_options,
         )
         .await
@@ -199,6 +200,10 @@ mod tests {
         let prefix_record =
             PrefixRecord::from_repodata_record(repodata_record, None, None, paths, None, None);
 
+        install_driver
+            .post_process(&vec![prefix_record.clone()], target_prefix)
+            .unwrap();
+
         return prefix_record;
     }