From c88ef1b47449d83119b3aa92978e54d68fe8e486 Mon Sep 17 00:00:00 2001
From: tmpolaczyk <44604217+tmpolaczyk@users.noreply.github.com>
Date: Wed, 22 May 2024 13:27:47 +0200
Subject: [PATCH] Remove empty folders after removing container db folder
 (#556)

* Remove empty folders after removing container db folder

We still leave one folder, "network", because it has a p2p key inside:
"secret_ed25519".
---
 node/src/container_chain_spawner.rs | 57 ++++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 6 deletions(-)

diff --git a/node/src/container_chain_spawner.rs b/node/src/container_chain_spawner.rs
index 8670ae9fd..31d30071b 100644
--- a/node/src/container_chain_spawner.rs
+++ b/node/src/container_chain_spawner.rs
@@ -742,16 +742,46 @@ fn open_and_maybe_delete_db(
     Ok(())
 }
 
-// TODO: this leaves some empty folders behind, because it is called with db_path:
-//     Collator2002-01/data/containers/chains/simple_container_2002/paritydb/full-container-2002
-// but we want to delete everything under
-//     Collator2002-01/data/containers/chains/simple_container_2002
+/// Remove the container chain database folder. This is called with db_path:
+///     `Collator2002-01/data/containers/chains/simple_container_2002/paritydb/full-container-2002`
+/// but we want to delete everything under
+///     `Collator2002-01/data/containers/chains/simple_container_2002`
+/// So we use `delete_empty_folders_recursive` to try to remove the parent folders as well, but only
+/// if they are empty. This is to avoid removing any secret keys or other important data.
 fn delete_container_chain_db(db_path: &Path) {
-    if db_path.exists() {
-        std::fs::remove_dir_all(db_path).expect("failed to remove old container chain db");
+    // Remove folder `full-container-2002`
+    let _ = std::fs::remove_dir_all(db_path);
+    // Remove all the empty folders inside `simple_container_2002`, including self
+    if let Some(parent) = db_path.ancestors().nth(2) {
+        let _ = delete_empty_folders_recursive(parent);
     }
 }
 
+/// Removes all empty folders in `path`, recursively. Then, if `path` is empty, it removes it as well.
+/// Ignores any IO errors.
+fn delete_empty_folders_recursive(path: &Path) {
+    let entry_iter = std::fs::read_dir(path);
+    let entry_iter = match entry_iter {
+        Ok(x) => x,
+        Err(_e) => return,
+    };
+
+    for entry in entry_iter {
+        let entry = match entry {
+            Ok(x) => x,
+            Err(_e) => continue,
+        };
+
+        let path = entry.path();
+        if path.is_dir() {
+            let _ = delete_empty_folders_recursive(&path);
+        }
+    }
+
+    // Try to remove dir. Returns an error if the directory is not empty, but we ignore it.
+    let _ = std::fs::remove_dir(path);
+}
+
 /// Parse a list of boot nodes in `Vec<u8>` format. Invalid boot nodes are filtered out.
 fn parse_boot_nodes_ignore_invalid(
     boot_nodes_raw: Vec<Vec<u8>>,
@@ -786,6 +816,7 @@ fn parse_boot_nodes_ignore_invalid(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::path::PathBuf;
 
     // Copy of ContainerChainSpawner with extra assertions for tests, and mocked spawn function.
     struct MockContainerChainSpawner {
@@ -1217,4 +1248,18 @@ mod tests {
             1
         );
     }
+
+    #[test]
+    fn path_ancestors() {
+        // Test the implementation of `delete_container_chain_db`
+        let db_path = PathBuf::from("/tmp/zombienet/Collator2002-01/data/containers/chains/simple_container_2002/paritydb/full-container-2002");
+        let parent = db_path.ancestors().nth(2).unwrap();
+
+        assert_eq!(
+            parent,
+            PathBuf::from(
+                "/tmp/zombienet/Collator2002-01/data/containers/chains/simple_container_2002"
+            )
+        )
+    }
 }