yapi_ai/auth/
lock.rs

1//! Cross-process safety for shared files: a lock compatible with
2//! `proper-lockfile`, which pi uses for `auth.json`, `settings.json` and
3//! `trust.json`, and an atomic replacement of a file's contents. The lock is a
4//! `<file>.lock` directory, and one whose modification time is older than the
5//! stale period may be taken over. A lock taken with [`FileLock::acquire`]
6//! keeps its modification time fresh while held, as `proper-lockfile` does.
7
8use std::io::Write as _;
9use std::path::{Path, PathBuf};
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::time::{Duration, Instant, SystemTime};
12
13use tokio_util::sync::CancellationToken;
14
15/// pi's `stale` option for `auth.json`.
16const STALE: Duration = Duration::from_secs(30);
17const MAX_DELAY: Duration = Duration::from_secs(1);
18
19/// Why the lock could not be taken.
20#[derive(Debug, thiserror::Error)]
21pub enum LockError {
22    /// The operation was cancelled while waiting.
23    #[error("Login cancelled")]
24    Cancelled,
25    /// Another holder kept the lock for the whole stale period.
26    #[error("Lock file is already being held")]
27    Held,
28    /// Creating the lock directory failed for another reason.
29    #[error("{0}")]
30    Io(#[from] std::io::Error),
31}
32
33/// A held lock. Dropping it releases the lock.
34#[derive(Debug)]
35pub struct FileLock {
36    dir: PathBuf,
37    /// The task that keeps the lock fresh.
38    refresh: Option<tokio::task::AbortHandle>,
39}
40
41impl FileLock {
42    /// Locks `file`, retrying with jittered backoff until the stale period has
43    /// passed. While held, the lock is kept fresh.
44    pub async fn acquire(file: &Path, cancel: &CancellationToken) -> Result<FileLock, LockError> {
45        let dir = lock_dir(file);
46        let deadline = Instant::now() + STALE;
47        let mut retry = 0u32;
48        loop {
49            if cancel.is_cancelled() {
50                return Err(LockError::Cancelled);
51            }
52            if let Some(mut lock) = try_lock(&dir)? {
53                lock.refresh = Some(keep_fresh(dir));
54                return Ok(lock);
55            }
56            let remaining = deadline.saturating_duration_since(Instant::now());
57            if remaining.is_zero() {
58                return Err(LockError::Held);
59            }
60            let base =
61                Duration::from_millis(10u64.saturating_mul(1 << retry.min(10))).min(MAX_DELAY / 2);
62            retry += 1;
63            let delay = (base + base.mul_f64(jitter())).min(remaining);
64            cancel
65                .run_until_cancelled(tokio::time::sleep(delay))
66                .await
67                .ok_or(LockError::Cancelled)?;
68        }
69    }
70
71    /// Locks `file` from synchronous code, as pi's `lockSync` with its retries
72    /// for `settings.json` and `trust.json`: ten attempts 20 ms apart, then
73    /// [`LockError::Held`]. The lock is not kept fresh, so hold it only briefly.
74    pub fn acquire_blocking(file: &Path) -> Result<FileLock, LockError> {
75        let dir = lock_dir(file);
76        for attempt in 1..=10 {
77            if let Some(lock) = try_lock(&dir)? {
78                return Ok(lock);
79            }
80            if attempt < 10 {
81                std::thread::sleep(Duration::from_millis(20));
82            }
83        }
84        Err(LockError::Held)
85    }
86}
87
88/// Touches the lock directory every half stale period, as `proper-lockfile`'s
89/// `update` does, until aborted.
90fn keep_fresh(dir: PathBuf) -> tokio::task::AbortHandle {
91    tokio::spawn(async move {
92        loop {
93            tokio::time::sleep(STALE / 2).await;
94            // A failed touch lets the lock go stale, as if its holder had died.
95            let _ = std::fs::File::open(&dir).and_then(|dir| dir.set_modified(SystemTime::now()));
96        }
97    })
98    .abort_handle()
99}
100
101fn lock_dir(file: &Path) -> PathBuf {
102    let mut name = file.as_os_str().to_owned();
103    name.push(".lock");
104    PathBuf::from(name)
105}
106
107/// Takes the lock, or `None` while another holder keeps it fresh.
108fn try_lock(dir: &Path) -> std::io::Result<Option<FileLock>> {
109    loop {
110        match std::fs::create_dir(dir) {
111            Ok(()) => {
112                return Ok(Some(FileLock {
113                    dir: dir.to_path_buf(),
114                    refresh: None,
115                }));
116            }
117            Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
118                // Another process may take a stale lock over first. Either way, try again.
119                if !is_stale(dir) || std::fs::remove_dir(dir).is_err() {
120                    return Ok(None);
121                }
122            }
123            Err(err) => return Err(err),
124        }
125    }
126}
127
128/// Replaces `path`'s contents so that readers see the old or the new bytes,
129/// never a partial file: writes a temporary file beside it, flushes it to disk
130/// and renames it over `path`. An existing file keeps its permissions, and a
131/// new one gets `create_mode` (Unix, less the umask). A symbolic link
132/// keeps pointing at its target, which receives the contents. On error the
133/// file is unchanged.
134pub fn write_atomic(path: &Path, contents: &[u8], create_mode: u32) -> std::io::Result<()> {
135    static NEXT: AtomicU64 = AtomicU64::new(0);
136    let target = std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf());
137    let mut name = std::ffi::OsString::from(".");
138    name.push(target.file_name().unwrap_or_default());
139    name.push(format!(
140        ".{}.{}.tmp",
141        std::process::id(),
142        NEXT.fetch_add(1, Ordering::Relaxed)
143    ));
144    let temp = target.with_file_name(name);
145    let written = write_temp(&temp, &target, contents, create_mode)
146        .and_then(|()| std::fs::rename(&temp, &target));
147    if written.is_err() {
148        let _ = std::fs::remove_file(&temp);
149    }
150    written
151}
152
153fn write_temp(
154    temp: &Path,
155    target: &Path,
156    contents: &[u8],
157    create_mode: u32,
158) -> std::io::Result<()> {
159    let existing = std::fs::metadata(target)
160        .ok()
161        .map(|meta| meta.permissions());
162    // A file a crashed process with the same id left behind would keep its mode.
163    let _ = std::fs::remove_file(temp);
164    let mut options = std::fs::OpenOptions::new();
165    options.write(true).create_new(true);
166    #[cfg(unix)]
167    std::os::unix::fs::OpenOptionsExt::mode(&mut options, create_mode);
168    #[cfg(not(unix))]
169    let _ = create_mode;
170    let mut file = options.open(temp)?;
171    if let Some(permissions) = existing {
172        file.set_permissions(permissions)?;
173    }
174    file.write_all(contents)?;
175    file.sync_all()
176}
177
178impl Drop for FileLock {
179    fn drop(&mut self) {
180        if let Some(refresh) = &self.refresh {
181            refresh.abort();
182        }
183        let _ = std::fs::remove_dir(&self.dir);
184    }
185}
186
187fn is_stale(dir: &Path) -> bool {
188    std::fs::metadata(dir)
189        .and_then(|metadata| metadata.modified())
190        .ok()
191        .and_then(|modified| SystemTime::now().duration_since(modified).ok())
192        .is_some_and(|age| age > STALE)
193}
194
195/// A value in `[0, 1)` that spreads retries of competing processes.
196fn jitter() -> f64 {
197    // Like pi's `Math.random`, jitter never fails. Without the generator it is 0.
198    let bytes = yapi_types::time::try_random_bytes().unwrap_or_default();
199    f64::from(u16::from_le_bytes(bytes)) / 65536.0
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[tokio::test]
207    async fn excludes_and_releases() {
208        let dir = std::env::temp_dir().join(format!("yapi-lock-{}", std::process::id()));
209        std::fs::create_dir_all(&dir).unwrap();
210        let file = dir.join("auth.json");
211        let cancel = CancellationToken::new();
212        let held = FileLock::acquire(&file, &cancel).await.unwrap();
213        assert!(dir.join("auth.json.lock").is_dir());
214        let waiting = tokio::spawn({
215            let file = file.clone();
216            let cancel = cancel.clone();
217            async move { FileLock::acquire(&file, &cancel).await.map(drop) }
218        });
219        tokio::time::sleep(Duration::from_millis(50)).await;
220        assert!(!waiting.is_finished());
221        drop(held);
222        waiting.await.unwrap().unwrap();
223        assert!(!dir.join("auth.json.lock").exists());
224        std::fs::remove_dir_all(&dir).unwrap();
225    }
226
227    #[tokio::test(start_paused = true)]
228    async fn stale_lock_taken_over() {
229        let dir = std::env::temp_dir().join(format!("yapi-lock-stale-{}", std::process::id()));
230        let lock = dir.join("auth.json.lock");
231        std::fs::create_dir_all(&lock).unwrap();
232        let age = || {
233            let old = SystemTime::now() - STALE - Duration::from_secs(1);
234            std::fs::File::open(&lock)
235                .unwrap()
236                .set_modified(old)
237                .unwrap();
238        };
239        age();
240        let (file, cancel) = (dir.join("auth.json"), CancellationToken::new());
241        let acquire = FileLock::acquire(&file, &cancel);
242        let held = tokio::time::timeout(Duration::from_secs(1), acquire).await;
243        assert!(matches!(held, Ok(Ok(_))));
244        // The new holder keeps its lock fresh.
245        age();
246        assert!(is_stale(&lock));
247        tokio::time::sleep(STALE / 2 + Duration::from_millis(1)).await;
248        assert!(!is_stale(&lock));
249        std::fs::remove_dir_all(&dir).unwrap();
250    }
251
252    #[test]
253    fn blocking_lock_gives_up() {
254        let dir = std::env::temp_dir().join(format!("yapi-lock-sync-{}", std::process::id()));
255        std::fs::create_dir_all(&dir).unwrap();
256        let file = dir.join("settings.json");
257        let held = FileLock::acquire_blocking(&file).unwrap();
258        assert!(matches!(
259            FileLock::acquire_blocking(&file),
260            Err(LockError::Held)
261        ));
262        drop(held);
263        FileLock::acquire_blocking(&file).unwrap();
264        std::fs::remove_dir_all(&dir).unwrap();
265    }
266
267    #[cfg(unix)]
268    #[test]
269    fn atomic_write_mode_links() {
270        use std::os::unix::fs::PermissionsExt as _;
271        let mode = |path: &Path| std::fs::metadata(path).unwrap().permissions().mode() & 0o777;
272        let dir = std::env::temp_dir().join(format!("yapi-atomic-{}", std::process::id()));
273        let _ = std::fs::remove_dir_all(&dir);
274        std::fs::create_dir_all(&dir).unwrap();
275
276        let auth = dir.join("auth.json");
277        write_atomic(&auth, b"{}", 0o600).unwrap();
278        assert_eq!(mode(&auth), 0o600);
279        std::fs::set_permissions(&auth, std::fs::Permissions::from_mode(0o640)).unwrap();
280        write_atomic(&auth, b"{\"a\":1}", 0o600).unwrap();
281        assert_eq!(std::fs::read_to_string(&auth).unwrap(), "{\"a\":1}");
282        assert_eq!(mode(&auth), 0o640);
283
284        let link = dir.join("link.json");
285        std::os::unix::fs::symlink(&auth, &link).unwrap();
286        write_atomic(&link, b"{\"b\":2}", 0o666).unwrap();
287        assert!(std::fs::symlink_metadata(&link).unwrap().is_symlink());
288        assert_eq!(std::fs::read_to_string(&auth).unwrap(), "{\"b\":2}");
289
290        // Without room for the temporary file, the old contents stay.
291        std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o555)).unwrap();
292        let failed = write_atomic(&auth, b"{}", 0o600);
293        std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o755)).unwrap();
294        assert!(failed.is_err());
295        assert_eq!(std::fs::read_to_string(&auth).unwrap(), "{\"b\":2}");
296        let names: Vec<_> = std::fs::read_dir(&dir).unwrap().collect();
297        assert_eq!(names.len(), 2);
298        std::fs::remove_dir_all(&dir).unwrap();
299    }
300}
yapi_ai/auth/lock.rs

yapi_ai/auth/
lock.rs
