packie/src/common/fs.rs

use rust_i18n::t;
use std::{
    fs::{FileType, ReadDir},
    path::{Path, PathBuf},
    time::SystemTime,
};

#[derive(Debug)]
pub struct TreeDir {
    stack: Vec<(ReadDir, PathBuf)>,
}
impl TreeDir {
    pub fn new(path: impl Into<PathBuf>) -> std::io::Result<Self> {
        let root_path = path.into();
        let read_dir = std::fs::read_dir(&root_path)?;
        Ok(Self {
            stack: vec![(read_dir, PathBuf::new())],
        })
    }
}
impl Iterator for TreeDir {
    type Item = std::io::Result<TreeDirEntry>;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some((read_dir, current_relpath)) = self.stack.last_mut() {
            match read_dir.next() {
                Some(Ok(entry)) => {
                    let file_name = entry.file_name();
                    let relpath = current_relpath.join(file_name);

                    let ty = match entry.file_type() {
                        Ok(ty) => ty,
                        Err(e) => return Some(Err(e)),
                    };

                    if ty.is_dir() {
                        match std::fs::read_dir(entry.path()) {
                            Ok(sub_read_dir) => {
                                self.stack.push((sub_read_dir, relpath.clone()));
                            }
                            Err(e) => {
                                return Some(Err(e));
                            }
                        }
                    }

                    return Some(Ok(TreeDirEntry { relpath, ty }));
                }
                Some(Err(e)) => {
                    return Some(Err(e));
                }
                None => {
                    self.stack.pop();
                }
            }
        }
        None
    }
}

#[derive(Debug)]
pub struct TreeDirEntry {
    pub relpath: PathBuf,
    pub ty: FileType,
}

pub fn du_dir(path: impl AsRef<Path>) -> std::io::Result<u64> {
    let mut size = 0;
    let path = path.as_ref();
    for ent in TreeDir::new(path)? {
        let ent = ent?;
        let metadata = std::fs::symlink_metadata(path.join(&ent.relpath))?;
        size += metadata.len();
    }
    Ok(size)
}

pub fn copy_dir(src: impl Into<PathBuf>, dst: impl AsRef<Path>) -> Result<(), CopyError> {
    let src = src.into();
    std::fs::create_dir(dst.as_ref()).map_err(|error| CopyError {
        src: src.clone(),
        dst: dst.as_ref().into(),
        error,
    })?;
    match copy_dir_raw(src, dst.as_ref()) {
        Ok(()) => Ok(()),
        Err(err) => {
            _ = std::fs::remove_dir_all(dst.as_ref());
            Err(err)
        }
    }
}

fn copy_dir_raw(src: PathBuf, dst: &Path) -> Result<(), CopyError> {
    let global_err = |error| CopyError {
        src: src.clone(),
        dst: dst.into(),
        error,
    };
    let src_tree = TreeDir::new(src.clone()).map_err(global_err)?;
    for ent in src_tree {
        let ent = ent.map_err(global_err)?;
        let src_full_path = src.join(&ent.relpath);
        let dst_full_path = dst.join(&ent.relpath);

        copy_fs_node(&src_full_path, &dst_full_path).map_err(|error| CopyError {
            src: src_full_path.clone(),
            dst: dst_full_path.clone(),
            error,
        })?;
    }
    Ok(())
}

fn copy_fs_node(src: &Path, dst: &Path) -> std::io::Result<()> {
    let metadata = std::fs::symlink_metadata(src)?;

    if metadata.file_type().is_dir() {
        std::fs::create_dir(dst)?;
    } else if metadata.is_symlink() {
        let link_to = std::fs::read_link(src)?;
        symlink(&link_to, dst)?;
        return Ok(());
    } else {
        std::fs::copy(&src, &dst)?;
    }

    std::fs::set_permissions(dst, metadata.permissions())?;
    set_modified_time(dst, metadata.modified()?)?;

    Ok(())
}

#[cfg(target_family = "unix")]
pub fn symlink(src: &Path, dst: &Path) -> std::io::Result<()> {
    std::os::unix::fs::symlink(src, dst)
}

#[cfg(target_family = "unix")]
fn set_modified_time(path: &Path, modified: SystemTime) -> std::io::Result<()> {
    let modified_duration = modified
        .duration_since(SystemTime::UNIX_EPOCH)
        .map_err(|_| std::io::ErrorKind::InvalidData)?;
    unsafe {
        nix::sys::stat::utimensat(
            std::os::fd::BorrowedFd::borrow_raw(nix::libc::AT_FDCWD),
            path,
            &nix::sys::time::TimeSpec::UTIME_NOW,
            &nix::sys::time::TimeSpec::from_duration(modified_duration),
            nix::sys::stat::UtimensatFlags::NoFollowSymlink,
        )
        .map_err(|errno| std::io::Error::from_raw_os_error(errno as _))
    }
}

#[derive(Debug, thiserror::Error)]
#[error("{}", t!("common.CopyError", src = src.display(), dst = dst.display(), error = error))]
pub struct CopyError {
    src: PathBuf,
    dst: PathBuf,
    error: std::io::Error,
}

#[cfg(target_family = "unix")]
#[derive(Debug)]
pub struct LockGuard {
    _inner: Option<nix::fcntl::Flock<std::fs::File>>,
}
#[cfg(target_family = "unix")]
impl LockGuard {
    pub fn noop() -> Self {
        Self { _inner: None }
    }

    pub fn open(path: impl AsRef<Path>, nonblocking: bool) -> std::io::Result<Self> {
        let file = std::fs::File::options()
            .create(true)
            .write(true)
            .open(path)?;
        let flags = if nonblocking {
            nix::fcntl::FlockArg::LockExclusiveNonblock
        } else {
            nix::fcntl::FlockArg::LockExclusive
        };
        let lock = nix::fcntl::Flock::lock(file, flags).map_err(|x| x.1)?;
        Ok(Self { _inner: Some(lock) })
    }
}