From 793d3c0c7014c17561f1715c77f187f52ca3c637 Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sat, 23 Mar 2024 18:41:59 -0500
Subject: [PATCH] Cleanup temporary directory on launch by default

---
 defaults.toml             |  3 ++-
 src/config/commandline.rs |  7 +++++++
 src/config/defaults.rs    |  2 ++
 src/config/file.rs        |  2 ++
 src/lib.rs                |  6 +++++-
 src/tmp_file.rs           | 31 ++++++++++++++++++++++++-------
 6 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/defaults.toml b/defaults.toml
index 4b6802b..60abfa3 100644
--- a/defaults.toml
+++ b/defaults.toml
@@ -4,6 +4,7 @@ read_only = false
 danger_dummy_mode = false
 max_file_count = 1
 temporary_directory = "/tmp"
+cleanup_temporary_directory = true
 
 [client]
 timeout = 30
@@ -46,7 +47,7 @@ proxy = "7d"
 [media.magick]
 max_width = 10000
 max_height = 10000
-max_area = 40000000
+max_area = 20000
 memory = 256
 map = 512
 disk = 1024
diff --git a/src/config/commandline.rs b/src/config/commandline.rs
index 035ecbd..d029709 100644
--- a/src/config/commandline.rs
+++ b/src/config/commandline.rs
@@ -55,6 +55,7 @@ impl Args {
                 address,
                 api_key,
                 temporary_directory,
+                no_cleanup_temporary_directory,
                 certificate,
                 private_key,
                 client_timeout,
@@ -122,6 +123,7 @@ impl Args {
                     danger_dummy_mode,
                     max_file_count,
                     temporary_directory,
+                    cleanup_temporary_directory: !no_cleanup_temporary_directory,
                     certificate,
                     private_key,
                 };
@@ -541,6 +543,7 @@ struct Server {
     max_file_count: Option<u32>,
     #[serde(skip_serializing_if = "Option::is_none")]
     temporary_directory: Option<PathBuf>,
+    cleanup_temporary_directory: bool,
     #[serde(skip_serializing_if = "Option::is_none")]
     certificate: Option<PathBuf>,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -973,6 +976,10 @@ struct Run {
     #[arg(long)]
     temporary_directory: Option<PathBuf>,
 
+    /// Whether to attempt to clean files left behind from a previous run of pict-rs
+    #[arg(long)]
+    no_cleanup_temporary_directory: bool,
+
     /// The path to the TLS certificate. Both the certificate and the private_key must be specified
     /// to enable TLS
     #[arg(long)]
diff --git a/src/config/defaults.rs b/src/config/defaults.rs
index aec1d98..c23a77d 100644
--- a/src/config/defaults.rs
+++ b/src/config/defaults.rs
@@ -24,6 +24,7 @@ struct ServerDefaults {
     danger_dummy_mode: bool,
     max_file_count: u32,
     temporary_directory: PathBuf,
+    cleanup_temporary_directory: bool,
 }
 
 #[derive(Clone, Debug, serde::Serialize)]
@@ -211,6 +212,7 @@ impl Default for ServerDefaults {
             danger_dummy_mode: false,
             max_file_count: 1,
             temporary_directory: std::env::temp_dir(),
+            cleanup_temporary_directory: true,
         }
     }
 }
diff --git a/src/config/file.rs b/src/config/file.rs
index 43b53bd..508878a 100644
--- a/src/config/file.rs
+++ b/src/config/file.rs
@@ -119,6 +119,8 @@ pub(crate) struct Server {
 
     pub(crate) temporary_directory: PathBuf,
 
+    pub(crate) cleanup_temporary_directory: bool,
+
     #[serde(skip_serializing_if = "Option::is_none")]
     pub(crate) certificate: Option<PathBuf>,
 
diff --git a/src/lib.rs b/src/lib.rs
index 5395e56..bc19baa 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1924,7 +1924,11 @@ impl PictRsConfiguration {
         // describe all the metrics pict-rs produces
         init_metrics::init_metrics();
 
-        let tmp_dir = TmpDir::init(&config.server.temporary_directory).await?;
+        let tmp_dir = TmpDir::init(
+            &config.server.temporary_directory,
+            config.server.cleanup_temporary_directory,
+        )
+        .await?;
         let policy_dir = magick::write_magick_policy(&config.media, &tmp_dir).await?;
 
         let client = build_client()?;
diff --git a/src/tmp_file.rs b/src/tmp_file.rs
index 46e9c14..aacca9a 100644
--- a/src/tmp_file.rs
+++ b/src/tmp_file.rs
@@ -16,9 +16,17 @@ pub(crate) struct TmpDir {
 }
 
 impl TmpDir {
-    pub(crate) async fn init<P: AsRef<Path>>(path: P) -> std::io::Result<Arc<Self>> {
-        let path = path.as_ref().join(Uuid::now_v7().to_string());
-        tokio::fs::create_dir(&path).await?;
+    pub(crate) async fn init<P: AsRef<Path>>(path: P, cleanup: bool) -> std::io::Result<Arc<Self>> {
+        let base_path = path.as_ref().join("pict-rs");
+
+        if cleanup && tokio::fs::metadata(&base_path).await.is_ok() {
+            tokio::fs::remove_dir_all(&base_path).await?;
+        }
+
+        let path = base_path.join(Uuid::now_v7().to_string());
+
+        tokio::fs::create_dir_all(&path).await?;
+
         Ok(Arc::new(TmpDir { path: Some(path) }))
     }
 
@@ -47,8 +55,13 @@ impl TmpDir {
     }
 
     pub(crate) async fn cleanup(self: Arc<Self>) -> std::io::Result<()> {
-        if let Some(path) = Arc::into_inner(self).and_then(|mut this| this.path.take()) {
-            tokio::fs::remove_dir_all(path).await?;
+        if let Some(mut path) = Arc::into_inner(self).and_then(|mut this| this.path.take()) {
+            tokio::fs::remove_dir_all(&path).await?;
+
+            if path.pop() {
+                // attempt to remove parent directory if it is empty
+                let _ = tokio::fs::remove_dir(path).await;
+            }
         }
 
         Ok(())
@@ -57,9 +70,13 @@ impl TmpDir {
 
 impl Drop for TmpDir {
     fn drop(&mut self) {
-        if let Some(path) = self.path.take() {
+        if let Some(mut path) = self.path.take() {
             tracing::warn!("TmpDir - Blocking remove of {path:?}");
-            std::fs::remove_dir_all(path).expect("Removed directory");
+            std::fs::remove_dir_all(&path).expect("Removed directory");
+            if path.pop() {
+                // attempt to remove parent directory if it is empty
+                let _ = std::fs::remove_dir(path);
+            }
         }
     }
 }