dosh/src/pty.rs

use anyhow::{Context, Result};
use portable_pty::{Child, CommandBuilder, MasterPty, NativePtySystem, PtySize, PtySystem};
use std::io::{Read, Write};
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::thread;
use tokio::sync::mpsc;

/// Backing for a PTY master held by the server.
///
/// `Owned` means this process spawned the shell as a child and is responsible
/// for it: dropping the handle kills the shell. This is the original,
/// non-persistent model and stays the default.
///
/// `Adopted` means the shell lives in a separate holder process and this handle
/// only borrows the master fd (received over a Unix socket via SCM_RIGHTS).
/// Dropping it must NOT kill the shell — it just closes our copy of the fd and
/// stops the reader thread, leaving the holder + shell alive so a server restart
/// can re-adopt them.
enum Backing {
    Owned {
        child: Mutex<Box<dyn Child + Send + Sync>>,
        _master: Box<dyn MasterPty + Send>,
    },
    Adopted {
        master: Mutex<std::fs::File>,
    },
}

pub struct PtyHandle {
    writer: Arc<Mutex<Box<dyn Write + Send>>>,
    backing: Backing,
}

impl PtyHandle {
    pub fn write_all(&self, bytes: &[u8]) -> Result<()> {
        let mut writer = self.writer.lock().expect("pty writer poisoned");
        writer.write_all(bytes)?;
        writer.flush()?;
        Ok(())
    }

    pub fn resize(&self, cols: u16, rows: u16) -> Result<()> {
        match &self.backing {
            Backing::Owned { _master, .. } => {
                _master.resize(PtySize {
                    rows,
                    cols,
                    pixel_width: 0,
                    pixel_height: 0,
                })?;
            }
            Backing::Adopted { master } => {
                let file = master.lock().expect("pty master poisoned");
                resize_fd(file.as_raw_fd(), cols, rows)?;
            }
        }
        Ok(())
    }

    /// True for a handle backed by a separate holder process. Such a handle must
    /// be detached (not killed) when the server lets go of a session, so the
    /// shell survives a server restart.
    pub fn is_persistent(&self) -> bool {
        matches!(self.backing, Backing::Adopted { .. })
    }

    /// Terminate the shell process backing this PTY and reap it.
    ///
    /// Only meaningful for an `Owned` backing (the server spawned the shell as a
    /// child). For an `Adopted` backing the shell belongs to the holder process,
    /// so this is a no-op here; reaping a persistent session is done by asking
    /// the holder to shut down (see `persist::request_shutdown`).
    pub fn kill(&self) {
        if let Backing::Owned { child, .. } = &self.backing
            && let Ok(mut child) = child.lock()
        {
            let _ = child.kill();
            let _ = child.wait();
        }
    }
}

impl Drop for PtyHandle {
    fn drop(&mut self) {
        // Adopted handles must NOT kill the shell: it lives in the holder so it
        // can outlive this server. Dropping just closes our fd / stops the
        // reader. Owned handles keep the original kill-on-drop behavior.
        if let Backing::Owned { .. } = &self.backing {
            self.kill();
        }
    }
}

fn resize_fd(fd: RawFd, cols: u16, rows: u16) -> Result<()> {
    let winsize = libc::winsize {
        ws_row: rows,
        ws_col: cols,
        ws_xpixel: 0,
        ws_ypixel: 0,
    };
    let rc = unsafe { libc::ioctl(fd, libc::TIOCSWINSZ, &winsize) };
    if rc != 0 {
        return Err(std::io::Error::last_os_error()).context("TIOCSWINSZ");
    }
    Ok(())
}

#[derive(Debug)]
pub struct PtyOutput {
    pub session: String,
    pub bytes: Vec<u8>,
    pub exited: bool,
}

pub fn spawn_pty_session(
    session: String,
    shell: &str,
    cols: u16,
    rows: u16,
    env: &[(String, String)],
    tx: mpsc::UnboundedSender<PtyOutput>,
) -> Result<PtyHandle> {
    let pty_system = NativePtySystem::default();
    let pair = pty_system
        .openpty(PtySize {
            rows,
            cols,
            pixel_width: 0,
            pixel_height: 0,
        })
        .context("open pty")?;
    let cmd = build_shell_command(shell, env);
    let child = pair.slave.spawn_command(cmd).context("spawn shell")?;
    drop(pair.slave);

    let writer = pair.master.take_writer().context("take pty writer")?;
    let reader = pair.master.try_clone_reader().context("clone pty reader")?;
    spawn_reader_thread(session, reader, tx)?;

    Ok(PtyHandle {
        writer: Arc::new(Mutex::new(writer)),
        backing: Backing::Owned {
            child: Mutex::new(child),
            _master: pair.master,
        },
    })
}

/// Build the [`CommandBuilder`] for a dosh shell, identically for the in-process
/// `spawn_pty_session` and the out-of-process holder, so a persistent session's
/// environment matches a non-persistent one.
pub fn build_shell_command(shell: &str, env: &[(String, String)]) -> CommandBuilder {
    let mut cmd = CommandBuilder::new(shell);
    cmd.env("TERM", "xterm-256color");
    cmd.env("COLORTERM", "truecolor");
    for (name, value) in env {
        cmd.env(name, value);
    }
    cmd.env("SHELL", shell);
    if let Some(home) = dirs::home_dir() {
        cmd.env("HOME", home.as_os_str());
        cmd.env("PWD", home.as_os_str());
        cmd.cwd(home.as_os_str());
    } else if let Some(parent) = Path::new(shell).parent() {
        cmd.env("PWD", parent.as_os_str());
    }
    cmd
}

/// Build a [`PtyHandle`] from a master fd received from a holder process.
///
/// `master_fd` is an fd this handle takes ownership of (it is wrapped in a
/// `File` and closed on drop). The shell is NOT a child of this process; it
/// belongs to the holder, so dropping this handle leaves it running.
pub fn adopt_pty_from_fd(
    session: String,
    master_fd: RawFd,
    tx: mpsc::UnboundedSender<PtyOutput>,
) -> Result<PtyHandle> {
    // Take ownership of the fd. A clone gives us an independent reader so the
    // reader thread and the writer/resize side hold separate `File`s and don't
    // get closed out from under each other.
    let master = unsafe { std::fs::File::from_raw_fd(master_fd) };
    let reader_file = master.try_clone().context("clone master for reader")?;
    let writer_file = master.try_clone().context("clone master for writer")?;
    spawn_reader_thread(session, Box::new(reader_file), tx)?;
    Ok(PtyHandle {
        writer: Arc::new(Mutex::new(Box::new(writer_file))),
        backing: Backing::Adopted {
            master: Mutex::new(master),
        },
    })
}

fn spawn_reader_thread(
    session: String,
    mut reader: Box<dyn Read + Send>,
    tx: mpsc::UnboundedSender<PtyOutput>,
) -> Result<()> {
    let reader_session = session.clone();
    thread::Builder::new()
        .name(format!("dosh-pty-{session}"))
        .spawn(move || {
            let mut buf = [0u8; 8192];
            loop {
                match reader.read(&mut buf) {
                    Ok(0) => {
                        let _ = tx.send(PtyOutput {
                            session: reader_session.clone(),
                            bytes: Vec::new(),
                            exited: true,
                        });
                        break;
                    }
                    Ok(n) => {
                        let _ = tx.send(PtyOutput {
                            session: reader_session.clone(),
                            bytes: buf[..n].to_vec(),
                            exited: false,
                        });
                    }
                    Err(_) => {
                        let _ = tx.send(PtyOutput {
                            session: reader_session.clone(),
                            bytes: Vec::new(),
                            exited: true,
                        });
                        break;
                    }
                }
            }
        })
        .context("spawn pty reader")?;
    Ok(())
}