Skip to main content

mars_agents/fs/
mod.rs

1use std::fs;
2use std::io::Write;
3use std::os::unix::io::AsRawFd;
4use std::path::Path;
5
6use crate::error::MarsError;
7use crate::types::ItemKind;
8
9/// Top-level source entries excluded when installing flat skill repositories.
10pub const FLAT_SKILL_EXCLUDED_TOP_LEVEL: &[&str] = &[
11    ".git",
12    ".mars",
13    "mars.toml",
14    "mars.lock",
15    "mars.local.toml",
16    ".gitignore",
17];
18
19/// Atomic file write: write to temp file in same directory, then rename.
20///
21/// The rename is atomic on POSIX. Temp files are in the same directory
22/// as the destination to guarantee same-filesystem atomic rename.
23pub fn atomic_write(dest: &Path, content: &[u8]) -> Result<(), MarsError> {
24    // Ensure parent directory exists
25    if let Some(parent) = dest.parent() {
26        fs::create_dir_all(parent)?;
27    }
28
29    let parent = dest.parent().unwrap_or(Path::new("."));
30    let mut tmp = tempfile::NamedTempFile::new_in(parent)?;
31    tmp.write_all(content)?;
32    tmp.as_file().sync_all()?;
33    #[cfg(unix)]
34    {
35        use std::os::unix::fs::PermissionsExt;
36        tmp.as_file()
37            .set_permissions(fs::Permissions::from_mode(0o644))?;
38    }
39    tmp.persist(dest).map_err(|e| e.error)?;
40    Ok(())
41}
42
43/// Atomic directory install: copy source tree to a temp dir in the same
44/// parent as `dest`, then rename into place.
45///
46/// Uses rename-old-then-rename-new to minimize the window where `dest`
47/// doesn't exist. If `dest` already exists, it's renamed to `.{name}.old`
48/// before the new content takes its place. Stale `.old` from prior crashes
49/// is cleaned up automatically.
50pub fn atomic_install_dir(src: &Path, dest: &Path) -> Result<(), MarsError> {
51    atomic_install_dir_impl(src, dest, &[])
52}
53
54/// Atomic directory install with optional top-level source entry exclusions.
55pub fn atomic_install_dir_filtered(
56    src: &Path,
57    dest: &Path,
58    excluded_top_level: &[&str],
59) -> Result<(), MarsError> {
60    atomic_install_dir_impl(src, dest, excluded_top_level)
61}
62
63fn atomic_install_dir_impl(
64    src: &Path,
65    dest: &Path,
66    excluded_top_level: &[&str],
67) -> Result<(), MarsError> {
68    let parent = dest.parent().unwrap_or(Path::new("."));
69    fs::create_dir_all(parent)?;
70
71    let tmp_dir = tempfile::TempDir::new_in(parent)?;
72    copy_dir_recursive(src, tmp_dir.path(), src, excluded_top_level)?;
73    let tmp_path = tmp_dir.keep();
74
75    if dest.exists() {
76        // Step 1: Rename old to .old (old content still accessible)
77        let old_path = parent.join(format!(
78            ".{}.old",
79            dest.file_name().unwrap_or_default().to_string_lossy()
80        ));
81        // Clean up stale .old from a prior crash
82        if old_path.exists() {
83            fs::remove_dir_all(&old_path)?;
84        }
85        // Atomic: old content moves to .old, dest slot is free
86        fs::rename(dest, &old_path)?;
87        // Atomic: new content takes dest slot
88        if let Err(e) = fs::rename(&tmp_path, dest) {
89            // Rollback: move old content back
90            let _ = fs::rename(&old_path, dest);
91            let _ = fs::remove_dir_all(&tmp_path);
92            return Err(e.into());
93        }
94        // Cleanup: remove old content (non-critical)
95        let _ = fs::remove_dir_all(&old_path);
96    } else {
97        fs::rename(&tmp_path, dest)?;
98    }
99
100    Ok(())
101}
102
103/// Recursively copy a directory tree.
104fn copy_dir_recursive(
105    src: &Path,
106    dest: &Path,
107    root: &Path,
108    excluded_top_level: &[&str],
109) -> Result<(), MarsError> {
110    for entry in fs::read_dir(src)? {
111        let entry = entry?;
112        let file_type = entry.file_type()?;
113        let src_path = entry.path();
114        let dest_path = dest.join(entry.file_name());
115
116        let rel_path = src_path
117            .strip_prefix(root)
118            .expect("copy traversal path should be under root");
119        if is_excluded_top_level(rel_path, excluded_top_level) {
120            continue;
121        }
122
123        if file_type.is_dir() {
124            fs::create_dir_all(&dest_path)?;
125            copy_dir_recursive(&src_path, &dest_path, root, excluded_top_level)?;
126        } else {
127            fs::copy(&src_path, &dest_path)?;
128        }
129    }
130    Ok(())
131}
132
133fn is_excluded_top_level(path: &Path, excluded_top_level: &[&str]) -> bool {
134    let Some(first) = path.components().next().map(|c| c.as_os_str()) else {
135        return false;
136    };
137    excluded_top_level.iter().any(|excluded| first == *excluded)
138}
139
140/// Remove a file or directory (skills are dirs).
141pub fn remove_item(path: &Path, kind: ItemKind) -> Result<(), MarsError> {
142    match kind {
143        ItemKind::Agent => fs::remove_file(path)?,
144        ItemKind::Skill => fs::remove_dir_all(path)?,
145    }
146    Ok(())
147}
148
149/// Advisory file lock (flock) for concurrent access.
150///
151/// Prevents concurrent `mars sync` from corrupting state.
152/// The lock is held start-to-end — acquired before fetching and held through completion.
153/// Dropping the `FileLock` closes the fd, which releases the advisory lock.
154pub struct FileLock {
155    _fd: fs::File,
156}
157
158impl FileLock {
159    /// Acquire an advisory file lock, blocking until available.
160    pub fn acquire(lock_path: &Path) -> Result<Self, MarsError> {
161        let file = Self::open_lock_file(lock_path)?;
162        let fd = file.as_raw_fd();
163
164        // SAFETY: fd is a valid open file descriptor
165        let ret = unsafe { libc::flock(fd, libc::LOCK_EX) };
166        if ret != 0 {
167            return Err(std::io::Error::last_os_error().into());
168        }
169
170        Ok(FileLock { _fd: file })
171    }
172
173    /// Try to acquire the lock without blocking.
174    /// Returns `Ok(Some(lock))` if acquired, `Ok(None)` if already held by another process.
175    pub fn try_acquire(lock_path: &Path) -> Result<Option<Self>, MarsError> {
176        let file = Self::open_lock_file(lock_path)?;
177        let fd = file.as_raw_fd();
178
179        // SAFETY: fd is a valid open file descriptor
180        let ret = unsafe { libc::flock(fd, libc::LOCK_EX | libc::LOCK_NB) };
181        if ret != 0 {
182            let err = std::io::Error::last_os_error();
183            if err.kind() == std::io::ErrorKind::WouldBlock {
184                return Ok(None);
185            }
186            return Err(err.into());
187        }
188
189        Ok(Some(FileLock { _fd: file }))
190    }
191
192    /// Open (or create) the lock file, creating parent dirs if needed.
193    fn open_lock_file(lock_path: &Path) -> Result<fs::File, MarsError> {
194        if let Some(parent) = lock_path.parent() {
195            fs::create_dir_all(parent)?;
196        }
197        let file = fs::OpenOptions::new()
198            .read(true)
199            .write(true)
200            .create(true)
201            .truncate(false)
202            .open(lock_path)?;
203        Ok(file)
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use tempfile::TempDir;
211
212    #[test]
213    fn atomic_write_creates_file_with_correct_content() {
214        let dir = TempDir::new().unwrap();
215        let dest = dir.path().join("output.txt");
216        let content = b"hello world";
217
218        atomic_write(&dest, content).unwrap();
219
220        assert_eq!(fs::read(&dest).unwrap(), content);
221    }
222
223    #[test]
224    fn atomic_write_creates_parent_dirs() {
225        let dir = TempDir::new().unwrap();
226        let dest = dir.path().join("nested").join("dir").join("file.txt");
227        let content = b"nested content";
228
229        atomic_write(&dest, content).unwrap();
230
231        assert_eq!(fs::read(&dest).unwrap(), content);
232    }
233
234    #[test]
235    fn atomic_write_overwrites_existing_file() {
236        let dir = TempDir::new().unwrap();
237        let dest = dir.path().join("output.txt");
238
239        atomic_write(&dest, b"first").unwrap();
240        atomic_write(&dest, b"second").unwrap();
241
242        assert_eq!(fs::read(&dest).unwrap(), b"second");
243    }
244
245    #[test]
246    fn atomic_install_dir_copies_tree() {
247        let dir = TempDir::new().unwrap();
248        let src = dir.path().join("src_dir");
249        let dest = dir.path().join("dest_dir");
250
251        // Create source tree
252        fs::create_dir_all(src.join("sub")).unwrap();
253        fs::write(src.join("a.txt"), "file a").unwrap();
254        fs::write(src.join("sub").join("b.txt"), "file b").unwrap();
255
256        atomic_install_dir(&src, &dest).unwrap();
257
258        assert_eq!(fs::read_to_string(dest.join("a.txt")).unwrap(), "file a");
259        assert_eq!(
260            fs::read_to_string(dest.join("sub").join("b.txt")).unwrap(),
261            "file b"
262        );
263    }
264
265    #[test]
266    fn atomic_install_dir_replaces_existing() {
267        let dir = TempDir::new().unwrap();
268        let src = dir.path().join("src_dir");
269        let dest = dir.path().join("dest_dir");
270
271        // Create initial dest
272        fs::create_dir_all(&dest).unwrap();
273        fs::write(dest.join("old.txt"), "old").unwrap();
274
275        // Create source
276        fs::create_dir_all(&src).unwrap();
277        fs::write(src.join("new.txt"), "new").unwrap();
278
279        atomic_install_dir(&src, &dest).unwrap();
280
281        assert!(dest.join("new.txt").exists());
282        assert!(!dest.join("old.txt").exists());
283    }
284
285    #[test]
286    fn atomic_install_dir_cleans_stale_old() {
287        let dir = TempDir::new().unwrap();
288        let src = dir.path().join("src_dir");
289        let dest = dir.path().join("dest_dir");
290
291        // Create initial dest
292        fs::create_dir_all(&dest).unwrap();
293        fs::write(dest.join("old.txt"), "old").unwrap();
294
295        // Create stale .old from a prior crash
296        let old_path = dir.path().join(".dest_dir.old");
297        fs::create_dir_all(&old_path).unwrap();
298        fs::write(old_path.join("stale.txt"), "stale").unwrap();
299
300        // Create source
301        fs::create_dir_all(&src).unwrap();
302        fs::write(src.join("new.txt"), "new").unwrap();
303
304        atomic_install_dir(&src, &dest).unwrap();
305
306        assert!(dest.join("new.txt").exists());
307        assert!(!dest.join("old.txt").exists());
308        assert!(!old_path.exists(), "stale .old should be cleaned up");
309    }
310
311    #[test]
312    fn atomic_install_dir_dest_exists_throughout() {
313        let dir = TempDir::new().unwrap();
314        let src = dir.path().join("src_dir");
315        let dest = dir.path().join("dest_dir");
316
317        // Create initial dest
318        fs::create_dir_all(&dest).unwrap();
319        fs::write(dest.join("v1.txt"), "v1").unwrap();
320
321        // Create source
322        fs::create_dir_all(&src).unwrap();
323        fs::write(src.join("v2.txt"), "v2").unwrap();
324
325        assert!(dest.exists(), "dest should exist before install");
326        atomic_install_dir(&src, &dest).unwrap();
327        assert!(dest.exists(), "dest should exist after install");
328        assert!(dest.join("v2.txt").exists());
329    }
330
331    #[test]
332    fn atomic_install_dir_filtered_excludes_top_level_entries() {
333        let dir = TempDir::new().unwrap();
334        let src = dir.path().join("src_dir");
335        let dest = dir.path().join("dest_dir");
336
337        fs::create_dir_all(src.join(".git")).unwrap();
338        fs::create_dir_all(src.join("resources")).unwrap();
339        fs::write(src.join("SKILL.md"), "skill").unwrap();
340        fs::write(src.join("mars.toml"), "ignored").unwrap();
341        fs::write(src.join(".gitignore"), "ignored").unwrap();
342        fs::write(src.join(".git").join("config"), "ignored").unwrap();
343        fs::write(src.join("resources").join("guide.md"), "kept").unwrap();
344
345        atomic_install_dir_filtered(&src, &dest, FLAT_SKILL_EXCLUDED_TOP_LEVEL).unwrap();
346
347        assert!(dest.join("SKILL.md").exists());
348        assert!(dest.join("resources").join("guide.md").exists());
349        assert!(!dest.join(".git").exists());
350        assert!(!dest.join("mars.toml").exists());
351        assert!(!dest.join(".gitignore").exists());
352    }
353
354    #[test]
355    fn remove_item_removes_file() {
356        let dir = TempDir::new().unwrap();
357        let file = dir.path().join("agent.md");
358        fs::write(&file, "agent content").unwrap();
359
360        remove_item(&file, ItemKind::Agent).unwrap();
361
362        assert!(!file.exists());
363    }
364
365    #[test]
366    fn remove_item_removes_directory() {
367        let dir = TempDir::new().unwrap();
368        let skill_dir = dir.path().join("my-skill");
369        fs::create_dir_all(skill_dir.join("sub")).unwrap();
370        fs::write(skill_dir.join("main.md"), "skill").unwrap();
371        fs::write(skill_dir.join("sub").join("helper.md"), "helper").unwrap();
372
373        remove_item(&skill_dir, ItemKind::Skill).unwrap();
374
375        assert!(!skill_dir.exists());
376    }
377
378    #[test]
379    fn file_lock_acquire_returns_lock() {
380        let dir = TempDir::new().unwrap();
381        let lock_path = dir.path().join("test.lock");
382
383        let lock = FileLock::acquire(&lock_path).unwrap();
384        assert!(lock_path.exists());
385        drop(lock);
386    }
387
388    #[test]
389    fn file_lock_released_on_drop() {
390        let dir = TempDir::new().unwrap();
391        let lock_path = dir.path().join("test.lock");
392
393        {
394            let _lock = FileLock::acquire(&lock_path).unwrap();
395            // Lock held here
396        }
397        // Lock dropped — should be acquirable again
398        let lock2 = FileLock::try_acquire(&lock_path).unwrap();
399        assert!(lock2.is_some());
400    }
401
402    #[test]
403    fn file_lock_creates_parent_dirs() {
404        let dir = TempDir::new().unwrap();
405        let lock_path = dir.path().join("nested").join("dir").join("test.lock");
406
407        let lock = FileLock::acquire(&lock_path).unwrap();
408        assert!(lock_path.exists());
409        drop(lock);
410    }
411}