1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

use crate::{AutoRemove, ContainingDirectory, ForksafeTempfile, Registration, NEXT_MAP_INDEX, REGISTER};
use std::{io, path::Path};
use tempfile::NamedTempFile;

impl Registration {
    /// Create a registered tempfile at the given `path`, where `path` includes the desired filename.
    ///
    /// Depending on the `directory` configuration, intermediate directories will be created, and depending on `cleanup` empty
    /// intermediate directories will be removed.
    pub fn at_path(
        path: impl AsRef<Path>,
        directory: ContainingDirectory,
        cleanup: AutoRemove,
    ) -> io::Result<Registration> {
        let path = path.as_ref();
        let tempfile = {
            let mut builder = tempfile::Builder::new();
            let dot_ext_storage;
            match path.file_stem() {
                Some(stem) => builder.prefix(stem),
                None => builder.prefix(""),
            };
            if let Some(ext) = path.extension() {
                dot_ext_storage = format!(".{}", ext.to_string_lossy());
                builder.suffix(&dot_ext_storage);
            }
            let parent_dir = path.parent().expect("parent directory is present");
            let parent_dir = directory.resolve(parent_dir)?;
            ForksafeTempfile::new(builder.rand_bytes(0).tempfile_in(parent_dir)?, cleanup)
        };
        let id = NEXT_MAP_INDEX.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        expect_none(REGISTER.insert(id, Some(tempfile)));
        Ok(Registration { id })
    }

    /// Create a registered tempfile within `containing_directory` with a name that won't clash, and clean it up as specified with `cleanup`.
    /// Control how to deal with intermediate directories with `directory`.
    pub fn new(
        containing_directory: impl AsRef<Path>,
        directory: ContainingDirectory,
        cleanup: AutoRemove,
    ) -> io::Result<Registration> {
        let id = NEXT_MAP_INDEX.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let containing_directory = directory.resolve(containing_directory.as_ref())?;
        expect_none(REGISTER.insert(
            id,
            Some(ForksafeTempfile::new(
                NamedTempFile::new_in(containing_directory)?,
                cleanup,
            )),
        ));
        Ok(Registration { id })
    }

    /// Take ownership of the temporary file.
    ///
    pub fn take(self) -> Option<NamedTempFile> {
        let res = REGISTER.remove(&self.id);
        std::mem::forget(self);
        res.and_then(|(_k, v)| v.map(|v| v.inner))
    }
}

impl ContainingDirectory {
    fn resolve(self, dir: &Path) -> std::io::Result<&Path> {
        match self {
            ContainingDirectory::Exists => Ok(dir),
            ContainingDirectory::CreateAllRaceProof(retries) => crate::create_dir::all(dir, retries),
        }
    }
}

fn expect_none<T>(v: Option<T>) {
    assert!(
        v.is_none(),
        "there should never be conflicts or old values as ids are never reused."
    );
}

impl Drop for Registration {
    fn drop(&mut self) {
        if let Some((_id, Some(tempfile))) = REGISTER.remove(&self.id) {
            let directory = tempfile
                .inner
                .path()
                .parent()
                .expect("every tempfile has a parent directory")
                .to_owned();
            drop(tempfile.inner);
            tempfile.cleanup.execute_best_effort(&directory);
        }
    }
}