Skip to main content

save_data/
lib.rs

1use std::{
2    ffi::{OsStr, OsString},
3    fs::{self, File},
4    io,
5    path::{Path, PathBuf},
6};
7
8fn is_invalid_path(file_path: &Path) -> io::Result<bool> {
9    match std::fs::symlink_metadata(file_path) {
10        Ok(metadata) => return Ok(!metadata.is_file()),
11        Err(ref e) if e.kind() == io::ErrorKind::NotFound => (),
12        Err(e) => return Err(e),
13    }
14    if file_path.file_name().is_none() {
15        return Ok(true);
16    }
17    let Some(parent) = file_path.parent() else {
18        // `file_path` is "/" or ""
19        return Ok(true);
20    };
21    Ok(!(parent == Path::new("") || parent.is_dir()))
22}
23
24fn parent_dir(file_path: &Path) -> &Path {
25    let mut parent_dir = file_path.parent().unwrap();
26    if parent_dir == Path::new("") {
27        parent_dir = Path::new(".");
28    }
29    parent_dir
30}
31
32fn open_dir(path: &Path) -> io::Result<File> {
33    Ok(File::from(nix::fcntl::open(
34        path,
35        nix::fcntl::OFlag::O_RDONLY | nix::fcntl::OFlag::O_DIRECTORY,
36        nix::sys::stat::Mode::empty(),
37    )?))
38}
39
40fn create_temp_file(parent_path: &Path) -> io::Result<(File, PathBuf)> {
41    let template = parent_path.join("tmp.XXXXXX");
42    let (fd, path) = nix::unistd::mkstemp(&template)?;
43    Ok((File::from(fd), path))
44}
45
46fn write_to_temp_file_with_sync<R: io::Read>(
47    parent_path: &Path,
48    mut src: R,
49) -> io::Result<PathBuf> {
50    let (mut temp_file, temp_file_path) = create_temp_file(parent_path)?;
51    io::copy(&mut src, &mut temp_file)?;
52    temp_file.sync_all()?;
53    Ok(temp_file_path)
54}
55
56fn add_extension<S: AsRef<OsStr>>(path: &Path, ext: S) -> PathBuf {
57    let mut file_name = path.file_name().unwrap().to_owned();
58    file_name.push(".");
59    file_name.push(ext.as_ref());
60    path.with_file_name(file_name)
61}
62
63fn generate_random_temp_file_name(len: usize) -> OsString {
64    use rand::distr::{Alphanumeric, SampleString};
65    let mut temp_file_name = OsString::from(Alphanumeric.sample_string(&mut rand::rng(), len));
66    temp_file_name.push(".tmp");
67    temp_file_name
68}
69
70fn generate_temp_file_path(parent: Option<&Path>) -> PathBuf {
71    let temp_file_name = generate_random_temp_file_name(10);
72    if let Some(parent) = parent {
73        parent.join(temp_file_name)
74    } else {
75        PathBuf::from(temp_file_name)
76    }
77}
78
79fn retry_loop<T, F>(file_path: &Path, func: F) -> io::Result<(T, PathBuf)>
80where
81    F: Fn(&Path) -> io::Result<T>,
82{
83    let parent = file_path.parent();
84    let mut retry_remain = 3;
85    loop {
86        let temp_file_path = generate_temp_file_path(parent);
87        match func(&temp_file_path) {
88            Ok(value) => return Ok((value, temp_file_path)),
89            Err(e) => {
90                if e.kind() == io::ErrorKind::AlreadyExists && retry_remain > 0 {
91                    // This is very rare case. We retry it.
92                    retry_remain -= 1;
93                } else {
94                    return Err(e);
95                }
96            }
97        }
98    }
99}
100
101fn hard_link_to_temp_file(file_path: &Path) -> io::Result<PathBuf> {
102    Ok(retry_loop(file_path, |path| fs::hard_link(file_path, path))?.1)
103}
104
105fn make_backup_file(file_path: &Path) -> io::Result<()> {
106    let backup_path = add_extension(file_path, "orig");
107    match hard_link_to_temp_file(file_path) {
108        Ok(temp_file_path) => {
109            fs::rename(temp_file_path, backup_path)?;
110            Ok(())
111        }
112        // If `file_path` not exists yet, we can't make backup file.
113        Err(ref e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
114        Err(e) => Err(e),
115    }
116}
117
118/// The data read from `src` is saved to `file_path`.
119///
120/// The file is updated safely.
121/// This means that even if the operation is interrupted,
122/// `file_path` is guaranteed to contain either the complete old data or the complete new data.
123/// However, it is assumed that the data is permanently written to storage when
124/// [`File::fsync_all()`](https://doc.rust-lang.org/std/fs/struct.File.html#method.sync_all)
125/// completes successfully.
126///
127/// If `file_path` exists, the data is first written to a temporary file,
128/// then the file is replaced using `rename`.
129/// Additionally, a backup file of the original file is created.
130/// The path of the backup file is `file_path` appended with `.orig`.
131pub fn save_data<P: AsRef<Path>, R: io::Read>(file_path: P, src: R) -> io::Result<()> {
132    let file_path = file_path.as_ref();
133    if is_invalid_path(file_path)? {
134        return Err(io::Error::from(io::ErrorKind::InvalidInput));
135    }
136    let parent_path = parent_dir(file_path);
137    let parent_dir = open_dir(parent_path)?;
138    let temp_file_path = write_to_temp_file_with_sync(parent_path, src)?;
139    parent_dir.sync_all()?;
140    make_backup_file(file_path)?;
141    parent_dir.sync_all()?;
142    fs::rename(temp_file_path, file_path)?;
143    parent_dir.sync_all()?;
144    Ok(())
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    fn write_read(path: &Path, data: &[u8]) {
152        save_data(path, data).unwrap();
153        let data2 = std::fs::read("foo.bin").unwrap();
154        assert_eq!(data, &data2);
155    }
156
157    #[test]
158    fn it_works() {
159        let mut data = [0u8; 4096];
160        for _ in 0..10 {
161            rand::fill(&mut data);
162            write_read("foo.bin".as_ref(), &data);
163        }
164    }
165
166    #[test]
167    fn path_test() {
168        assert!(is_invalid_path(Path::new("/a/b/..")).unwrap());
169
170        assert!(is_invalid_path(Path::new("/")).unwrap());
171        assert!(is_invalid_path(Path::new("")).unwrap());
172    }
173}