ghostty-rs/crates/ghostty-os/src/flatpak.rs

use std::collections::HashMap;
use std::os::fd::{AsFd, OwnedFd};
use std::sync::{Arc, Condvar, Mutex};

use zbus::blocking::{Connection, MessageIterator};
use zbus::zvariant::Fd;
use zbus::{MatchRule, message::Type as MessageType};

const BUS_NAME: &str = "org.freedesktop.Flatpak";
const OBJECT_PATH: &str = "/org/freedesktop/Flatpak/Development";
const INTERFACE: &str = "org.freedesktop.Flatpak.Development";

/// Returns true if we're running in a Flatpak environment.
pub fn is_flatpak() -> bool {
    std::path::Path::new("/.flatpak-info").exists()
}

/// Errors that are possible from us.
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum Error {
    #[error("flatpak host command must be started")]
    FlatpakMustBeStarted,
    #[error("flatpak spawn failed")]
    FlatpakSpawnFail,
    #[error("flatpak setup failed")]
    FlatpakSetupFail,
    #[error("flatpak rpc failed")]
    FlatpakRPCFail,
}

/// State the process is in.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
    /// Initial state
    Init,
    /// Error starting. The error message is only available via logs.
    Err,
    /// Process started with the given pid on the host.
    Started { pid: u32 },
    /// Process exited
    Exited { pid: u32, status: u8 },
}

/// State shared with the dedicated thread, protected by the lock and
/// condition variable.
struct Shared {
    state: Mutex<State>,
    cv: Condvar,
}

impl Shared {
    fn update(&self, state: State) {
        let mut guard = self
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        *guard = state;
        self.cv.notify_all();
    }
}

/// A struct to help execute commands on the host via the
/// org.freedesktop.Flatpak.Development DBus module. This uses `zbus`
/// (blocking API) rather than GIO/GLib.
///
/// This always spawns its own thread to wait for the exit signal. This makes
/// it easy for the command to behave synchronously similar to
/// `std::process::Child`.
///
/// There are lots of chances for low-hanging improvements here (automatic
/// pipes, /dev/null, etc.) but this was purpose built for a specific need so
/// it doesn't have all of those.
///
/// NOTE: this can't be exercised in the test environment (needs a Flatpak
/// session bus with the Development interface); it is a best-effort port.
pub struct FlatpakHostCommand {
    /// Argv are the arguments to call on the host with argv[0] being
    /// the command to execute.
    pub argv: Vec<String>,

    /// The cwd for the new process. If this is not set then it will use
    /// the current cwd of the calling process.
    pub cwd: Option<String>,

    /// Environment variables for the child process. If this is None, this
    /// does not send any environment variables.
    pub env: Option<HashMap<String, String>>,

    /// File descriptors to send to the child process. It is up to the
    /// caller to create the file descriptors and set them up. They are
    /// duplicated when spawning so the caller keeps ownership.
    pub stdin: OwnedFd,
    pub stdout: OwnedFd,
    pub stderr: OwnedFd,

    /// Connection kept for sending signals.
    shared: Arc<Shared>,
}

impl FlatpakHostCommand {
    pub fn new(argv: Vec<String>, stdin: OwnedFd, stdout: OwnedFd, stderr: OwnedFd) -> Self {
        Self {
            argv,
            cwd: None,
            env: None,
            stdin,
            stdout,
            stderr,
            shared: Arc::new(Shared {
                state: Mutex::new(State::Init),
                cv: Condvar::new(),
            }),
        }
    }

    fn state(&self) -> State {
        *self
            .shared
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Spawn the command. This will start the host command. On return,
    /// the pid will be available. This must only be called with the
    /// state in "init".
    pub fn spawn(&mut self) -> Result<u32, Error> {
        let request = Request {
            argv: self.argv.clone(),
            cwd: self.cwd.clone(),
            env: self.env.clone().unwrap_or_default(),
            fds: [
                self.stdin
                    .try_clone()
                    .map_err(|_| Error::FlatpakSetupFail)?,
                self.stdout
                    .try_clone()
                    .map_err(|_| Error::FlatpakSetupFail)?,
                self.stderr
                    .try_clone()
                    .map_err(|_| Error::FlatpakSetupFail)?,
            ],
        };
        let shared = Arc::clone(&self.shared);

        // We don't track this thread, it will terminate on its own on
        // command exit.
        std::thread::Builder::new()
            .name("flatpak-host-command".into())
            .spawn(move || thread_main(&shared, request))
            .map_err(|_| Error::FlatpakSpawnFail)?;

        // Wait for the process to start or error.
        let mut guard = self
            .shared
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        while *guard == State::Init {
            guard = self
                .shared
                .cv
                .wait(guard)
                .unwrap_or_else(std::sync::PoisonError::into_inner);
        }

        match *guard {
            State::Init => unreachable!(),
            State::Err => Err(Error::FlatpakSpawnFail),
            State::Started { pid } | State::Exited { pid, .. } => Ok(pid),
        }
    }

    /// Wait for the process to end and return the exit status. This
    /// can only be called ONCE. Once this returns, the state is reset.
    pub fn wait(&mut self) -> Result<u8, Error> {
        let mut guard = self
            .shared
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        loop {
            match *guard {
                State::Init => return Err(Error::FlatpakMustBeStarted),
                State::Err => return Err(Error::FlatpakSpawnFail),
                State::Started { .. } => {}
                State::Exited { status, .. } => {
                    *guard = State::Init;
                    self.shared.cv.notify_all();
                    return Ok(status);
                }
            }

            guard = self
                .shared
                .cv
                .wait(guard)
                .unwrap_or_else(std::sync::PoisonError::into_inner);
        }
    }

    /// Wait for the process to end asynchronously. The callback is invoked
    /// from a helper thread with the result. This can only be called ONCE.
    ///
    /// This replaces `waitXev` (libxev completion) of the Zig version.
    pub fn wait_callback(
        mut self,
        callback: impl FnOnce(Result<u8, Error>) + Send + 'static,
    ) -> std::io::Result<()> {
        std::thread::Builder::new()
            .name("flatpak-host-wait".into())
            .spawn(move || callback(self.wait()))?;
        Ok(())
    }

    /// Send a signal to the started command. This does nothing if the
    /// command is not in the started state.
    pub fn signal(&self, sig: u8, pg: bool) -> Result<(), Error> {
        let State::Started { pid } = self.state() else {
            return Ok(());
        };

        // Get our bus connection.
        let bus = Connection::session().map_err(|e| {
            log::warn!("signal error getting bus: {e}");
            Error::FlatpakSetupFail
        })?;

        if let Err(e) = bus.call_method(
            Some(BUS_NAME),
            OBJECT_PATH,
            Some(INTERFACE),
            "HostCommandSignal",
            &(pid, u32::from(sig), pg),
        ) {
            log::warn!("signal send error: {e}");
        }
        Ok(())
    }
}

struct Request {
    argv: Vec<String>,
    cwd: Option<String>,
    env: HashMap<String, String>,
    fds: [OwnedFd; 3],
}

fn thread_main(shared: &Shared, request: Request) {
    // Get our bus connection. This has to remain active until we exit
    // the thread otherwise our signals won't be called.
    let bus = match Connection::session() {
        Ok(b) => b,
        Err(e) => {
            log::warn!("spawn error getting bus: {e}");
            shared.update(State::Err);
            return;
        }
    };

    // Spawn the command first. This will setup all our IO.
    match start(shared, &bus, request) {
        Ok(iter_and_pid) => {
            let (iter, pid) = iter_and_pid;
            wait_for_exit(shared, iter, pid);
        }
        Err(err) => {
            log::warn!("error starting host command: {err}");
            shared.update(State::Err);
        }
    }
}

/// Start the command. This will start the host command and set the
/// pid on success. This will not wait for completion.
fn start(
    shared: &Shared,
    bus: &Connection,
    request: Request,
) -> Result<(MessageIterator, u32), Error> {
    // Build our file descriptor map: 0 -> stdin, 1 -> stdout, 2 -> stderr.
    let mut fds: HashMap<u32, Fd<'_>> = HashMap::new();
    for (i, fd) in request.fds.iter().enumerate() {
        fds.insert(i as u32, Fd::from(fd.as_fd()));
    }

    // Build our args (arrays of NUL-terminated byte strings).
    let args: Vec<Vec<u8>> = request
        .argv
        .iter()
        .map(|a| {
            let mut v = a.as_bytes().to_vec();
            v.push(0);
            v
        })
        .collect();

    // Get the cwd in case we don't have ours set.
    let cwd = match request.cwd {
        Some(c) => c,
        None => std::env::current_dir()
            .map_err(|_| Error::FlatpakSetupFail)?
            .to_string_lossy()
            .into_owned(),
    };
    let mut cwd = cwd.into_bytes();
    cwd.push(0);

    // Subscribe to exit notifications (before the call so we can't miss it).
    let rule = MatchRule::builder()
        .msg_type(MessageType::Signal)
        .sender(BUS_NAME)
        .and_then(|b| b.interface(INTERFACE))
        .and_then(|b| b.member("HostCommandExited"))
        .and_then(|b| b.path(OBJECT_PATH))
        .map(zbus::match_rule::Builder::build)
        .map_err(|e| {
            log::warn!("error building match rule: {e}");
            Error::FlatpakSetupFail
        })?;
    let iter = MessageIterator::for_match_rule(rule, bus, None).map_err(|e| {
        log::warn!("error subscribing to HostCommandExited: {e}");
        Error::FlatpakSetupFail
    })?;

    // Go!
    let reply = bus
        .call_method(
            Some(BUS_NAME),
            OBJECT_PATH,
            Some(INTERFACE),
            "HostCommand",
            &(cwd, args, fds, request.env, 0u32),
        )
        .map_err(|e| {
            log::warn!("Flatpak.HostCommand failed: {e}");
            Error::FlatpakRPCFail
        })?;

    let pid: u32 = reply.body().deserialize().map_err(|e| {
        log::warn!("Flatpak.HostCommand bad reply: {e}");
        Error::FlatpakRPCFail
    })?;
    log::debug!("HostCommand started pid={pid}");

    shared.update(State::Started { pid });
    Ok((iter, pid))
}

/// Block on the exit signal for our pid and publish the exit status.
fn wait_for_exit(shared: &Shared, iter: MessageIterator, pid: u32) {
    for msg in iter {
        let Ok(msg) = msg else { continue };
        let Ok((exit_pid, exit_status_raw)) = msg.body().deserialize::<(u32, u32)>() else {
            continue;
        };
        if exit_pid != pid {
            continue;
        }

        // posix.W.EXITSTATUS
        let status = ((exit_status_raw >> 8) & 0xff) as u8;
        shared.update(State::Exited { pid, status });
        log::debug!("HostCommand exited pid={pid} status={status}");
        return;
    }

    // The connection closed without an exit signal.
    log::warn!("host command signal stream ended without exit pid={pid}");
    shared.update(State::Err);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn is_flatpak_matches_flatpak_info() {
        assert_eq!(
            is_flatpak(),
            std::path::Path::new("/.flatpak-info").exists()
        );
    }

    #[test]
    fn wait_before_spawn_errors() {
        let (r, w) = crate::pipe::pipe().unwrap();
        let w2 = w.try_clone().unwrap();
        let mut cmd = FlatpakHostCommand::new(vec!["true".into()], r, w, w2);
        assert_eq!(cmd.wait(), Err(Error::FlatpakMustBeStarted));
        // Signalling an unstarted command is a no-op.
        assert_eq!(cmd.signal(15, false), Ok(()));
    }

    #[test]
    fn exit_status_extraction() {
        // Mirrors posix.W.EXITSTATUS on a raw wait status.
        assert_eq!(((3u32 << 8) >> 8) & 0xff, 3);
    }
}