Skip to main content

purple_ssh/
fs_util.rs

1use std::fs;
2use std::io;
3use std::path::{Path, PathBuf};
4
5/// Advisory file lock using a `.lock` file.
6/// The lock is released when the `FileLock` is dropped.
7pub struct FileLock {
8    lock_path: PathBuf,
9    #[cfg(unix)]
10    _file: fs::File,
11}
12
13impl FileLock {
14    /// Acquire an advisory lock for the given path.
15    /// Creates a `.purple_lock` file alongside the target and holds an `flock` on it.
16    /// Blocks until the lock is acquired (or returns an error on failure).
17    pub fn acquire(path: &Path) -> io::Result<Self> {
18        let mut lock_name = path
19            .file_name()
20            .unwrap_or_default()
21            .to_os_string();
22        lock_name.push(".purple_lock");
23        let lock_path = path.with_file_name(lock_name);
24
25        #[cfg(unix)]
26        {
27            use std::os::unix::fs::OpenOptionsExt;
28            let file = fs::OpenOptions::new()
29                .write(true)
30                .create(true)
31                .truncate(false)
32                .mode(0o600)
33                .open(&lock_path)?;
34
35            // LOCK_EX = exclusive, blocks until acquired
36            let ret = unsafe { libc::flock(std::os::unix::io::AsRawFd::as_raw_fd(&file), libc::LOCK_EX) };
37            if ret != 0 {
38                return Err(io::Error::last_os_error());
39            }
40
41            Ok(FileLock { lock_path, _file: file })
42        }
43
44        #[cfg(not(unix))]
45        {
46            // On non-Unix, use a simple lock file (best-effort)
47            let file = fs::OpenOptions::new()
48                .write(true)
49                .create_new(true)
50                .open(&lock_path)
51                .or_else(|_| {
52                    // If it already exists, wait briefly and retry
53                    std::thread::sleep(std::time::Duration::from_millis(100));
54                    fs::remove_file(&lock_path).ok();
55                    fs::OpenOptions::new()
56                        .write(true)
57                        .create_new(true)
58                        .open(&lock_path)
59                })?;
60            Ok(FileLock { lock_path, _file: file })
61        }
62    }
63}
64
65impl Drop for FileLock {
66    fn drop(&mut self) {
67        // On Unix, flock is released when the file descriptor is closed (automatic).
68        // Clean up the lock file.
69        let _ = fs::remove_file(&self.lock_path);
70    }
71}
72
73/// Atomic write: write content to a PID-suffixed temp file with chmod 600, then rename.
74/// Uses O_EXCL (create_new) to prevent symlink attacks on the temp file path.
75/// Cleans up the temp file on failure.
76pub fn atomic_write(path: &Path, content: &[u8]) -> io::Result<()> {
77    // Ensure parent directory exists
78    if let Some(parent) = path.parent() {
79        fs::create_dir_all(parent)?;
80    }
81
82    let mut tmp_name = path
83        .file_name()
84        .unwrap_or_default()
85        .to_os_string();
86    tmp_name.push(format!(".purple_tmp.{}", std::process::id()));
87    let tmp_path = path.with_file_name(tmp_name);
88
89    #[cfg(unix)]
90    {
91        use std::io::Write;
92        use std::os::unix::fs::OpenOptionsExt;
93        // Try O_EXCL first. If a stale tmp file exists from a crashed run, remove
94        // it and retry once. This avoids a TOCTOU gap from removing before creating.
95        let open = || {
96            fs::OpenOptions::new()
97                .write(true)
98                .create_new(true)
99                .mode(0o600)
100                .open(&tmp_path)
101        };
102        let mut file = match open() {
103            Ok(f) => f,
104            Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
105                let _ = fs::remove_file(&tmp_path);
106                open().map_err(|e| {
107                    io::Error::new(
108                        e.kind(),
109                        format!("Failed to create temp file {}: {}", tmp_path.display(), e),
110                    )
111                })?
112            }
113            Err(e) => {
114                return Err(io::Error::new(
115                    e.kind(),
116                    format!("Failed to create temp file {}: {}", tmp_path.display(), e),
117                ));
118            }
119        };
120        if let Err(e) = file.write_all(content) {
121            drop(file);
122            let _ = fs::remove_file(&tmp_path);
123            return Err(e);
124        }
125        if let Err(e) = file.sync_all() {
126            drop(file);
127            let _ = fs::remove_file(&tmp_path);
128            return Err(e);
129        }
130    }
131
132    #[cfg(not(unix))]
133    {
134        if let Err(e) = fs::write(&tmp_path, content) {
135            let _ = fs::remove_file(&tmp_path);
136            return Err(e);
137        }
138        // sync_all via reopen since fs::write doesn't return a File handle
139        match fs::File::open(&tmp_path) {
140            Ok(f) => {
141                if let Err(e) = f.sync_all() {
142                    let _ = fs::remove_file(&tmp_path);
143                    return Err(e);
144                }
145            }
146            Err(e) => {
147                let _ = fs::remove_file(&tmp_path);
148                return Err(e);
149            }
150        }
151    }
152
153    let result = fs::rename(&tmp_path, path);
154    if result.is_err() {
155        let _ = fs::remove_file(&tmp_path);
156    }
157    result
158}