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
9pub 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
19pub fn atomic_write(dest: &Path, content: &[u8]) -> Result<(), MarsError> {
24 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
43pub fn atomic_install_dir(src: &Path, dest: &Path) -> Result<(), MarsError> {
51 atomic_install_dir_impl(src, dest, &[])
52}
53
54pub 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 let old_path = parent.join(format!(
78 ".{}.old",
79 dest.file_name().unwrap_or_default().to_string_lossy()
80 ));
81 if old_path.exists() {
83 fs::remove_dir_all(&old_path)?;
84 }
85 fs::rename(dest, &old_path)?;
87 if let Err(e) = fs::rename(&tmp_path, dest) {
89 let _ = fs::rename(&old_path, dest);
91 let _ = fs::remove_dir_all(&tmp_path);
92 return Err(e.into());
93 }
94 let _ = fs::remove_dir_all(&old_path);
96 } else {
97 fs::rename(&tmp_path, dest)?;
98 }
99
100 Ok(())
101}
102
103fn 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
140pub 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
149pub struct FileLock {
155 _fd: fs::File,
156}
157
158impl FileLock {
159 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 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 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 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 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 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 fs::create_dir_all(&dest).unwrap();
273 fs::write(dest.join("old.txt"), "old").unwrap();
274
275 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 fs::create_dir_all(&dest).unwrap();
293 fs::write(dest.join("old.txt"), "old").unwrap();
294
295 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 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 fs::create_dir_all(&dest).unwrap();
319 fs::write(dest.join("v1.txt"), "v1").unwrap();
320
321 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 }
397 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}