fix(cli-core): use safe-write for atomic replace, upgrade to 0.2.0

dstack/Cargo.toml

@@ -143,7 +143,7 @@ rand = "0.8.5"
 regorus = { version = "0.10.1", default-features = false, features = ["full-opa", "arc"] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
-safe-write = "0.1.3"
+safe-write = "0.2.0"
 rustix = { version = "0.38", features = ["fs"] }
 nix = "0.29.0"
 # Vendored: 6.0.2 does not build for musl. See vendor/README.md.

dstack/crates/dstack-cli-core/Cargo.toml

@@ -14,6 +14,7 @@ http-client = { workspace = true, features = ["prpc"] }
 dstack-vmm-rpc.workspace = true
 dstack-types.workspace = true
 serde_json.workspace = true
+safe-write.workspace = true
 # advisory file locking (flock) for the allowlist/state read-modify-write;
 # already in the dependency tree transitively, so no extra compile cost.
 rustix = { version = "0.38", features = ["fs"] }

dstack/crates/dstack-cli-core/src/fsutil.rs

@@ -10,73 +10,58 @@
 //! auth webhook fails *closed* on invalid JSON, so a half-written allowlist
 //! denies keys to every app on the host. These helpers make the write atomic
 //! and serialize concurrent writers.
+//!
+//! The atomic replace itself is delegated to the `safe-write` crate, which the
+//! rest of the workspace already uses. The two concerns stay separate:
+//! `safe-write` makes a single write atomic and durable, while
+//! [`lock_exclusive`] is what serializes a *read*-modify-write so two processes
+//! cannot each publish a complete file built from the same stale read.
 
 use anyhow::{Context, Result};
 use std::ffi::OsString;
 use std::fs::{File, OpenOptions};
-use std::io::Write;
 use std::path::{Path, PathBuf};
 
 /// `path` with `suffix` appended to its full name (not replacing the extension,
-/// so `a/b.json` + `.tmp` → `a/b.json.tmp`, a sibling in the same directory).
+/// so `a/b.json` + `.lock` → `a/b.json.lock`, a sibling in the same directory).
 fn sibling(path: &Path, suffix: &str) -> PathBuf {
     let mut s: OsString = path.as_os_str().to_os_string();
     s.push(suffix);
     PathBuf::from(s)
 }
 
-/// atomically replace `path`'s contents: write a sibling temp file, fsync it,
-/// rename it over the target, then fsync the directory. A reader (or a crash)
-/// sees either the old file or the new one, never a fragment, and the rename is
-/// durable across a power loss. `tmp` and `path` are in the same directory so
-/// the rename is atomic.
+/// atomically replace `path`'s contents: write a uniquely named temp file in
+/// the same directory, fsync it, rename it over the target, then fsync the
+/// directory. A reader (or a crash) sees either the old file or the new one,
+/// never a fragment, and the rename is durable across a power loss.
+///
+/// Concurrent writers to the same path all succeed and each publishes its
+/// complete content; the winner is whoever renames last. That is *not* a
+/// substitute for [`lock_exclusive`] when the write is part of a
+/// read-modify-write — see the module docs.
 pub fn write_atomic(path: &Path, contents: &str) -> Result<()> {
-    write_atomic_inner(path, contents, None)
+    safe_write::safe_write(path, contents).with_context(|| format!("writing {}", path.display()))
 }
 
-/// like [`write_atomic`], but the temp file is created with `mode` (Unix
-/// permission bits) *before* any content is written, so a secret never exists
-/// on disk with broader-than-intended permissions — not even transiently
-/// between the rename and a follow-up `chmod`. Use for credential files
-/// (`0o600`). The final file keeps `mode` because `rename` preserves it.
+/// like [`write_atomic`], but the file is created with `mode` (Unix permission
+/// bits) *before* any content is written, so a secret never exists on disk with
+/// broader-than-intended permissions — not even transiently between the rename
+/// and a follow-up `chmod`. Use for credential files (`0o600`).
+///
+/// `mode` is subject to the process umask, matching
+/// [`std::os::unix::fs::OpenOptionsExt::mode`], so the result is never wider
+/// than requested. On non-Unix platforms `mode` is ignored, as before.
 pub fn write_atomic_mode(path: &Path, contents: &str, mode: u32) -> Result<()> {
-    write_atomic_inner(path, contents, Some(mode))
-}
-
-fn write_atomic_inner(path: &Path, contents: &str, mode: Option<u32>) -> Result<()> {
-    let tmp = sibling(path, ".tmp");
-    let mut opts = OpenOptions::new();
-    opts.write(true).create(true).truncate(true);
-    #[cfg(unix)]
-    if let Some(mode) = mode {
-        use std::os::unix::fs::OpenOptionsExt;
-        opts.mode(mode);
-    }
-    let mut f = opts
-        .open(&tmp)
-        .with_context(|| format!("creating temp file {}", tmp.display()))?;
-    // if a stale temp file survived a crash, `create` reused it without
-    // resetting its mode; tighten it before writing the secret.
     #[cfg(unix)]
-    if let Some(mode) = mode {
-        use std::os::unix::fs::PermissionsExt;
-        std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(mode))
-            .with_context(|| format!("setting mode on {}", tmp.display()))?;
+    {
+        safe_write::safe_write_with_mode(path, contents, mode)
+            .with_context(|| format!("writing {}", path.display()))
     }
-    f.write_all(contents.as_bytes())
-        .with_context(|| format!("writing {}", tmp.display()))?;
-    f.sync_all()
-        .with_context(|| format!("syncing {}", tmp.display()))?;
-    drop(f);
-    std::fs::rename(&tmp, path)
-        .with_context(|| format!("renaming {} -> {}", tmp.display(), path.display()))?;
-    // fsync the containing directory so the rename itself survives a crash.
-    if let Some(dir) = path.parent().filter(|d| !d.as_os_str().is_empty()) {
-        if let Ok(d) = File::open(dir) {
-            let _ = d.sync_all();
-        }
+    #[cfg(not(unix))]
+    {
+        let _ = mode;
+        write_atomic(path, contents)
     }
-    Ok(())
 }
 
 /// acquire an exclusive advisory lock tied to `path` (held on a sibling
@@ -101,17 +86,77 @@ pub fn lock_exclusive(path: &Path) -> Result<File> {
 mod tests {
     use super::*;
 
+    /// every entry in `dir`, so "no temp file left behind" can be asserted
+    /// without knowing the temp file's name.
+    fn entries(dir: &Path) -> Vec<String> {
+        let mut v: Vec<String> = std::fs::read_dir(dir)
+            .unwrap()
+            .map(|e| e.unwrap().file_name().to_string_lossy().into_owned())
+            .collect();
+        v.sort();
+        v
+    }
+
     #[test]
     fn atomic_write_replaces_contents() {
         let dir = std::env::temp_dir().join(format!("dstack-fsutil-{}", std::process::id()));
+        let _ = std::fs::remove_dir_all(&dir);
         std::fs::create_dir_all(&dir).unwrap();
         let p = dir.join("x.json");
         write_atomic(&p, "one").unwrap();
         assert_eq!(std::fs::read_to_string(&p).unwrap(), "one");
         write_atomic(&p, "two").unwrap();
         assert_eq!(std::fs::read_to_string(&p).unwrap(), "two");
-        // no temp file left behind.
-        assert!(!sibling(&p, ".tmp").exists());
+        // no temp file left behind, whatever it was called.
+        assert_eq!(entries(&dir), vec!["x.json".to_string()]);
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+
+    /// Two writers racing on one path must both succeed, and the result must be
+    /// exactly one of them — not a mixture, and not a spurious failure. The
+    /// previous hand-rolled implementation used a fixed `<path>.tmp`, so
+    /// writers truncated each other's temp file and 3 in 4 writes failed.
+    #[test]
+    fn concurrent_writers_do_not_clobber_each_other() {
+        const LEN: usize = 256 * 1024;
+        let dir = std::env::temp_dir().join(format!("dstack-fsrace-{}", std::process::id()));
+        let _ = std::fs::remove_dir_all(&dir);
+        std::fs::create_dir_all(&dir).unwrap();
+        let p = dir.join("state.json");
+
+        for _ in 0..10 {
+            std::thread::scope(|s| {
+                for c in ['a', 'b', 'c', 'd'] {
+                    let p = &p;
+                    s.spawn(move || {
+                        let body: String = std::iter::repeat_n(c, LEN).collect();
+                        write_atomic(p, &body).expect("concurrent write must not fail");
+                    });
+                }
+            });
+            let got = std::fs::read_to_string(&p).unwrap();
+            assert_eq!(got.len(), LEN, "torn write: wrong length");
+            let distinct: std::collections::BTreeSet<char> = got.chars().collect();
+            assert_eq!(distinct.len(), 1, "torn write: mixed two writers' content");
+        }
+        assert_eq!(entries(&dir), vec!["state.json".to_string()]);
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+
+    /// Rewriting a credential file through the plain helper must not widen its
+    /// permissions. The previous implementation reset them to `0o666 & !umask`.
+    #[cfg(unix)]
+    #[test]
+    fn rewrite_preserves_existing_permissions() {
+        use std::os::unix::fs::PermissionsExt;
+        let dir = std::env::temp_dir().join(format!("dstack-fsperm-{}", std::process::id()));
+        let _ = std::fs::remove_dir_all(&dir);
+        std::fs::create_dir_all(&dir).unwrap();
+        let p = dir.join("token");
+        write_atomic_mode(&p, "secret", 0o600).unwrap();
+        write_atomic(&p, "rotated").unwrap();
+        let mode = std::fs::metadata(&p).unwrap().permissions().mode() & 0o777;
+        assert_eq!(mode, 0o600, "rewrite widened a credential file to {mode:o}");
         let _ = std::fs::remove_dir_all(&dir);
     }