1use std::fs;
2use std::io::Write;
3use std::path::Path;
4
5use crate::error::MarsError;
6use crate::types::ItemKind;
7
8pub const FLAT_SKILL_EXCLUDED_TOP_LEVEL: &[&str] = &[
10 ".git",
11 ".mars",
12 "mars.toml",
13 "mars.lock",
14 "mars.local.toml",
15 ".gitignore",
16];
17
18pub fn atomic_write(dest: &Path, content: &[u8]) -> Result<(), MarsError> {
23 if let Some(parent) = dest.parent() {
25 fs::create_dir_all(parent)?;
26 }
27
28 let parent = dest.parent().unwrap_or(Path::new("."));
29 let mut tmp = tempfile::NamedTempFile::new_in(parent)?;
30 tmp.write_all(content)?;
31 tmp.as_file().sync_all()?;
32 #[cfg(unix)]
33 {
34 use std::os::unix::fs::PermissionsExt;
35 tmp.as_file()
36 .set_permissions(fs::Permissions::from_mode(0o644))?;
37 }
38 tmp.persist(dest).map_err(|e| e.error)?;
39 Ok(())
40}
41
42pub fn atomic_install_dir(src: &Path, dest: &Path) -> Result<(), MarsError> {
50 atomic_install_dir_impl(src, dest, &[])
51}
52
53pub fn atomic_install_dir_filtered(
55 src: &Path,
56 dest: &Path,
57 excluded_top_level: &[&str],
58) -> Result<(), MarsError> {
59 atomic_install_dir_impl(src, dest, excluded_top_level)
60}
61
62fn atomic_install_dir_impl(
63 src: &Path,
64 dest: &Path,
65 excluded_top_level: &[&str],
66) -> Result<(), MarsError> {
67 let parent = dest.parent().unwrap_or(Path::new("."));
68 fs::create_dir_all(parent)?;
69
70 let tmp_dir = tempfile::TempDir::new_in(parent)?;
71 copy_dir_recursive(src, tmp_dir.path(), src, excluded_top_level)?;
72 let tmp_path = tmp_dir.keep();
73
74 if dest.exists() {
75 let old_path = parent.join(format!(
77 ".{}.old",
78 dest.file_name().unwrap_or_default().to_string_lossy()
79 ));
80 if old_path.exists() {
82 fs::remove_dir_all(&old_path)?;
83 }
84 fs::rename(dest, &old_path)?;
86 if let Err(e) = fs::rename(&tmp_path, dest) {
88 let _ = fs::rename(&old_path, dest);
90 let _ = fs::remove_dir_all(&tmp_path);
91 return Err(e.into());
92 }
93 let _ = fs::remove_dir_all(&old_path);
95 } else {
96 fs::rename(&tmp_path, dest)?;
97 }
98
99 Ok(())
100}
101
102fn copy_dir_recursive(
104 src: &Path,
105 dest: &Path,
106 root: &Path,
107 excluded_top_level: &[&str],
108) -> Result<(), MarsError> {
109 for entry in fs::read_dir(src)? {
110 let entry = entry?;
111 let file_type = entry.file_type()?;
112 let src_path = entry.path();
113 let dest_path = dest.join(entry.file_name());
114
115 let rel_path = src_path
116 .strip_prefix(root)
117 .expect("copy traversal path should be under root");
118 if is_excluded_top_level(rel_path, excluded_top_level) {
119 continue;
120 }
121
122 if file_type.is_dir() {
123 fs::create_dir_all(&dest_path)?;
124 copy_dir_recursive(&src_path, &dest_path, root, excluded_top_level)?;
125 } else {
126 fs::copy(&src_path, &dest_path)?;
127 }
128 }
129 Ok(())
130}
131
132fn is_excluded_top_level(path: &Path, excluded_top_level: &[&str]) -> bool {
133 let Some(first) = path.components().next().map(|c| c.as_os_str()) else {
134 return false;
135 };
136 excluded_top_level.iter().any(|excluded| first == *excluded)
137}
138
139pub fn remove_item(path: &Path, kind: ItemKind) -> Result<(), MarsError> {
141 match kind {
142 ItemKind::Agent => fs::remove_file(path)?,
143 ItemKind::Skill => fs::remove_dir_all(path)?,
144 }
145 Ok(())
146}
147
148#[cfg(windows)]
149#[allow(clippy::permissions_set_readonly_false)]
150pub fn clear_readonly(path: &Path) -> std::io::Result<()> {
151 if let Ok(metadata) = std::fs::metadata(path) {
152 let mut perms = metadata.permissions();
153 if perms.readonly() {
154 perms.set_readonly(false);
155 std::fs::set_permissions(path, perms)?;
156 }
157 }
158 Ok(())
159}
160
161pub struct FileLock {
167 _fd: fs::File,
168}
169
170impl FileLock {
171 pub fn acquire(lock_path: &Path) -> Result<Self, MarsError> {
173 let file = Self::open_lock_file(lock_path)?;
174 platform::lock_exclusive(&file)?;
175 Ok(FileLock { _fd: file })
176 }
177
178 pub fn try_acquire(lock_path: &Path) -> Result<Option<Self>, MarsError> {
181 let file = Self::open_lock_file(lock_path)?;
182 match platform::try_lock_exclusive(&file) {
183 Ok(true) => Ok(Some(FileLock { _fd: file })),
184 Ok(false) => Ok(None),
185 Err(err) => Err(err.into()),
186 }
187 }
188
189 fn open_lock_file(lock_path: &Path) -> Result<fs::File, MarsError> {
191 if let Some(parent) = lock_path.parent() {
192 fs::create_dir_all(parent)?;
193 }
194 let file = fs::OpenOptions::new()
195 .read(true)
196 .write(true)
197 .create(true)
198 .truncate(false)
199 .open(lock_path)?;
200 Ok(file)
201 }
202}
203
204#[cfg(unix)]
205mod platform {
206 use std::fs;
207 use std::os::unix::io::AsRawFd;
208
209 pub fn lock_exclusive(file: &fs::File) -> std::io::Result<()> {
210 let ret = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX) };
212 if ret != 0 {
213 Err(std::io::Error::last_os_error())
214 } else {
215 Ok(())
216 }
217 }
218
219 pub fn try_lock_exclusive(file: &fs::File) -> std::io::Result<bool> {
220 let ret = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
222 if ret != 0 {
223 let err = std::io::Error::last_os_error();
224 if err.kind() == std::io::ErrorKind::WouldBlock {
225 Ok(false)
226 } else {
227 Err(err)
228 }
229 } else {
230 Ok(true)
231 }
232 }
233}
234
235#[cfg(windows)]
236mod platform {
237 use std::fs;
238 use std::os::windows::io::AsRawHandle;
239
240 use windows_sys::Win32::Foundation::HANDLE;
241 use windows_sys::Win32::Storage::FileSystem::{
242 LOCKFILE_EXCLUSIVE_LOCK, LOCKFILE_FAIL_IMMEDIATELY, LockFileEx,
243 };
244
245 const ERROR_LOCK_VIOLATION: i32 = 33;
246
247 pub fn lock_exclusive(file: &fs::File) -> std::io::Result<()> {
248 let handle = file.as_raw_handle() as HANDLE;
249 let mut overlapped = unsafe { std::mem::zeroed() };
252 let ret =
254 unsafe { LockFileEx(handle, LOCKFILE_EXCLUSIVE_LOCK, 0, !0, !0, &mut overlapped) };
255 if ret == 0 {
256 Err(std::io::Error::last_os_error())
257 } else {
258 Ok(())
259 }
260 }
261
262 pub fn try_lock_exclusive(file: &fs::File) -> std::io::Result<bool> {
263 let handle = file.as_raw_handle() as HANDLE;
264 let mut overlapped = unsafe { std::mem::zeroed() };
267 let ret = unsafe {
269 LockFileEx(
270 handle,
271 LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY,
272 0,
273 !0,
274 !0,
275 &mut overlapped,
276 )
277 };
278 if ret == 0 {
279 let err = std::io::Error::last_os_error();
280 if err.raw_os_error() == Some(ERROR_LOCK_VIOLATION) {
281 Ok(false)
282 } else {
283 Err(err)
284 }
285 } else {
286 Ok(true)
287 }
288 }
289}
290
291#[cfg(test)]
292mod tests {
293 use super::*;
294 use tempfile::TempDir;
295
296 #[test]
297 fn atomic_write_creates_file_with_correct_content() {
298 let dir = TempDir::new().unwrap();
299 let dest = dir.path().join("output.txt");
300 let content = b"hello world";
301
302 atomic_write(&dest, content).unwrap();
303
304 assert_eq!(fs::read(&dest).unwrap(), content);
305 }
306
307 #[test]
308 fn atomic_write_creates_parent_dirs() {
309 let dir = TempDir::new().unwrap();
310 let dest = dir.path().join("nested").join("dir").join("file.txt");
311 let content = b"nested content";
312
313 atomic_write(&dest, content).unwrap();
314
315 assert_eq!(fs::read(&dest).unwrap(), content);
316 }
317
318 #[test]
319 fn atomic_write_overwrites_existing_file() {
320 let dir = TempDir::new().unwrap();
321 let dest = dir.path().join("output.txt");
322
323 atomic_write(&dest, b"first").unwrap();
324 atomic_write(&dest, b"second").unwrap();
325
326 assert_eq!(fs::read(&dest).unwrap(), b"second");
327 }
328
329 #[test]
330 fn atomic_install_dir_copies_tree() {
331 let dir = TempDir::new().unwrap();
332 let src = dir.path().join("src_dir");
333 let dest = dir.path().join("dest_dir");
334
335 fs::create_dir_all(src.join("sub")).unwrap();
337 fs::write(src.join("a.txt"), "file a").unwrap();
338 fs::write(src.join("sub").join("b.txt"), "file b").unwrap();
339
340 atomic_install_dir(&src, &dest).unwrap();
341
342 assert_eq!(fs::read_to_string(dest.join("a.txt")).unwrap(), "file a");
343 assert_eq!(
344 fs::read_to_string(dest.join("sub").join("b.txt")).unwrap(),
345 "file b"
346 );
347 }
348
349 #[test]
350 fn atomic_install_dir_replaces_existing() {
351 let dir = TempDir::new().unwrap();
352 let src = dir.path().join("src_dir");
353 let dest = dir.path().join("dest_dir");
354
355 fs::create_dir_all(&dest).unwrap();
357 fs::write(dest.join("old.txt"), "old").unwrap();
358
359 fs::create_dir_all(&src).unwrap();
361 fs::write(src.join("new.txt"), "new").unwrap();
362
363 atomic_install_dir(&src, &dest).unwrap();
364
365 assert!(dest.join("new.txt").exists());
366 assert!(!dest.join("old.txt").exists());
367 }
368
369 #[test]
370 fn atomic_install_dir_cleans_stale_old() {
371 let dir = TempDir::new().unwrap();
372 let src = dir.path().join("src_dir");
373 let dest = dir.path().join("dest_dir");
374
375 fs::create_dir_all(&dest).unwrap();
377 fs::write(dest.join("old.txt"), "old").unwrap();
378
379 let old_path = dir.path().join(".dest_dir.old");
381 fs::create_dir_all(&old_path).unwrap();
382 fs::write(old_path.join("stale.txt"), "stale").unwrap();
383
384 fs::create_dir_all(&src).unwrap();
386 fs::write(src.join("new.txt"), "new").unwrap();
387
388 atomic_install_dir(&src, &dest).unwrap();
389
390 assert!(dest.join("new.txt").exists());
391 assert!(!dest.join("old.txt").exists());
392 assert!(!old_path.exists(), "stale .old should be cleaned up");
393 }
394
395 #[test]
396 fn atomic_install_dir_dest_exists_throughout() {
397 let dir = TempDir::new().unwrap();
398 let src = dir.path().join("src_dir");
399 let dest = dir.path().join("dest_dir");
400
401 fs::create_dir_all(&dest).unwrap();
403 fs::write(dest.join("v1.txt"), "v1").unwrap();
404
405 fs::create_dir_all(&src).unwrap();
407 fs::write(src.join("v2.txt"), "v2").unwrap();
408
409 assert!(dest.exists(), "dest should exist before install");
410 atomic_install_dir(&src, &dest).unwrap();
411 assert!(dest.exists(), "dest should exist after install");
412 assert!(dest.join("v2.txt").exists());
413 }
414
415 #[test]
416 fn atomic_install_dir_filtered_excludes_top_level_entries() {
417 let dir = TempDir::new().unwrap();
418 let src = dir.path().join("src_dir");
419 let dest = dir.path().join("dest_dir");
420
421 fs::create_dir_all(src.join(".git")).unwrap();
422 fs::create_dir_all(src.join("resources")).unwrap();
423 fs::write(src.join("SKILL.md"), "skill").unwrap();
424 fs::write(src.join("mars.toml"), "ignored").unwrap();
425 fs::write(src.join(".gitignore"), "ignored").unwrap();
426 fs::write(src.join(".git").join("config"), "ignored").unwrap();
427 fs::write(src.join("resources").join("guide.md"), "kept").unwrap();
428
429 atomic_install_dir_filtered(&src, &dest, FLAT_SKILL_EXCLUDED_TOP_LEVEL).unwrap();
430
431 assert!(dest.join("SKILL.md").exists());
432 assert!(dest.join("resources").join("guide.md").exists());
433 assert!(!dest.join(".git").exists());
434 assert!(!dest.join("mars.toml").exists());
435 assert!(!dest.join(".gitignore").exists());
436 }
437
438 #[test]
439 fn remove_item_removes_file() {
440 let dir = TempDir::new().unwrap();
441 let file = dir.path().join("agent.md");
442 fs::write(&file, "agent content").unwrap();
443
444 remove_item(&file, ItemKind::Agent).unwrap();
445
446 assert!(!file.exists());
447 }
448
449 #[test]
450 fn remove_item_removes_directory() {
451 let dir = TempDir::new().unwrap();
452 let skill_dir = dir.path().join("my-skill");
453 fs::create_dir_all(skill_dir.join("sub")).unwrap();
454 fs::write(skill_dir.join("main.md"), "skill").unwrap();
455 fs::write(skill_dir.join("sub").join("helper.md"), "helper").unwrap();
456
457 remove_item(&skill_dir, ItemKind::Skill).unwrap();
458
459 assert!(!skill_dir.exists());
460 }
461
462 #[test]
463 fn file_lock_acquire_returns_lock() {
464 let dir = TempDir::new().unwrap();
465 let lock_path = dir.path().join("test.lock");
466
467 let lock = FileLock::acquire(&lock_path).unwrap();
468 assert!(lock_path.exists());
469 drop(lock);
470 }
471
472 #[test]
473 fn file_lock_released_on_drop() {
474 let dir = TempDir::new().unwrap();
475 let lock_path = dir.path().join("test.lock");
476
477 {
478 let _lock = FileLock::acquire(&lock_path).unwrap();
479 }
481 let lock2 = FileLock::try_acquire(&lock_path).unwrap();
483 assert!(lock2.is_some());
484 }
485
486 #[test]
487 fn file_lock_creates_parent_dirs() {
488 let dir = TempDir::new().unwrap();
489 let lock_path = dir.path().join("nested").join("dir").join("test.lock");
490
491 let lock = FileLock::acquire(&lock_path).unwrap();
492 assert!(lock_path.exists());
493 drop(lock);
494 }
495}