moli/moli-browser-profile/src/atomic_write.rs

use std::{
    ffi::OsString,
    fs::{self, OpenOptions},
    io::Write,
    path::{Path, PathBuf},
    sync::atomic::{AtomicU64, Ordering},
    time::{SystemTime, UNIX_EPOCH},
};

use anyhow::{Context, Result};

static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);

/// Atomically replaces a profile file and makes both its bytes and directory
/// entry durable before reporting success.
pub fn write_file_atomically(path: &Path, bytes: &[u8], label: &str) -> Result<()> {
    if let Some(parent) = path.parent()
        && !parent.as_os_str().is_empty()
    {
        fs::create_dir_all(parent)
            .with_context(|| format!("failed to create {label} dir `{}`", parent.display()))?;
    }

    let tmp_path = unique_temp_path(path);
    let result = (|| {
        let mut tmp = OpenOptions::new()
            .create_new(true)
            .write(true)
            .open(&tmp_path)
            .with_context(|| format!("failed to create {label} `{}`", tmp_path.display()))?;
        tmp.write_all(bytes)
            .with_context(|| format!("failed to write {label} `{}`", tmp_path.display()))?;
        tmp.sync_all()
            .with_context(|| format!("failed to sync {label} `{}`", tmp_path.display()))?;
        drop(tmp);
        fs::rename(&tmp_path, path).with_context(|| {
            format!(
                "failed to replace {label} `{}` from `{}`",
                path.display(),
                tmp_path.display()
            )
        })?;
        sync_parent_directory(path, label)
    })();
    if result.is_err() {
        let _ = fs::remove_file(&tmp_path);
    }
    result
}

fn sync_parent_directory(path: &Path, label: &str) -> Result<()> {
    let parent = path
        .parent()
        .filter(|parent| !parent.as_os_str().is_empty())
        .unwrap_or_else(|| Path::new("."));
    fs::File::open(parent)
        .and_then(|directory| directory.sync_all())
        .with_context(|| {
            format!(
                "failed to sync {label} parent directory `{}`",
                parent.display()
            )
        })
}

fn unique_temp_path(path: &Path) -> PathBuf {
    let nonce = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_nanos())
        .unwrap_or(0);
    let counter = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
    let file_name = path
        .file_name()
        .map(OsString::from)
        .unwrap_or_else(|| OsString::from("profile-write"));
    let mut tmp_name = OsString::from(".");
    tmp_name.push(file_name);
    tmp_name.push(format!(".{}.{}.{}.tmp", std::process::id(), nonce, counter));
    path.with_file_name(tmp_name)
}

#[cfg(test)]
mod tests {
    use std::{
        fs,
        path::PathBuf,
        time::{SystemTime, UNIX_EPOCH},
    };

    use anyhow::Result;

    use super::write_file_atomically;

    struct TempDir {
        path: PathBuf,
    }

    impl TempDir {
        fn new(name: &str) -> Self {
            let nonce = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .expect("system clock should be after epoch")
                .as_nanos();
            let path = std::env::temp_dir().join(format!(
                "moli-atomic-write-{name}-{}-{nonce}",
                std::process::id()
            ));
            Self { path }
        }
    }

    impl Drop for TempDir {
        fn drop(&mut self) {
            let _ = fs::remove_dir_all(&self.path);
        }
    }

    #[test]
    fn atomic_write_uses_unique_temp_name_instead_of_fixed_tmp_path() -> Result<()> {
        let temp = TempDir::new("unique");
        let target = temp.path.join("profile.json");
        let mut fixed_tmp = target.as_os_str().to_owned();
        fixed_tmp.push(".tmp");
        let fixed_tmp = PathBuf::from(fixed_tmp);
        fs::create_dir_all(&temp.path)?;
        fs::write(&target, b"old profile")?;
        fs::write(&fixed_tmp, b"stale fixed tmp")?;

        write_file_atomically(&target, b"new profile", "profile test")?;

        assert_eq!(fs::read(&target)?, b"new profile");
        assert_eq!(fs::read(&fixed_tmp)?, b"stale fixed tmp");
        Ok(())
    }

    #[test]
    fn atomic_write_removes_unique_temp_after_replace_error() -> Result<()> {
        let temp = TempDir::new("replace-error");
        let target = temp.path.join("profile.json");
        fs::create_dir_all(&target)?;

        assert!(write_file_atomically(&target, b"new profile", "profile test").is_err());

        let generated_temps = fs::read_dir(&temp.path)?
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.file_type().is_ok_and(|file_type| file_type.is_file()))
            .count();
        assert_eq!(generated_temps, 0);
        Ok(())
    }
}