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); }