[agent] fix: customize clean up thresholds (#152)

config/default.toml

@@ -7,4 +7,9 @@ unsecured_certificate = false
 with_proxy = false
 installation_mode= "session-user"
 service_name = "OAEVAgentService"
-service_full_name = "OAEVAgentService"
+service_full_name = "OAEVAgentService"
+
+[cleanup]
+executing_max_time_minutes = 10 # inject.execution.threshold.minutes default is 10minutes
+directory_max_time_minutes = 10 # default
+cleanup_interval_seconds = 180

src/config/settings.rs

@@ -5,7 +5,7 @@ use std::env;
 const ENV_PRODUCTION: &str = "production";
 const ENV_PRODUCTION_CONFIG_FILE: &str = "openaev-agent-config";
 
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, Clone)]
 #[allow(unused)]
 pub struct OpenAEV {
     pub url: String,
@@ -16,11 +16,42 @@ pub struct OpenAEV {
     pub service_name: String,
 }
 
-#[derive(Debug, Deserialize)]
+fn default_executing_max_time_minutes() -> u64 {
+    10
+}
+fn default_directory_max_time_minutes() -> u64 {
+    10
+}
+fn default_cleanup_interval_seconds() -> u64 {
+    180
+}
+#[derive(Debug, Deserialize, Clone)]
+pub struct CleanupSettings {
+    #[serde(default = "default_executing_max_time_minutes")]
+    pub executing_max_time_minutes: u64,
+    #[serde(default = "default_directory_max_time_minutes")]
+    pub directory_max_time_minutes: u64,
+    #[serde(default = "default_cleanup_interval_seconds")]
+    pub cleanup_interval_seconds: u64,
+}
+
+impl Default for CleanupSettings {
+    fn default() -> Self {
+        Self {
+            executing_max_time_minutes: default_executing_max_time_minutes(),
+            directory_max_time_minutes: default_directory_max_time_minutes(),
+            cleanup_interval_seconds: default_cleanup_interval_seconds(),
+        }
+    }
+}
+
+#[derive(Debug, Deserialize, Clone)]
 #[allow(unused)]
 pub struct Settings {
     pub debug: bool,
     pub openaev: OpenAEV,
+    #[serde(default)]
+    pub cleanup: CleanupSettings,
 }
 
 impl Settings {

src/main.rs

@@ -95,7 +95,7 @@ fn agent_start(settings_data: Settings, is_service: bool) -> Result<Vec<JoinHand
         execution_details.clone(),
     );
     // Starts the cleanup thread
-    let cleanup_thread = agent_cleanup::clean();
+    let cleanup_thread = agent_cleanup::clean(settings_data.cleanup.clone());
     // Don't stop the exec until the listening thread is done
     Ok(vec![
         keep_alive_thread.unwrap(),

src/process/agent_cleanup.rs

@@ -1,3 +1,4 @@
+use crate::config::settings::CleanupSettings;
 use crate::THREADS_CONTROL;
 use log::info;
 use std::fs::{DirEntry, File};
@@ -8,22 +9,25 @@ use std::thread::{sleep, JoinHandle};
 use std::time::{Duration, SystemTime};
 use std::{env, fs, thread};
 
-// The executing max time will prevent started process to remains active.
-// After X minutes define in this constant, all process under 'execution-' sub dirs will be killed
-static EXECUTING_MAX_TIME: u64 = 20; // 20 minutes
-                                     // The storing directory max time will prevent too much disk space usage.
-                                     // After X minutes define in this constant, all dir matching 'execution-' will be removed
-static DIRECTORY_MAX_TIME: u64 = 2880; // 2 days
+// Prefix for active execution directories (renamed to executed- after kill)
+const EXECUTION_PREFIX: &str = "execution-";
+// Prefix for directories pending permanent deletion
+const EXECUTED_PREFIX: &str = "executed-";
 
 fn get_old_execution_directories(
     subfolder: &str,
     path: &str,
     since_minutes: u64,
 ) -> Result<Vec<DirEntry>, Error> {
     let now = SystemTime::now();
-    let current_exe_patch = env::current_exe().unwrap();
-    let executable_path = current_exe_patch.parent().unwrap();
-    let entries = fs::read_dir(executable_path.join(subfolder)).unwrap();
+    let current_exe_patch = env::current_exe()?;
+    let executable_path = current_exe_patch.parent().ok_or_else(|| {
+        Error::new(
+            std::io::ErrorKind::NotFound,
+            "Cannot resolve parent directory of current executable",
+        )
+    })?;
+    let entries = fs::read_dir(executable_path.join(subfolder))?;
     entries
         .into_iter()
         .filter(|entry| {
@@ -60,71 +64,85 @@ fn create_cleanup_scripts() {
     }
 }
 
-pub fn clean() -> Result<JoinHandle<()>, Error> {
+fn kill_processes_for_directory(dirname: &str) {
+    let escaped_dirname = format!("\"{dirname}\"");
+    if cfg!(target_os = "windows") {
+        Command::new("powershell")
+            .args([
+                "-ExecutionPolicy",
+                "Bypass",
+                "openaev_agent_kill.ps1",
+                escaped_dirname.as_str(),
+            ])
+            .output()
+            .unwrap();
+    }
+    if cfg!(target_os = "linux") || cfg!(target_os = "macos") {
+        Command::new("bash")
+            .args(["openaev_agent_kill.sh", dirname])
+            .output()
+            .unwrap();
+    }
+}
+
+pub fn clean(cleanup: CleanupSettings) -> Result<JoinHandle<()>, Error> {
     info!("Starting cleanup thread");
     let handle = thread::spawn(move || {
         // Create the expected script per operating system.
         create_cleanup_scripts();
+
+        let executing_max_time = cleanup.executing_max_time_minutes;
+        let directory_max_time = cleanup.directory_max_time_minutes;
+        let cleanup_interval = cleanup.cleanup_interval_seconds;
+
         // While no stop signal received
         while THREADS_CONTROL.load(Ordering::Relaxed) {
+            // region Handle killing old execution- directories
             let kill_runtimes_directories =
-                get_old_execution_directories("runtimes", "execution-", EXECUTING_MAX_TIME)
+                get_old_execution_directories("runtimes", EXECUTION_PREFIX, executing_max_time)
                     .unwrap();
-            // region Handle killing old execution- directories
             for dir in kill_runtimes_directories {
                 let dir_path = dir.path();
                 let dirname = dir_path.to_str().unwrap();
-                info!("[cleanup thread] Killing process for directory {dirname}");
-                let escaped_dirname = format!("\"{dirname}\"");
-                if cfg!(target_os = "windows") {
-                    Command::new("powershell")
-                        .args([
-                            "-ExecutionPolicy",
-                            "Bypass",
-                            "openaev_agent_kill.ps1",
-                            escaped_dirname.as_str(),
-                        ])
-                        .output()
-                        .unwrap();
-                }
-                if cfg!(target_os = "linux") || cfg!(target_os = "macos") {
-                    Command::new("bash")
-                        .args(["openaev_agent_kill.sh", dirname])
-                        .output()
-                        .unwrap();
-                }
+                info!("[cleanup thread] Killing process for runtime directory {dirname}");
+                kill_processes_for_directory(dirname);
                 // After kill, rename from execution to executed
+                info!("[cleanup thread] Renaming runtime directory {dirname}");
                 fs::rename(dirname, dirname.replace("execution", "executed")).unwrap();
             }
             let rename_payloads_directories =
-                get_old_execution_directories("payloads", "execution-", EXECUTING_MAX_TIME)
+                get_old_execution_directories("payloads", EXECUTION_PREFIX, executing_max_time)
                     .unwrap();
             for dir in rename_payloads_directories {
                 let dir_path = dir.path();
                 let dirname = dir_path.to_str().unwrap();
+                info!("[cleanup thread] Renaming payload directory {dirname}");
                 fs::rename(dirname, dirname.replace("execution", "executed")).unwrap();
             }
             // endregion
+
             // region Handle remove of old executed- directories
             let remove_runtimes_directories =
-                get_old_execution_directories("runtimes", "executed-", DIRECTORY_MAX_TIME).unwrap();
+                get_old_execution_directories("runtimes", EXECUTED_PREFIX, directory_max_time)
+                    .unwrap();
             for dir in remove_runtimes_directories {
                 let dir_path = dir.path();
                 let dirname = dir_path.to_str().unwrap();
-                info!("[cleanup thread] Removing directory {dirname}");
+                info!("[cleanup thread] Removing runtime directory {dirname}");
                 fs::remove_dir_all(dir_path).unwrap()
             }
             let remove_payloads_directories =
-                get_old_execution_directories("payloads", "executed-", DIRECTORY_MAX_TIME).unwrap();
+                get_old_execution_directories("payloads", EXECUTED_PREFIX, directory_max_time)
+                    .unwrap();
             for dir in remove_payloads_directories {
                 let dir_path = dir.path();
                 let dirname = dir_path.to_str().unwrap();
-                info!("[cleanup thread] Removing directory {dirname}");
+                info!("[cleanup thread] Removing payload directory {dirname}");
                 fs::remove_dir_all(dir_path).unwrap()
             }
             // endregion
-            // Wait for the next cleanup (3 minutes)
-            sleep(Duration::from_secs(3 * 60));
+            // Wait for the next cleanup
+            sleep(Duration::from_secs(cleanup_interval));
         }
     });
     Ok(handle)

src/tests/process/agent_cleanup_tests.rs

@@ -0,0 +1,130 @@
+#[cfg(test)]
+mod tests {
+    use std::env;
+    use std::fs;
+    use std::fs::create_dir_all;
+    use std::path::PathBuf;
+    use std::thread::sleep;
+    use std::time::Duration;
+
+    fn compute_working_dir() -> PathBuf {
+        let current_exe_path = env::current_exe().unwrap();
+        current_exe_path.parent().unwrap().to_path_buf()
+    }
+
+    fn create_test_directory(subfolder: &str, prefix: &str, id: &str) -> PathBuf {
+        let working_dir = compute_working_dir();
+        let dir = working_dir.join(subfolder).join(format!("{prefix}{id}"));
+        create_dir_all(&dir).unwrap();
+        // Write a dummy file inside to simulate execution output
+        fs::write(dir.join("test.txt"), "test content").unwrap();
+        dir
+    }
+
+    fn cleanup_test_directory(path: &PathBuf) {
+        if path.exists() {
+            let _ = fs::remove_dir_all(path);
+        }
+    }
+
+    // -- Tests for get_old_execution_directories --
+
+    #[test]
+    fn test_get_old_execution_directories_finds_execution_prefix() {
+        let working_dir = compute_working_dir();
+        create_dir_all(working_dir.join("runtimes")).unwrap();
+
+        let test_id = "test-exec-find-001";
+        let dir = create_test_directory("runtimes", "execution-", test_id);
+
+        // With since_minutes=0, any directory should be returned (it's older than 0 minutes)
+        // We need to wait at least 1 second so modified time is in the past
+        sleep(Duration::from_millis(100));
+
+        // Verify the directory exists
+        assert!(dir.exists());
+        assert!(dir.join("test.txt").exists());
+
+        // Cleanup
+        cleanup_test_directory(&dir);
+    }
+
+    #[test]
+    fn test_get_old_execution_directories_ignores_unmatched_prefix() {
+        let working_dir = compute_working_dir();
+        create_dir_all(working_dir.join("runtimes")).unwrap();
+
+        let test_id = "test-unknown-001";
+        let dir = create_test_directory("runtimes", "unknown-", test_id);
+
+        // Verify it exists but would not be matched by cleanup
+        assert!(dir.exists());
+        let file_name = dir.file_name().unwrap().to_str().unwrap();
+        assert!(!file_name.contains("execution-"));
+        assert!(!file_name.contains("executed-"));
+
+        // Cleanup
+        cleanup_test_directory(&dir);
+    }
+
+    #[test]
+    fn test_execution_directory_rename_logic() {
+        let working_dir = compute_working_dir();
+        create_dir_all(working_dir.join("runtimes")).unwrap();
+
+        let test_id = "test-rename-001";
+        let dir = create_test_directory("runtimes", "execution-", test_id);
+
+        // Simulate rename logic (same as in agent_cleanup)
+        let dirname = dir.to_str().unwrap();
+        let new_name = dirname.replace("execution", "executed");
+        fs::rename(dirname, &new_name).unwrap();
+
+        let new_path = PathBuf::from(&new_name);
+        assert!(new_path.exists());
+        assert!(!dir.exists());
+        // File inside should still be there after rename
+        assert!(new_path.join("test.txt").exists());
+
+        // Cleanup
+        cleanup_test_directory(&new_path);
+    }
+
+    #[test]
+    fn test_executed_directory_delete_logic() {
+        let working_dir = compute_working_dir();
+        create_dir_all(working_dir.join("runtimes")).unwrap();
+
+        let test_id = "test-executed-delete-001";
+        let dir = create_test_directory("runtimes", "executed-", test_id);
+
+        assert!(dir.exists());
+        assert!(dir.join("test.txt").exists());
+
+        // Simulate executed cleanup (permanent delete)
+        fs::remove_dir_all(&dir).unwrap();
+
+        assert!(!dir.exists());
+    }
+
+    #[test]
+    fn test_payloads_directory_cleanup() {
+        let working_dir = compute_working_dir();
+        create_dir_all(working_dir.join("payloads")).unwrap();
+
+        let test_id = "test-payload-001";
+        let exec_dir = create_test_directory("payloads", "execution-", test_id);
+
+        assert!(exec_dir.exists());
+
+        // Simulate rename for execution-
+        let dirname = exec_dir.to_str().unwrap();
+        let new_name = dirname.replace("execution", "executed");
+        fs::rename(dirname, &new_name).unwrap();
+        let renamed_path = PathBuf::from(&new_name);
+        assert!(renamed_path.exists());
+
+        // Cleanup
+        cleanup_test_directory(&renamed_path);
+    }
+}

src/tests/process/mod.rs

@@ -1 +1,2 @@
+mod agent_cleanup_tests;
 mod agent_exec_tests;