Guard Temporary Directories with RAII Cleanup

use std::fs;
use std::io;
use std::path::{Component, Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};

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

struct TempDir {
    path: PathBuf,
}

impl TempDir {
    fn new(prefix: &str) -> io::Result<Self> {
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
        let path = std::env::temp_dir().join(format!("{prefix}-{}-{id}", std::process::id()));
        fs::create_dir(&path)?;
        Ok(Self { path })
    }

    fn path(&self) -> &Path {
        &self.path
    }

    fn write(&self, name: &str, contents: &str) -> io::Result<PathBuf> {
        let relative = Path::new(name);
        let escapes = relative.components().any(|component| {
            matches!(
                component,
                Component::ParentDir | Component::Prefix(_) | Component::RootDir
            )
        });
        if relative.is_absolute() || escapes {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "path escapes the temporary directory",
            ));
        }

        let target = self.path.join(relative);
        if let Some(parent) = target.parent() {
            fs::create_dir_all(parent)?;
        }
        fs::write(&target, contents)?;
        Ok(target)
    }
}

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

fn main() -> io::Result<()> {
    let dir = TempDir::new("devexamples")?;
    let file = dir.write("nested/config.txt", "ok")?;
    println!("wrote {}", file.display());
    assert_eq!(fs::read_to_string(&file)?, "ok");
    assert!(dir.write("../escape.txt", "no").is_err());

    let path = dir.path().to_path_buf();
    drop(dir);
    assert!(!path.exists(), "temporary directory should be removed");

    let second = TempDir::new("devexamples")?;
    assert!(second.path().exists());
    Ok(())
}