rust/utils: Add AtomicFile

This type helps to perform atomic operations by writing to a temporary
file and renaming it to the actual filename when the
`AtomicFile::finish` function is called. If the `AtomicFile::finish`
function is never called, the temporary file is automatically removed
when it goes out of scope. It utilizes the `renameat2` [1] libc function
and its semantics.

[1] https://man7.org/linux/man-pages/man2/renameat.2.html
Reviewed-by: Steffen Eiden <seiden@linux.ibm.com>
Signed-off-by: Marc Hartmayer <mhartmay@linux.ibm.com>
Signed-off-by: Steffen Eiden <seiden@linux.ibm.com>

rust/utils/src/file.rs

@@ -0,0 +1,303 @@
+use std::{
+    ffi::{CString, OsStr},
+    fs::{File, OpenOptions},
+    io::{self, Seek, SeekFrom, Write},
+    os::unix::{ffi::OsStrExt, fs::OpenOptionsExt},
+    path::Path,
+};
+
+use pv::{Error, FileAccessErrorType, PvCoreError, Result};
+
+/// Rust wrapper for `libc::renameat2`
+fn renameat2<P: AsRef<Path>, Q: AsRef<Path>>(oldpath: P, newpath: Q, flags: u32) -> io::Result<()> {
+    let oldpath_cstr = CString::new(oldpath.as_ref().as_os_str().as_bytes())?;
+    let oldpath_raw = oldpath_cstr.into_raw();
+    let newpath_cstr = CString::new(newpath.as_ref().as_os_str().as_bytes())?;
+    let newpath_raw = newpath_cstr.into_raw();
+    unsafe {
+        // SAFETY: oldpath_raw and newpath_raw are valid CStrings because they were
+        // generated by the `CString::new` function.
+        let ret = libc::renameat2(
+            libc::AT_FDCWD,
+            oldpath_raw,
+            libc::AT_FDCWD,
+            newpath_raw,
+            flags,
+        );
+        // SAFETY: libc::renameat2 does not modify `newpath_raw` and is
+        // therefore still valid.
+        let _ = CString::from_raw(newpath_raw);
+        // SAFETY: libc::renameat2 does not modify `oldpath_raw` and is
+        // therefore still valid.
+        let _ = CString::from_raw(oldpath_raw);
+        if ret == -1 {
+            return Err(io::Error::last_os_error());
+        }
+        Ok(())
+    }
+}
+
+/// This type helps to perform atomic operations by writing to a temporary file
+/// and renaming it to the actual filename when the [`AtomicFile::finish`]
+/// function is called. If the [`AtomicFile::finish`] function is never called,
+/// the temporary file is automatically removed when it goes out of scope. It
+/// utilizes the `renameat2` libc function and its semantics.
+#[derive(Debug)]
+pub struct AtomicFile<F = File> {
+    /// Do not change the order! See comment in [`TempPath::drop`].
+    file: F,
+    path: TempPath,
+}
+
+impl<F> AsRef<F> for AtomicFile<F> {
+    fn as_ref(&self) -> &F {
+        &self.file
+    }
+}
+
+impl<F> AsMut<F> for AtomicFile<F> {
+    fn as_mut(&mut self) -> &mut F {
+        &mut self.file
+    }
+}
+
+/// Enum used for more verbosity.
+#[derive(Debug)]
+pub enum AtomicFileOperation {
+    /// Replace existing file
+    Replace,
+    /// Do not replace existing file
+    NoReplace,
+}
+
+impl<F: Write> AtomicFile<F> {}
+
+impl AtomicFile<File> {
+    /// Creates a new [`AtomicFile`] at the given `output` path using
+    /// `open_options`.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if the temporary file could not be
+    /// created.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use utils::AtomicFile;
+    ///
+    /// let file = AtomicFile::new("test", &mut std::fs::OpenOptions::new()).unwrap();
+    /// ```
+    pub fn new<P: AsRef<Path>>(output: P, open_options: &mut OpenOptions) -> Result<Self> {
+        Self::with_extension(output, "part", open_options)
+    }
+
+    /// Creates a new [`AtomicFile`] at the given `output` path using
+    /// `open_options` and `suffix` as the suffix for the temporary file.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if the temporary file could not be
+    /// created.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use utils::AtomicFile;
+    ///
+    /// let file = AtomicFile::with_extension("test", ".incomplete", &mut std::fs::OpenOptions::new()).unwrap();
+    /// ```
+    pub fn with_extension<P: AsRef<Path>, S: AsRef<OsStr>>(
+        output: P,
+        tmp_suffix: S,
+        open_options: &mut OpenOptions,
+    ) -> Result<Self> {
+        let path = TempPath::new(output, tmp_suffix);
+
+        open_options
+            .read(true)
+            .write(true)
+            .create_new(true)
+            .mode(0o600);
+        match open_options.open(&path.temp_path) {
+            Ok(file) => Ok(Self { path, file }),
+            Err(err) => {
+                let path_copy = path.temp_path.to_path_buf();
+                path.forget();
+                Err(Error::PvCore(PvCoreError::FileAccess {
+                    ty: FileAccessErrorType::Create,
+                    path: path_copy,
+                    source: err,
+                }))
+            }
+        }
+    }
+
+    /// Renames the file to the actual file name. If `overwrite` is set to
+    /// [`AtomicFileOperation::Replace`] the file is overwritten.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if the rename operation fails.
+    pub fn finish(mut self, overwrite: AtomicFileOperation) -> Result<()> {
+        self.file.flush()?;
+        self.file.sync_all()?;
+        drop(self.file);
+        self.path.persist(overwrite)
+    }
+
+    /// Discard all changes and delete the temporary file.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if the temporary file could not be
+    /// deleted.
+    pub fn discard(mut self) -> io::Result<()> {
+        self.file.flush()?;
+        self.file.sync_all()?;
+        drop(self.file);
+        self.path.remove_file()
+    }
+}
+
+impl<S: Seek> Seek for AtomicFile<S> {
+    /// Seek to the offset (in bytes) in the underlying file.
+    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
+        self.as_mut().seek(pos)
+    }
+}
+
+impl Write for AtomicFile {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.as_mut().write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.as_mut().flush()
+    }
+}
+
+#[derive(Debug)]
+struct TempPath {
+    temp_path: Box<Path>,
+    path: Box<Path>,
+}
+
+impl TempPath {
+    fn new<P: AsRef<Path>, S: AsRef<OsStr>>(path: P, extension: S) -> Self {
+        let mut temp_path = path.as_ref().to_path_buf();
+        match temp_path.extension() {
+            Some(ext) => {
+                let mut ext = ext.to_os_string();
+                ext.push(".");
+                ext.push(extension.as_ref());
+                temp_path.set_extension(ext)
+            }
+            None => temp_path.set_extension(extension.as_ref()),
+        };
+        Self {
+            temp_path: temp_path.into(),
+            path: path.as_ref().into(),
+        }
+    }
+
+    fn forget(self) {
+        // Make sure that [`Self::drop`] is not called.
+        Self { .. } = self;
+    }
+
+    /// Removes the temporary file
+    pub fn remove_file(self) -> io::Result<()> {
+        let result = std::fs::remove_file(&self.temp_path);
+        self.forget();
+        result
+    }
+
+    fn persist(self, operation: AtomicFileOperation) -> Result<()> {
+        let options = match operation {
+            AtomicFileOperation::Replace => 0,
+            AtomicFileOperation::NoReplace => libc::RENAME_NOREPLACE,
+        };
+
+        renameat2(&self.temp_path, &self.path, options).map_err(|e| {
+            PvCoreError::FileAccessRename {
+                src: self.temp_path.as_ref().to_str().unwrap().to_string(),
+                dst: self.path.as_ref().to_str().unwrap().to_string(),
+                source: e,
+            }
+        })?;
+        self.forget();
+        Ok(())
+    }
+}
+
+impl Drop for TempPath {
+    // When this is called while dropping an [`AtomicFile`],
+    // the field [`AtomicFile::file`] is already dropped.
+    fn drop(&mut self) {
+        let _ = std::fs::remove_file(&self.temp_path);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::Write;
+
+    use super::{AtomicFile, AtomicFileOperation};
+    use crate::TemporaryDirectory;
+
+    #[test]
+    fn atomicfile_basic_functionality() {
+        let tmp_dir = TemporaryDirectory::new().expect("should work");
+        let tmp_file = tmp_dir.path().join("atomic_basic");
+        let data = b"test_data";
+        assert!(!tmp_file.exists());
+        let mut writer =
+            AtomicFile::with_extension(&tmp_file, "basic", &mut std::fs::OpenOptions::new())
+                .unwrap();
+
+        writer.write_all(data).unwrap();
+        writer.finish(AtomicFileOperation::Replace).unwrap();
+        assert!(tmp_file.exists());
+
+        let writer =
+            AtomicFile::with_extension(&tmp_file, "basic", &mut std::fs::OpenOptions::new())
+                .unwrap();
+        writer.finish(AtomicFileOperation::NoReplace).expect_err(
+            "Creating
+            the file is expected to fail",
+        );
+        std::fs::remove_file(tmp_file).expect("should not fail");
+    }
+
+    #[test]
+    fn atomicfile_discard() {
+        let tmp_dir = TemporaryDirectory::new().expect("should work");
+        let tmp_file = tmp_dir.path().join("atomic_close");
+        let data = b"test_data";
+        assert!(!tmp_file.exists());
+        let mut writer =
+            AtomicFile::with_extension(&tmp_file, "close", &mut std::fs::OpenOptions::new())
+                .unwrap();
+
+        writer.write_all(data).unwrap();
+
+        writer.discard().expect("bla");
+        assert!(!tmp_file.exists());
+    }
+
+    #[test]
+    fn atomicfile_drop() {
+        let tmp_dir = TemporaryDirectory::new().expect("should work");
+        let tmp_file = tmp_dir.path().join("atomic_close");
+        let data = b"test_data";
+        assert!(!tmp_file.exists());
+        let mut writer =
+            AtomicFile::with_extension(&tmp_file, "close", &mut std::fs::OpenOptions::new())
+                .unwrap();
+
+        writer.write_all(data).unwrap();
+        drop(writer);
+        assert!(!tmp_file.exists());
+    }
+}

rust/utils/src/lib.rs

@@ -4,19 +4,24 @@
 //!
 //! Copyright IBM Corp. 2023, 2024
 mod cli;
+mod file;
 mod hexslice;
 mod log;
 mod tmpfile;
 
-pub use crate::cli::{get_reader_from_cli_file_arg, get_writer_from_cli_file_arg};
-pub use crate::cli::{print_cli_error, print_error};
-pub use crate::cli::{CertificateOptions, DeprecatedVerbosityOptions, VerbosityOptions};
-pub use crate::cli::{STDIN, STDOUT};
-pub use crate::hexslice::HexSlice;
-pub use crate::log::PvLogger;
-pub use crate::tmpfile::TemporaryDirectory;
 pub use ::log::LevelFilter;
 
+pub use crate::{
+    cli::{
+        get_reader_from_cli_file_arg, get_writer_from_cli_file_arg, print_cli_error, print_error,
+        CertificateOptions, DeprecatedVerbosityOptions, VerbosityOptions, STDIN, STDOUT,
+    },
+    file::{AtomicFile, AtomicFileOperation},
+    hexslice::HexSlice,
+    log::PvLogger,
+    tmpfile::TemporaryDirectory,
+};
+
 /// Get the s390-tools release string
 ///
 /// Provides the s390-tools release string.