From 821f0e6472879d4e998cc7327ceec6fbe359ad4f Mon Sep 17 00:00:00 2001
From: Sosthene <contact@sosthene.net>
Date: Tue, 19 Nov 2024 13:38:20 +0100
Subject: [PATCH] Add cleanup routine

---
 src/main.rs | 101 ++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 94 insertions(+), 7 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 84fba2e..db36845 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,14 +1,80 @@
-use async_std::fs::{create_dir_all, File};
+use async_std::fs::{create_dir_all, read_dir, read_to_string, remove_file, File};
+use std::time::{SystemTime, UNIX_EPOCH, Duration};
+use serde::{Serialize, Deserialize};
+
 use async_std::io::WriteExt;
-use serde::{Deserialize, Serialize};
+use async_std::path::Path;
+use async_std::stream::{IntoStream, StreamExt};
+use async_std::task;
+use sdk_common::log;
 use serde_json::json;
 use tide::{Request, Response, StatusCode};
 use base64;
-use std::time::{Duration, SystemTime};
 
 const STORAGE_DIR: &str = "./storage";
 const DEFAULT_TTL: u64 = 86400;
 
+/// Scans storage and removes expired files
+async fn cleanup_expired_files() {
+    loop {
+        log::debug!("Running cleanup routine...");
+        // Traverse storage directory
+        let mut entries = match read_dir(STORAGE_DIR).await {
+            Ok(entry) => entry,
+            Err(e) => {
+                log::error!("Failed to read storage dir: {}", e);
+                task::sleep(Duration::from_secs(60)).await;
+                continue;
+            }
+        };
+
+        while let Some(entry) = entries.next().await {
+            let e = match entry {
+                Ok(e) => e,
+                Err(e) => {
+                    log::error!("entry returned error: {}", e);
+                    continue;
+                }
+            };
+            let path = e.path();
+            if path.is_dir().await {
+                if let Ok(mut sub_entries) = read_dir(&path).await {
+                    while let Some(sub_entry) = sub_entries.next().await {
+                        if let Ok(sub_entry) = sub_entry {
+                            let file_path = sub_entry.path();
+                            if file_path.extension() == Some("meta".as_ref()) {
+                                if let Err(err) = handle_file_cleanup(&file_path).await {
+                                    log::error!("Error cleaning file {:?}: {}", file_path, err);
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        // Sleep for 1 minute before next cleanup
+        // task::sleep(Duration::from_secs(60)).await;
+        task::sleep(Duration::from_secs(10)).await;
+    }
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+struct Metadata {
+    expires_at: u64,
+}
+
+/// Converts a `SystemTime` to a UNIX timestamp (seconds since UNIX epoch).
+fn system_time_to_unix(system_time: SystemTime) -> u64 {
+    system_time.duration_since(UNIX_EPOCH)
+        .expect("SystemTime before UNIX_EPOCH!")
+        .as_secs()
+}
+
+/// Converts a UNIX timestamp (seconds since UNIX epoch) back to `SystemTime`.
+fn unix_to_system_time(unix_timestamp: u64) -> SystemTime {
+    UNIX_EPOCH + Duration::from_secs(unix_timestamp)
+}
+
 #[derive(Deserialize)]
 struct StoreRequest {
     key: String,
@@ -38,10 +104,13 @@ async fn store_data(key: &str, value: &[u8], expires_at: SystemTime) -> Result<(
     let mut file = File::create(&file_path).await.map_err(|e| e.to_string())?;
     file.write_all(value).await.map_err(|e| e.to_string())?;
 
-    let metadata = serde_json::to_string(&json!({ "expires_at": expires_at }))
-        .map_err(|e| e.to_string())?;
+    let metadata = Metadata {
+        expires_at: system_time_to_unix(expires_at)
+    };
+
+    let metadata_json = serde_json::to_string(&metadata).map_err(|e| e.to_string())?;
     let mut meta_file = File::create(&metadata_path).await.map_err(|e| e.to_string())?;
-    meta_file.write_all(metadata.as_bytes()).await.map_err(|e| e.to_string())?;
+    meta_file.write_all(metadata_json.as_bytes()).await.map_err(|e| e.to_string())?;
 
     Ok(())
 }
@@ -113,16 +182,34 @@ async fn handle_retrieve(req: Request<()>) -> tide::Result<Response> {
     }
 }
 
+/// Checks a metadata file and deletes the associated data file if expired
+async fn handle_file_cleanup(meta_path: &Path) -> Result<(), String> {
+    let meta_content = read_to_string(meta_path).await.map_err(|e| format!("Failed to read metadata: {}", e.to_string()))?;
+    let metadata: Metadata = serde_json::from_str(&meta_content).map_err(|e| format!("Failed to parse metadata: {}", e.to_string()))?;
+
+    let now = system_time_to_unix(SystemTime::now());
+    if metadata.expires_at < now {
+        let data_file_path = meta_path.with_extension(""); 
+        remove_file(&data_file_path).await.map_err(|e| format!("Failed to remove data file: {}", e.to_string()))?;
+        remove_file(meta_path).await.map_err(|e| format!("Failed to remove metadata file: {}", e.to_string()))?;
+        log::debug!("Removed expired file: {:?}", data_file_path);
+    }
+    Ok(())
+}
+
 #[async_std::main]
 async fn main() -> tide::Result<()> {
+    sdk_common::env_logger::init();
     create_dir_all(STORAGE_DIR).await.expect("Failed to create storage directory.");
 
+    task::spawn(cleanup_expired_files());
+
     let mut app = tide::new();
     app.at("/store").post(handle_store);
     app.at("/retrieve/:key").get(handle_retrieve);
     app.listen("0.0.0.0:8080").await?;
 
-    println!("Server running at http://0.0.0.0:8080");
+    log::info!("Server running at http://0.0.0.0:8080");
 
     Ok(())
 }