clan-sidebus/sidebus-broker/src/portal/file_transfer.rs

use std::{collections::HashMap, path::PathBuf};

use tokio::sync::broadcast;
use tokio_stream::StreamExt;
use tracing::{debug, error};
use zbus::{
    Connection, fdo::Result, names::OwnedUniqueName, object_server::SignalEmitter, zvariant,
};

use super::{documents::DocumentsProxy, file_chooser::FileTransformer};

#[derive(Clone)]
pub struct FileTransfer {
    host: FileTransferProxy<'static>,
    file_transformer: FileTransformer,
    tx: broadcast::Sender<ForwarderCommand>,
}

#[derive(Clone, Debug)]
enum ForwarderCommand {
    Add(String, OwnedUniqueName),
    // Remove(String),
}

#[zbus::interface(
    name = "org.freedesktop.portal.FileTransfer",
    proxy(
        default_service = "org.freedesktop.portal.Documents",
        default_path = "/org/freedesktop/portal/documents"
    )
)]
impl FileTransfer {
    async fn add_files(
        &self,
        _key: &str,
        _fds: Vec<zbus::zvariant::Fd<'_>>,
        _options: HashMap<&str, zvariant::Value<'_>>,
    ) -> Result<()> {
        Err(zbus::fdo::Error::NotSupported(
            "Adding files to transfer out is not yet implemented".to_owned(),
        ))
    }

    async fn retrieve_files(
        &self,
        key: &str,
        options: HashMap<&str, zvariant::Value<'_>>,
    ) -> Result<Vec<String>> {
        let host_paths = self.host.retrieve_files(key, options).await?;
        let mut result = Vec::with_capacity(host_paths.len());
        for host_path in host_paths {
            if let Some(guest_path) = self
                .file_transformer
                .add_path_as_doc(PathBuf::from(&host_path))
                .await
            {
                result.push(guest_path.to_string_lossy().into_owned());
            } else {
                debug!("could not expose path {host_path}");
            }
        }
        Ok(result)
    }

    async fn start_transfer(
        &self,
        #[zbus(header)] hdr: zbus::message::Header<'_>,
        options: HashMap<&str, zvariant::Value<'_>>,
    ) -> Result<String> {
        let sender = hdr
            .sender()
            .ok_or_else(|| zbus::Error::MissingField)?
            .to_owned();
        let key = self.host.start_transfer(options).await?;
        debug!("start_transfer: {key}");
        if let Err(err) = self
            .tx
            .send(ForwarderCommand::Add(key.clone(), sender.into()))
        {
            error!("file_transfer internal channel error: {err:?}");
            return Err(zbus::fdo::Error::IOError("channel error".to_owned()));
        }
        Ok(key)
    }

    async fn stop_transfer(&self, key: &str) -> Result<()> {
        debug!("stop_transfer: {key}");
        self.host.stop_transfer(key).await
    }

    #[zbus(signal)]
    async fn transfer_closed(signal_emitter: &SignalEmitter<'_>, key: &str) -> zbus::Result<()>;

    #[zbus(property, name = "version")]
    fn version(&self) -> Result<u32> {
        Ok(1)
    }
}

impl FileTransfer {
    pub async fn new(
        host_session_conn: &Connection,
        priv_conn: &Connection,
        guest_root: PathBuf,
    ) -> Result<Self> {
        let host = FileTransferProxy::builder(host_session_conn)
            .build()
            .await?;
        let docs = DocumentsProxy::builder(priv_conn).build().await?;
        let file_transformer = FileTransformer {
            docs,
            guest_root,
            for_save: false,
            persistent: false,
            directory: false,
        };
        let (tx, _) = broadcast::channel(8);
        Ok(FileTransfer {
            host,
            file_transformer,
            tx,
        })
    }

    pub async fn forward_transfer_closed(
        &self,
        mut signal_emitter: SignalEmitter<'static>,
    ) -> Result<()> {
        let mut stream = self.host.receive_transfer_closed().await?;
        let mut cmds = self.tx.subscribe();
        let mut receivers = HashMap::new();

        loop {
            tokio::select! {
                Ok(cmd) = cmds.recv() => match cmd {
                    ForwarderCommand::Add(key, receiver) => { receivers.insert(key, receiver); },
                    // ForwarderCommand::Remove(key) => { receivers.remove(&key); },
                },
                Some(signal) = stream.next() => {
                    debug!("transfer closed {signal:?}");
                    if let Ok((key,)) = signal.0.deserialize::<(&str,)>() {
                        if let Some(bus_name) = receivers.remove(key) {
                            signal_emitter = signal_emitter.set_destination(zbus::names::BusName::Unique(bus_name.clone().into()));
                            if let Err(err) = FileTransfer::transfer_closed(&signal_emitter, key).await {
                                error!("could not forward signal for key {key}: {err:?}");
                            }
                        } else {
                            error!("got a signal for unknown key {key}");
                        }
                    } else {
                        error!("could not deserialize transfer closed signal");
                    };
                }
            }
        }
    }
}