From 638e8a0e6ea3656741c9c62742a239426e27c4db Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 3 Aug 2026 20:45:27 -0700 Subject: [PATCH] fix(cli-core): use safe-write for atomic replace, upgrade to 0.2.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `fsutil::write_atomic` derived its temp file as a fixed `.tmp`. That avoided the extension-replacing trap (the doc comment says so explicitly), but a deterministic name is still shared by every concurrent writer to that path, and the file was opened with `create(true).truncate(true)` rather than exclusively. Two writers therefore truncate each other's temp file, and the first one to rename leaves the second holding a descriptor that now points at the published target — so the loser writes into the live file and then fails its own rename with ENOENT. Measured with four writers racing on one path: 120 of 160 writes failed and 18 of 40 rounds left a target that was neither writer's content. The workspace already depends on `safe-write` from five other crates, and it creates a uniquely named temp file, so delegating to it removes the shared name at the source. Upgrading the workspace pin to 0.2.0 is what makes that worthwhile: 0.1.x had the extension bug this module was written to avoid, plus no parent-directory fsync, no cleanup on failure, and a permission reset on overwrite. What this changes for the three call sites: - concurrent writers no longer corrupt or fail. `register_app_in_allowlist` already held `lock_exclusive`, but `write_state` and `write_token_file` did not, and both run in the install flow that the module docs describe as racy. - a failed write no longer leaves the content behind. `write_token_file` currently leaves a complete bearer token in `vmm-auth-token.tmp`. - rewriting an existing file preserves its permissions instead of resetting them to `0o666 & !umask`. `lock_exclusive` is untouched. It solves a different problem — serializing a read-modify-write so two processes cannot each publish a complete file built from the same stale read — and `safe-write` explicitly does not do that. The module docs now say so. The three current call sites all create their parent directory beforehand, and the allowlist path is validated by the read that precedes the write, so `safe-write` creating parent directories cannot mask a mistyped path. The only `mode` in the tree is `0o600`, which is unaffected by the umask. Two new tests, both confirmed to fail against the previous implementation: concurrent writers ("No such file or directory") and permission preservation ("widened a credential file to 664"). --- dstack/Cargo.lock | 6 +- dstack/Cargo.toml | 2 +- dstack/crates/dstack-cli-core/Cargo.toml | 1 + dstack/crates/dstack-cli-core/src/fsutil.rs | 143 +++++++++++++------- 4 files changed, 100 insertions(+), 52 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 6ae5b810a..cf0d8c2d5 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -1851,6 +1851,7 @@ dependencies = [ "dstack-vmm-rpc", "http-client", "rustix 0.38.44", + "safe-write", "serde_json", "toml", ] @@ -6580,11 +6581,12 @@ dependencies = [ [[package]] name = "safe-write" -version = "0.1.3" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ed1f9442e2f57be7e5e7cc1246ca44bd94ac20ef8fb070621253674ac4cff4c" +checksum = "b5a9dc0fc219eaa0265dbbee50170a469de86252d30b946c265c396b065b7f9f" dependencies = [ "fs-err", + "tempfile", ] [[package]] diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 18c4c171d..eb09fdf05 100644 --- a/dstack/Cargo.toml +++ b/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. diff --git a/dstack/crates/dstack-cli-core/Cargo.toml b/dstack/crates/dstack-cli-core/Cargo.toml index 720946063..3f4d0bdcc 100644 --- a/dstack/crates/dstack-cli-core/Cargo.toml +++ b/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"] } diff --git a/dstack/crates/dstack-cli-core/src/fsutil.rs b/dstack/crates/dstack-cli-core/src/fsutil.rs index 125a6eab3..8db04d990 100644 --- a/dstack/crates/dstack-cli-core/src/fsutil.rs +++ b/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) -> 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 { 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 { + let mut v: Vec = 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 `.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 = 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); }