refactor: replace unsafe global environment variable manipulation with thread-local mocks and pure sanitization functions for improved test safety.

Author: fajarhide

src/guard/env.rs

@@ -28,7 +28,11 @@ pub const DENYLIST: &[&str] = &[
 ];
 
 pub fn sanitize_env() -> Vec<(String, String)> {
-    env::vars()
+    sanitize_vars(env::vars())
+}
+
+pub fn sanitize_vars(vars: impl IntoIterator<Item = (String, String)>) -> Vec<(String, String)> {
+    vars.into_iter()
         .filter(|(k, _)| !DENYLIST.iter().any(|d| d.eq_ignore_ascii_case(k)))
         .collect()
 }
@@ -44,46 +48,37 @@ mod tests {
 
     #[test]
     fn test_sanitize_env_menghapus_ld_preload() {
-        unsafe {
-            std::env::set_var("LD_PRELOAD", "bad.so");
-        }
-        let sanitized = sanitize_env();
+        let mock_env = vec![
+            ("LD_PRELOAD".to_string(), "bad.so".to_string()),
+            ("NORMAL_VAR".to_string(), "123".to_string()),
+        ];
+        let sanitized = sanitize_vars(mock_env);
         let contains = sanitized.iter().any(|(k, _)| k == "LD_PRELOAD");
         assert!(!contains);
-        unsafe {
-            std::env::remove_var("LD_PRELOAD");
-        }
     }
 
     #[test]
     fn test_sanitize_env_menghapus_semua_denylist_entries() {
-        for key in DENYLIST {
-            unsafe {
-                std::env::set_var(key, "malicious_payload");
-            }
-        }
+        let mock_env: Vec<(String, String)> = DENYLIST
+            .iter()
+            .map(|key| (key.to_string(), "malicious_payload".to_string()))
+            .collect();
 
-        let sanitized = sanitize_env();
+        let sanitized = sanitize_vars(mock_env);
 
         for (k, _) in sanitized {
             assert!(!DENYLIST.iter().any(|d| d.eq_ignore_ascii_case(&k)));
         }
-
-        for key in DENYLIST {
-            unsafe {
-                std::env::remove_var(key);
-            }
-        }
     }
 
     #[test]
     fn test_sanitize_env_mempertahankan_path_and_normal_vars() {
-        unsafe {
-            std::env::set_var("PATH", "/usr/bin:/bin");
-            std::env::set_var("NORMAL_VAR", "123");
-        }
+        let mock_env = vec![
+            ("PATH".to_string(), "/usr/bin:/bin".to_string()),
+            ("NORMAL_VAR".to_string(), "123".to_string()),
+        ];
 
-        let sanitized = sanitize_env();
+        let sanitized = sanitize_vars(mock_env);
         let has_path = sanitized.iter().any(|(k, _)| k.to_uppercase() == "PATH");
         let has_normal = sanitized
             .iter()

src/hooks/dispatcher.rs

@@ -146,10 +146,6 @@ mod tests {
             "workingDirectory": "/tmp"
         });
 
-        unsafe {
-            std::env::set_var("OMNI_CONTINUE", "1");
-            std::env::set_var("OMNI_FRESH", "0");
-        }
         let out = process_payload(&input.to_string(), store, session);
 
         assert!(out.is_some());

src/hooks/session_start.rs

@@ -192,9 +192,9 @@ mod tests {
         let db_path = dir.path().join("omni.db");
         // Set transcript dir to clean temp dir so find_pending() doesn't interfere
         let transcript_dir = dir.path().join("transcripts");
-        unsafe {
-            std::env::set_var("OMNI_TRANSCRIPT_DIR", transcript_dir.to_str().unwrap());
-        }
+        crate::store::transcript::MOCK_TRANSCRIPT_DIR.with(|d| {
+            *d.borrow_mut() = Some(transcript_dir);
+        });
         (
             Arc::new(Store::open_path(&db_path).expect("must succeed")),
             dir,

src/store/transcript.rs

@@ -337,7 +337,17 @@ pub fn cleanup_old(days: u32) {
 
 // ─── Helpers ────────────────────────────────────────────
 
+#[cfg(test)]
+thread_local! {
+    pub static MOCK_TRANSCRIPT_DIR: std::cell::RefCell<Option<PathBuf>> = const { std::cell::RefCell::new(None) };
+}
+
 fn transcripts_dir() -> PathBuf {
+    #[cfg(test)]
+    if let Some(mock) = MOCK_TRANSCRIPT_DIR.with(|d| d.borrow().clone()) {
+        return mock;
+    }
+
     if let Ok(custom) = std::env::var("OMNI_TRANSCRIPT_DIR") {
         return PathBuf::from(custom);
     }
@@ -368,19 +378,12 @@ fn truncate_payload(payload: &str, max_bytes: usize) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::sync::Mutex;
     use tempfile::tempdir;
 
-    // Serialize all transcript tests — they share the OMNI_TRANSCRIPT_DIR env var
-    static TEST_LOCK: Mutex<()> = Mutex::new(());
-
-    fn setup_test_dir() -> (tempfile::TempDir, std::sync::MutexGuard<'static, ()>) {
-        let guard = TEST_LOCK.lock().unwrap();
+    fn setup_test_dir() -> tempfile::TempDir {
         let dir = tempdir().unwrap();
-        unsafe {
-            std::env::set_var("OMNI_TRANSCRIPT_DIR", dir.path().to_str().unwrap());
-        }
-        (dir, guard)
+        MOCK_TRANSCRIPT_DIR.with(|d| *d.borrow_mut() = Some(dir.path().to_path_buf()));
+        dir
     }
 
     #[test]
@@ -414,7 +417,7 @@ mod tests {
 
     #[test]
     fn test_save_and_load_roundtrip() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut t = Transcript::new("roundtrip_1", "/project");
         let entry = TranscriptEntry::new_input("git status", Some("git"));
@@ -428,7 +431,7 @@ mod tests {
 
     #[test]
     fn test_atomic_write_no_corrupt_file() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let t = Transcript::new("atomic_1", "/project");
         t.save().unwrap();
@@ -444,7 +447,7 @@ mod tests {
 
     #[test]
     fn test_mark_last_completed() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut t = Transcript::new("complete_1", "/project");
         t.append_entry(TranscriptEntry::new_input("input1", None))
@@ -461,7 +464,7 @@ mod tests {
 
     #[test]
     fn test_mark_last_failed() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut t = Transcript::new("fail_1", "/project");
         t.append_entry(TranscriptEntry::new_input("bad input", None))
@@ -497,13 +500,13 @@ mod tests {
 
     #[test]
     fn test_find_pending_with_no_transcripts() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
         assert!(find_pending().is_none());
     }
 
     #[test]
     fn test_find_pending_finds_interrupted_session() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut t = Transcript::new("interrupted_1", "/project");
         t.append_entry(TranscriptEntry::new_input("pending work", None))
@@ -516,7 +519,7 @@ mod tests {
 
     #[test]
     fn test_find_pending_skips_completed_sessions() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut t = Transcript::new("done_1", "/project");
         t.append_entry(TranscriptEntry::new_input("done work", None))
@@ -528,7 +531,7 @@ mod tests {
 
     #[test]
     fn test_load_or_new_creates_if_missing() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let t = Transcript::load_or_new("new_session", "/project");
         assert_eq!(t.session_id, "new_session");
@@ -537,7 +540,7 @@ mod tests {
 
     #[test]
     fn test_load_or_new_loads_if_exists() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut original = Transcript::new("existing_1", "/project");
         original
@@ -551,7 +554,7 @@ mod tests {
 
     #[test]
     fn test_cleanup_old_removes_stale_transcripts() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         // Create a transcript with very old updated_at
         let mut t = Transcript::new("old_1", "/project");
@@ -570,7 +573,7 @@ mod tests {
 
     #[test]
     fn test_list_recent_returns_all() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let t1 = Transcript::new("list_a", "/project");
         t1.save().unwrap();
@@ -600,7 +603,7 @@ mod tests {
 
     #[test]
     fn test_snapshot_state_persists() {
-        let (_dir, _lock) = setup_test_dir();
+        let _dir = setup_test_dir();
 
         let mut t = Transcript::new("state_1", "/project");
         let mut state = SessionState::new();

tests/security_tests.rs

@@ -11,32 +11,24 @@ fn run_pipeline(input: &str) -> String {
 
 #[test]
 fn test_env_sanitization_denylist() {
-    use omni::guard::env::{DENYLIST, sanitize_env};
+    use omni::guard::env::{DENYLIST, sanitize_vars};
 
-    // Set some dangerous env vars (unsafe in Rust 2024)
+    // Set some dangerous env vars in a mock environment
+    let mut mock_env: Vec<(String, String)> = Vec::new();
     for var in DENYLIST.iter().take(3) {
-        unsafe {
-            std::env::set_var(var, "INJECTED_VALUE");
-        }
+        mock_env.push((var.to_string(), "INJECTED_VALUE".to_string()));
     }
 
-    let sanitized = sanitize_env();
+    let sanitized = sanitize_vars(mock_env);
 
     // Verify denylist vars are NOT in sanitized output
     for var in DENYLIST {
         assert!(
             !sanitized.iter().any(|(k, _)| k.eq_ignore_ascii_case(var)),
-            "Denylist variable {} should be removed by sanitize_env",
+            "Denylist variable {} should be removed by sanitize_vars",
             var
         );
     }
-
-    // Cleanup
-    for var in DENYLIST.iter().take(3) {
-        unsafe {
-            std::env::remove_var(var);
-        }
-    }
 }
 
 #[test]
@@ -115,16 +107,17 @@ fn test_pipeline_deterministic() {
 
 #[test]
 fn test_env_sanitization_removes_dangerous_vars() {
-    use omni::guard::env::{DENYLIST, sanitize_env};
-
-    // Set beberapa dangerous vars
-    unsafe {
-        std::env::set_var("LD_PRELOAD", "malicious.so");
-        std::env::set_var("BASH_ENV", "evil_script.sh");
-        std::env::set_var("NODE_OPTIONS", "--require=evil");
-    }
+    use omni::guard::env::{DENYLIST, sanitize_vars};
+
+    // Set beberapa dangerous vars in a mock env
+    let mock_env = vec![
+        ("LD_PRELOAD".to_string(), "malicious.so".to_string()),
+        ("BASH_ENV".to_string(), "evil_script.sh".to_string()),
+        ("NODE_OPTIONS".to_string(), "--require=evil".to_string()),
+        ("PATH".to_string(), "/usr/bin:/bin".to_string()),
+    ];
 
-    let sanitized = sanitize_env();
+    let sanitized = sanitize_vars(mock_env);
 
     // Verify semua DENYLIST entries hilang
     for key in DENYLIST {
@@ -140,13 +133,6 @@ fn test_env_sanitization_removes_dangerous_vars() {
         sanitized.iter().any(|(k, _)| k.to_uppercase() == "PATH"),
         "PATH should still be in sanitized env"
     );
-
-    // Cleanup
-    unsafe {
-        std::env::remove_var("LD_PRELOAD");
-        std::env::remove_var("BASH_ENV");
-        std::env::remove_var("NODE_OPTIONS");
-    }
 }
 
 #[test]