cosmic-files/src/operation/recursive.rs

// Copyright 2023 System76 <info@system76.com>
// SPDX-License-Identifier: GPL-3.0-only

use super::{Controller, OperationSelection, ReplaceResult, copy_unique_path};
use crate::operation::{OperationError, sync_to_disk};
use anyhow::Context as AnyhowContext;
use compio::BufResult;
use compio::buf::{IntoInner, IoBuf};
use compio::driver::ToSharedFd;
use compio::driver::op::AsyncifyFd;
use compio::io::{AsyncReadAt, AsyncWriteAt};
use cosmic::iced::futures;
use futures::{FutureExt, StreamExt};
use std::cell::Cell;
use std::error::Error;
use std::fs;
use std::future::Future;
use std::ops::ControlFlow;
use std::path::PathBuf;
use std::pin::Pin;
use std::rc::Rc;
use std::time::Instant;
use walkdir::WalkDir;

#[cfg(feature = "gvfs")]
use gio::prelude::FileExtManual;

#[derive(thiserror::Error, Debug)]
pub enum GioCopyError {
    #[error("controller state")]
    Controller(OperationError),
    #[cfg(feature = "gvfs")]
    #[error("gio copy failed")]
    GLib(#[from] glib::Error),
}

pub enum Method {
    Copy,
    Move { cross_device_copy: bool },
}

pub struct Context {
    buf: Vec<u8>,
    controller: Controller,
    on_progress: Box<dyn OnProgress>,
    on_replace: Pin<Box<dyn OnReplace>>,
    pub(crate) op_sel: OperationSelection,
    replace_result_opt: Option<ReplaceResult>,
    remaining_conflicts: usize,
}

pub trait OnProgress: Fn(&Op, &Progress) + 'static {}
impl<F> OnProgress for F where F: Fn(&Op, &Progress) + 'static {}

pub trait OnReplace:
    for<'a> Fn(&'a Op, usize) -> Pin<Box<dyn Future<Output = ReplaceResult> + 'a>> + 'static
{
}
impl<F> OnReplace for F where
    F: for<'a> Fn(&'a Op, usize) -> Pin<Box<dyn Future<Output = ReplaceResult> + 'a>> + 'static
{
}

impl Context {
    pub fn new(controller: Controller) -> Self {
        Self {
            // 128K is the optimal upper size of a buffer.
            buf: vec![0u8; 128 * 1024],
            controller,
            on_progress: Box::new(|_op, _progress| {}),
            on_replace: Box::pin(|_op, _count| Box::pin(async { ReplaceResult::Cancel })),
            op_sel: OperationSelection::default(),
            replace_result_opt: None,
            remaining_conflicts: 0,
        }
    }

    pub async fn recursive_copy_or_move(
        &mut self,
        from_to_pairs: impl IntoIterator<Item = (PathBuf, PathBuf)>,
        method: Method,
    ) -> Result<bool, OperationError> {
        let mut ops = Vec::new();
        let mut cleanup_ops = Vec::new();
        let mut written_files = Vec::new();
        let mut target_dirs = std::collections::HashSet::new();
        for (from_parent, to_parent) in from_to_pairs {
            self.controller
                .check()
                .await
                .map_err(|s| OperationError::from_state(s, &self.controller))?;

            if from_parent == to_parent {
                // Skip matching source and destination
                continue;
            }

            for entry in WalkDir::new(&from_parent) {
                self.controller
                    .check()
                    .await
                    .map_err(|s| OperationError::from_state(s, &self.controller))?;

                let entry = entry.map_err(|err| {
                    OperationError::from_err(
                        format!(
                            "failed to walk directory {}: {}",
                            from_parent.display(),
                            err
                        ),
                        &self.controller,
                    )
                })?;
                let file_type = entry.file_type();
                let from = entry.into_path();
                let kind = if file_type.is_dir() {
                    OpKind::Mkdir
                } else if file_type.is_file() {
                    match method {
                        Method::Copy => OpKind::Copy,
                        Method::Move { cross_device_copy } => OpKind::Move { cross_device_copy },
                    }
                } else if file_type.is_symlink() {
                    let target = fs::read_link(&from).map_err(|err| {
                        OperationError::from_err(
                            format!("failed to read link {}: {}", from_parent.display(), err),
                            &self.controller,
                        )
                    })?;
                    OpKind::Symlink { target }
                } else {
                    //TODO: present dialog and allow continue
                    return Err(OperationError::from_err(
                        format!("{} is not a known file type", from.display()),
                        &self.controller,
                    ));
                };
                let to = if from == from_parent {
                    // When copying a file, from matches from_parent, and to_parent must be used
                    to_parent.clone()
                } else {
                    let relative = from.strip_prefix(&from_parent).map_err(|err| {
                        OperationError::from_err(
                            format!(
                                "failed to remove prefix {} from {}: {}",
                                from_parent.display(),
                                from.display(),
                                err
                            ),
                            &self.controller,
                        )
                    })?;
                    //TODO: ensure to is inside of to_parent?
                    to_parent.join(relative)
                };
                let op = Op {
                    kind,
                    from,
                    to,
                    skipped: Rc::new(Skip {
                        normal: Cell::new(false),
                        cleanup: Cell::new(false),
                    }),
                    is_cleanup: false,
                };
                if matches!(method, Method::Move { .. })
                    && let Some(cleanup_op) = op.move_cleanup_op()
                {
                    cleanup_ops.push(cleanup_op);
                }
                if let Some(parent) = op.to.parent() {
                    target_dirs.insert(parent.to_path_buf());
                }
                ops.push(op);
            }

            self.op_sel.ignored.push(from_parent);
        }

        // Add cleanup ops after standard ops, in reverse
        cleanup_ops.reverse();
        ops.append(&mut cleanup_ops);

        // Count potential conflicts (files that would need replacement)
        self.remaining_conflicts = ops
            .iter()
            .filter(|op| {
                matches!(
                    op.kind,
                    OpKind::Copy | OpKind::Move { .. } | OpKind::Symlink { .. }
                ) && op.to.is_file()
            })
            .count();

        let total_ops = ops.len();
        for (current_ops, mut op) in ops.into_iter().enumerate() {
            self.controller
                .check()
                .await
                .map_err(|s| OperationError::from_state(s, &self.controller))?;

            let progress = Progress {
                current_ops,
                total_ops,
                current_bytes: 0,
                total_bytes: None,
            };
            (self.on_progress)(&op, &progress);
            if op.run(self, progress).await.map_err(|err| {
                OperationError::from_err(
                    format!(
                        "failed to {:?} {} to {}: {}",
                        op.kind,
                        op.from.display(),
                        op.to.display(),
                        err
                    ),
                    &self.controller,
                )
            })? {
                if matches!(
                    op.kind,
                    OpKind::Copy
                        | OpKind::Move {
                            cross_device_copy: true
                        }
                ) {
                    written_files.push(op.to.clone());
                }
                // The from path is ignored in the operation selection if it is a top level item
                if self.op_sel.ignored.contains(&op.from) {
                    // So add the to path to the selection
                    self.op_sel.selected.push(op.to);
                }
            } else {
                // Cancelled
                return Ok(false);
            }
        }

        // Flush files to disk
        sync_to_disk(written_files, target_dirs).await;

        Ok(true)
    }

    pub fn on_progress<F: OnProgress>(mut self, f: F) -> Self {
        self.on_progress = Box::new(f);
        self
    }

    pub fn on_replace(mut self, f: impl OnReplace + 'static) -> Self {
        self.on_replace = Box::pin(f);
        self
    }

    async fn replace(&mut self, op: &Op) -> Result<ControlFlow<bool, PathBuf>, Box<dyn Error>> {
        let replace_result = match self.replace_result_opt {
            Some(result) => result,
            None => (self.on_replace)(op, self.remaining_conflicts).await,
        };

        match replace_result {
            ReplaceResult::Replace(apply_to_all) => {
                if apply_to_all {
                    self.replace_result_opt = Some(replace_result);
                }
                compio::fs::remove_file(&op.to).await?;
                Ok(ControlFlow::Continue(op.to.clone()))
            }
            ReplaceResult::KeepBoth => match op.to.parent() {
                Some(to_parent) => Ok(ControlFlow::Continue(copy_unique_path(&op.from, to_parent))),
                None => Err(format!("failed to get parent of {}", op.to.display()).into()),
            },
            ReplaceResult::Skip(apply_to_all) => {
                if apply_to_all {
                    self.replace_result_opt = Some(replace_result);
                }
                op.skipped.normal.set(true);
                Ok(ControlFlow::Break(true))
            }
            ReplaceResult::Cancel => Ok(ControlFlow::Break(false)),
        }
    }
}

#[derive(Debug)]
pub struct Progress {
    pub current_ops: usize,
    pub total_ops: usize,
    pub current_bytes: u64,
    pub total_bytes: Option<u64>,
}

#[derive(Debug)]
pub enum OpKind {
    Copy,
    Move { cross_device_copy: bool },
    Mkdir,
    Remove,
    Rmdir,
    Symlink { target: PathBuf },
}

#[derive(Debug)]
pub struct Skip {
    /// Normal operation should be skipped
    pub normal: Cell<bool>,
    /// Cleanup operation should be skipped
    pub cleanup: Cell<bool>,
}

#[derive(Debug)]
pub struct Op {
    pub kind: OpKind,
    pub from: PathBuf,
    pub to: PathBuf,
    pub skipped: Rc<Skip>,
    pub is_cleanup: bool,
}

impl Op {
    fn move_cleanup_op(&self) -> Option<Self> {
        let kind = match self.kind {
            OpKind::Copy | OpKind::Move { .. } | OpKind::Symlink { .. } => OpKind::Remove,
            OpKind::Mkdir => OpKind::Rmdir,
            OpKind::Remove | OpKind::Rmdir => return None,
        };
        Some(Self {
            kind,
            from: self.from.clone(),
            //TODO: it is strange to have `to` here
            to: self.to.clone(),
            skipped: self.skipped.clone(),
            is_cleanup: true,
        })
    }

    async fn run(&mut self, ctx: &mut Context, progress: Progress) -> Result<bool, Box<dyn Error>> {
        if self.skipped.normal.get() || (self.is_cleanup && self.skipped.cleanup.get()) {
            return Ok(true);
        }
        match self.kind {
            OpKind::Copy => {
                crate::operation::actively_writing_add(self.to.clone());
                let result = self.copy(ctx, progress).await;

                if result.is_err() {
                    _ = compio::fs::remove_file(&self.to).await;
                }

                crate::operation::actively_writing_remove(&self.to);
                return result;
            }
            OpKind::Move { cross_device_copy } => {
                // Do not clean up if cross_device_copy is set
                if cross_device_copy {
                    self.skipped.cleanup.set(true);
                }

                // Remove `to` if overwriting and it is an existing file
                if self.to.is_file() {
                    match ctx.replace(self).await? {
                        ControlFlow::Continue(to) => {
                            self.to = to;
                        }
                        ControlFlow::Break(ret) => {
                            return Ok(ret);
                        }
                    }
                }
                // This is atomic and ensures `to` is not created by any other process
                match compio::fs::hard_link(&self.from, &self.to).await {
                    Ok(()) => {}
                    Err(err) => {
                        // https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_NOT_SAME_DEVICE.html
                        #[cfg(windows)]
                        const EXDEV: i32 = 17;
                        #[cfg(unix)]
                        const EXDEV: i32 = libc::EXDEV as _;

                        if err.raw_os_error() == Some(EXDEV) {
                            if cross_device_copy {
                                // Do not clean up if cross_device_copy is set
                                self.skipped.cleanup.set(true);
                            }
                            // Try standard copy if hard link fails with cross device error
                            let mut copy_op = Self {
                                kind: OpKind::Copy,
                                from: self.from.clone(),
                                to: self.to.clone(),
                                skipped: self.skipped.clone(),
                                is_cleanup: self.is_cleanup,
                            };
                            return Box::pin(copy_op.run(ctx, progress)).await;
                        }
                        return Err(err.into());
                    }
                }
            }
            OpKind::Mkdir => {
                compio::fs::create_dir_all(&self.to).await?;
            }
            OpKind::Remove => {
                compio::fs::remove_file(&self.from).await?;
            }
            OpKind::Rmdir => {
                compio::fs::remove_dir(&self.from).await?;
            }
            OpKind::Symlink { ref target } => {
                // Remove `to` if overwriting and it is an existing file
                if self.to.is_file() {
                    match ctx.replace(self).await? {
                        ControlFlow::Continue(to) => {
                            self.to = to;
                        }
                        ControlFlow::Break(ret) => {
                            return Ok(ret);
                        }
                    }
                }
                #[cfg(unix)]
                {
                    std::os::unix::fs::symlink(target, &self.to)?;
                }
                #[cfg(windows)]
                {
                    if target.is_dir() {
                        std::os::windows::fs::symlink_dir(target, &self.to)?;
                    } else {
                        std::os::windows::fs::symlink_file(target, &self.to)?;
                    }
                }
            }
        }
        Ok(true)
    }

    async fn copy(
        &mut self,
        ctx: &mut Context,
        mut progress: Progress,
    ) -> Result<bool, Box<dyn Error>> {
        // Remove `to` if overwriting and it is an existing file
        if self.to.is_file() {
            match ctx.replace(self).await? {
                ControlFlow::Continue(to) => {
                    self.to = to;
                }
                ControlFlow::Break(ret) => {
                    return Ok(ret);
                }
            }
        }

        let (from_file, metadata, to_file) = cosmic::iced::futures::join!(
            async {
                compio::fs::OpenOptions::new()
                    .read(true)
                    .open(&self.from)
                    .await
                    .with_context(|| format!("failed to open {} for reading", self.from.display(),))
            },
            async { compio::fs::metadata(&self.from).await.ok() },
            // This is atomic and ensures `to` is not created by any other process
            async {
                compio::fs::OpenOptions::new()
                    .create_new(true)
                    .write(true)
                    .open(&self.to)
                    .await
                    .with_context(|| format!("failed to open {} for writing", self.to.display()))
            }
        );

        let from_file = from_file?;
        let mut to_file = to_file?;
        progress.total_bytes = metadata.as_ref().map(|m| m.len());
        (ctx.on_progress)(self, &progress);

        if let Some(metadata) = metadata.as_ref()
            && let Err(why) = to_file.set_permissions(metadata.permissions()).await
        {
            // This error is not propagated upwards as some filesystems do not support setting permissions
            if !matches!(why.kind(), std::io::ErrorKind::Unsupported) {
                tracing::warn!(?why, "failed to set permissions for {}", self.to.display(),);
            }
        }

        // Prevent spamming the progress callbacks.
        let mut last_progress_update = Instant::now();
        // io_uring/IOCP requires transferring ownership of the buffer to the kernel.
        let mut buf_in = std::mem::take(&mut ctx.buf);
        // Track where the current read/write position is at.
        let mut pos = 0;

        loop {
            let BufResult(result, buf_out) = from_file.read_at(buf_in, pos).await;

            let count = match result {
                Ok(0) => {
                    buf_in = buf_out;
                    break;
                }
                Ok(count) => count,
                Err(why) => {
                    ctx.buf = buf_out;
                    tracing::error!("failed to read: {:?}", why);
                    _ = futures::future::join(from_file.close(), to_file.close()).await;
                    return Err(why).context("failed to read")?;
                }
            };

            let BufResult(result, buf_out_slice) =
                to_file.write_at(buf_out.slice(..count), pos).await;
            let buf_out = buf_out_slice.into_inner();

            if let Err(why) = result {
                #[cfg(feature = "gvfs")]
                if let std::io::ErrorKind::Unsupported = why.kind() {
                    ctx.buf = buf_out;
                    _ = futures::future::join(from_file.close(), to_file.close()).await;
                    return self
                        .gio_file_copy(ctx, progress)
                        .await
                        .map(|_| true)
                        .map_err(Into::into);
                }

                tracing::error!("failed to write: {:?}", why);
                ctx.buf = buf_out;
                _ = futures::future::join(from_file.close(), to_file.close()).await;
                return Err(why).context("failed to write")?;
            }

            progress.current_bytes += count as u64;
            pos += count as u64;

            // Avoid spamming progress messages too early.
            let current = Instant::now();
            if current.duration_since(last_progress_update).as_millis() > 49 {
                last_progress_update = current;
                (ctx.on_progress)(self, &progress);

                // Also check if the progress was cancelled.
                if let Err(state) = ctx.controller.check().await {
                    ctx.buf = buf_out;
                    tracing::warn!(
                        "operation to copy from {:?} to {:?} cancelled",
                        self.from,
                        self.to
                    );
                    _ = futures::future::join(from_file.close(), to_file.close()).await;
                    return Err(OperationError::from_state(state, &ctx.controller).into());
                }
            }

            buf_in = buf_out;
        }

        ctx.buf = buf_in;

        if let Some(metadata) = metadata.as_ref() {
            let mut times = fs::FileTimes::new();
            if let Ok(time) = metadata.modified() {
                times = times.set_modified(time);
            }
            if let Ok(time) = metadata.accessed() {
                times = times.set_accessed(time);
            }
            //TODO: upstream set_times implementation to compio?
            let op = AsyncifyFd::new(to_file.to_shared_fd(), move |file: &std::fs::File| {
                BufResult(file.set_times(times).map(|_| 0), ())
            });
            match compio::runtime::submit(op).await.0.map(|_| ()) {
                Ok(()) => {
                    tracing::info!("set times for {} to {:?}", self.to.display(), times);
                }
                Err(why) => {
                    if !matches!(why.kind(), std::io::ErrorKind::Unsupported) {
                        tracing::error!(?why, "failed to set times for {}", self.to.display());
                    }
                }
            }
        }

        _ = to_file.close().await;

        Ok(true)
    }

    /// Fallback mechanism in the event that unsupported I/O error errors occur.
    /// Fixes unsupported errors when copying large files over MTP.
    /// TODO: Find what Gio.File does to work around this.
    #[cfg(feature = "gvfs")]
    async fn gio_file_copy(
        &self,
        ctx: &mut Context,
        mut progress: Progress,
    ) -> Result<(), GioCopyError> {
        _ = compio::fs::remove_file(&self.to).await;

        let from = gio::File::for_path(&self.from);
        let to = gio::File::for_path(&self.to);
        let (progress_tx, mut progress_rx) = tokio::sync::mpsc::unbounded_channel();
        let (tx, rx) = tokio::sync::oneshot::channel();
        let (pause_tx, mut pause_rx) = tokio::sync::watch::channel(false);

        let task = compio::runtime::spawn_blocking(move || {
            let glib_context = glib::MainContext::new();
            let glib_loop = glib::MainLoop::new(Some(&glib_context), false);
            glib_context.with_thread_default(move || {
                let glib_loop2 = glib_loop.clone();
                glib::MainContext::ref_thread_default().spawn_local(async move {
                    // Create a future for copying the file with `gio::File`. This also creates a progress stream.
                    let (gio_copy_fut, mut progress_stream) = from.copy_future(
                        &to,
                        gio::FileCopyFlags::OVERWRITE | gio::FileCopyFlags::ALL_METADATA,
                        glib::Priority::LOW,
                    );

                    let mut copy_fut = gio_copy_fut
                        .map(|result| result.map_err(GioCopyError::GLib))
                        .fuse();

                    let progress_fut = std::pin::pin!(async {
                        while let Some((current_bytes, _)) = progress_stream.next().await {
                            _ = progress_tx.send(current_bytes);
                        }

                        drop(progress_tx);
                        futures::future::pending::<()>().await;
                    });

                    let mut progress_fut = progress_fut.fuse();
                    let mut pause_rx2 = pause_rx.clone();

                    loop {
                        let until_paused = std::pin::pin!(pause_rx.wait_for(|paused| *paused));
                        futures::select! {
                            _ = &mut progress_fut => {},

                            result = &mut copy_fut => {
                                _ = tx.send(result.map(|_| ()));
                                glib_loop2.quit();
                                return;
                            }

                            _ = until_paused.fuse() => {
                                _ = pause_rx2.wait_for(|paused| !*paused).await;
                            }
                        }
                    }
                });

                glib_loop.run();
            })
        });

        let mut last_progress_update = Instant::now();
        let mut task = task.fuse();
        let mut rx = rx.fuse();

        loop {
            let until_paused = std::pin::pin!(ctx.controller.until_paused());
            futures::select! {
                value = progress_rx.recv().fuse() => {
                    if let Some(current_bytes) = value {
                        progress.current_bytes = current_bytes as u64;
                        let current = Instant::now();
                        if current.duration_since(last_progress_update).as_millis() > 49 {
                            last_progress_update = current;
                            (ctx.on_progress)(self, &progress);
                            // Also check if the progress was cancelled.
                            if let Err(state) = ctx.controller.check().await {
                                tracing::warn!(
                                    "operation to copy from {:?} to {:?} cancelled",
                                    self.from,
                                    self.to
                                );
                                return Err::<(), GioCopyError>(GioCopyError::Controller(
                                    OperationError::from_state(state, &ctx.controller),
                                ));
                            }
                        }
                    }
                }

                result = rx => return result.unwrap(),

                _ = task => (),

                _ = until_paused.fuse() => {
                    // Pauses an active copy while the controller state is paused.
                    _ = pause_tx.send(true);
                    ctx.controller.until_unpaused().await;
                    _ = pause_tx.send(false);
                }
            }
        }
    }
}