regy 0.1.0

Private-by-default desktop agent for the Regy web interface
use std::fs::File;
use std::os::fd::AsRawFd;
#[cfg(test)]
use std::os::fd::RawFd;
use std::os::unix::process::CommandExt;
use std::process::{Child, Command, Stdio};
use std::sync::Mutex;
use std::sync::atomic::{AtomicI32, Ordering};

use nix::libc;
use nix::pty::{Winsize, openpty};
use tokio::time::{Duration, sleep};

use crate::domain::errors::{AgentError, AgentResult, ErrorCode};
use crate::domain::resize::TerminalSize;
use crate::infrastructure::pty_io::{PtyMaster, set_fd_cloexec};
use crate::infrastructure::pty_process::{configure_child_terminal, terminate_child};

#[derive(Debug)]
pub struct PtySession {
    session_id: String,
    pid: u32,
    pgid: AtomicI32,
    master: PtyMaster,
    child: Mutex<Option<Child>>,
}

impl PtySession {
    pub async fn spawn(
        session_id: String,
        command: Vec<String>,
        cwd: Option<String>,
        size: TerminalSize,
    ) -> AgentResult<Self> {
        let program = command
            .first()
            .cloned()
            .ok_or_else(|| start_err("command must not be empty"))?;
        let winsize = Winsize {
            ws_col: size.cols,
            ws_row: size.rows,
            ws_xpixel: 0,
            ws_ypixel: 0,
        };
        let pty = openpty(Some(&winsize), None).map_err(|e| start_err(e.to_string()))?;
        set_fd_cloexec(pty.master.as_raw_fd()).map_err(|e| start_err(format!("cloexec: {e}")))?;
        let master: File = pty.master.into();
        let slave: File = pty.slave.into();
        let slave_fd = slave.as_raw_fd();
        let master = PtyMaster::new(master).map_err(|e| start_err(e.to_string()))?;

        let mut cmd = Command::new(&program);
        cmd.args(command.into_iter().skip(1));
        if let Some(cwd) = cwd {
            cmd.current_dir(cwd);
        }
        cmd.env("TERM", "xterm-256color");
        unsafe {
            cmd.pre_exec(move || configure_child_terminal(slave_fd));
        }
        let stdin = slave
            .try_clone()
            .map_err(|e| start_err(format!("stdin clone: {e}")))?;
        let stdout = slave
            .try_clone()
            .map_err(|e| start_err(format!("stdout clone: {e}")))?;
        cmd.stdin(Stdio::from(stdin));
        cmd.stdout(Stdio::from(stdout));
        cmd.stderr(Stdio::from(slave));

        let child = cmd.spawn().map_err(|e| start_err(e.to_string()))?;
        let pid = child.id();
        Ok(Self {
            session_id,
            pid,
            pgid: AtomicI32::new(pid as i32),
            master,
            child: Mutex::new(Some(child)),
        })
    }

    pub fn pid(&self) -> u32 {
        self.pid
    }
    pub fn session_id(&self) -> &str {
        &self.session_id
    }

    #[cfg(test)]
    pub(crate) fn master_fd(&self) -> RawFd {
        self.master.raw_fd()
    }

    #[cfg(test)]
    pub(crate) fn child_is_some(&self) -> bool {
        self.child.lock().unwrap().is_some()
    }

    #[cfg(test)]
    pub(crate) fn set_pgid_for_test(&self, pgid: i32) {
        self.pgid.store(pgid, Ordering::SeqCst);
    }

    pub async fn read_next(&self) -> AgentResult<Vec<u8>> {
        self.master.read_next().await
    }

    pub async fn write_input(&self, bytes: &[u8]) -> AgentResult<()> {
        self.master.write_input(bytes).await
    }

    pub async fn resize(&self, size: TerminalSize) -> AgentResult<()> {
        self.master.resize(size)
    }

    pub async fn kill(&self) -> AgentResult<()> {
        let Some(mut child) = self.take_child()? else {
            return Ok(());
        };
        match terminate_child(&mut child, self.pgid.load(Ordering::SeqCst)).await {
            Ok(()) => Ok(()),
            Err(err) => {
                self.restore_child(child)?;
                Err(err)
            }
        }
    }

    pub(crate) async fn reap_exited_after_eof(&self) -> AgentResult<()> {
        for _ in 0..10 {
            if self.reap_exited()? {
                return Ok(());
            }
            sleep(Duration::from_millis(10)).await;
        }
        Ok(())
    }

    pub(crate) fn reap_exited(&self) -> AgentResult<bool> {
        let Some(mut child) = self.take_child()? else {
            return Ok(true);
        };
        if child
            .try_wait()
            .map_err(|e| kill_err(e.to_string()))?
            .is_none()
        {
            self.restore_child(child)?;
            return Ok(false);
        }
        Ok(true)
    }

    fn take_child(&self) -> AgentResult<Option<Child>> {
        self.child
            .lock()
            .map(|mut child| child.take())
            .map_err(|_| kill_err("child lock poisoned"))
    }

    fn restore_child(&self, child: Child) -> AgentResult<()> {
        self.child
            .lock()
            .map(|mut slot| *slot = Some(child))
            .map_err(|_| kill_err("child lock poisoned"))
    }
}

fn start_err(message: impl Into<String>) -> AgentError {
    AgentError::new(ErrorCode::PtyStartFailed, message)
}

fn kill_err(message: impl Into<String>) -> AgentError {
    AgentError::new(ErrorCode::PtyKillFailed, message)
}

pub(crate) fn signal_failure_should_be_ignored(errno: Option<i32>, _leader_exited: bool) -> bool {
    matches!(errno, Some(libc::ESRCH))
}