s390-tools/rust/utils/src/file.rs

// SPDX-License-Identifier: MIT
//
// Copyright IBM Corp. 2024

use std::{
    ffi::OsStr,
    fs::{rename, File, OpenOptions},
    io::{self, Seek, SeekFrom, Write},
    os::unix::fs::OpenOptionsExt,
    path::Path,
};

use pv::{Error, FileAccessErrorType, PvCoreError, Result};

/// 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, PartialEq, Eq)]
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<()> {
        if let Ok(true) = self.path.try_exists() {
            if operation == AtomicFileOperation::NoReplace {
                return Err(Error::Io(io::Error::from(io::ErrorKind::AlreadyExists)));
            }
        }

        rename(&self.temp_path, &self.path).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());
    }
}