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, PathBuf};
use std::sync::{Arc, Mutex};
use std::thread;
use tokio::sync::mpsc;

// Keep live terminal output comfortably below common path MTUs after Dosh's
// protocol header, AEAD tag, UDP/IP headers, and bincode framing. Full-screen
// TUIs often write several KiB on the first draw; sending that as one UDP
// datagram can fragment and vanish, leaving only a blank alternate screen.
const PTY_OUTPUT_CHUNK_BYTES: usize = 1024;

/// 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,
        },
    })
}

/// Whether this host has a terminfo entry for `term`, searching the standard
/// ncurses directories. Both the legacy single-letter (`x/xterm`) and the
/// hashed (`78/xterm`) subdirectory layouts are checked.
fn terminfo_available(term: &str) -> bool {
    let Some(first) = term.chars().next() else {
        return false;
    };
    let letter = first.to_string();
    let hashed = format!("{:x}", first as u32);
    let mut dirs: Vec<PathBuf> = Vec::new();
    if let Ok(t) = std::env::var("TERMINFO")
        && !t.is_empty()
    {
        dirs.push(PathBuf::from(t));
    }
    if let Some(home) = dirs::home_dir() {
        dirs.push(home.join(".terminfo"));
    }
    if let Ok(td) = std::env::var("TERMINFO_DIRS") {
        for d in td.split(':').filter(|d| !d.is_empty()) {
            dirs.push(PathBuf::from(d));
        }
    }
    for d in [
        "/etc/terminfo",
        "/lib/terminfo",
        "/usr/share/terminfo",
        "/usr/lib/terminfo",
        "/usr/share/lib/terminfo",
    ] {
        dirs.push(PathBuf::from(d));
    }
    dirs.iter()
        .any(|dir| dir.join(&letter).join(term).exists() || dir.join(&hashed).join(term).exists())
}

/// 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("COLORTERM", "truecolor");
    for (name, value) in env {
        cmd.env(name, value);
    }
    // The client's TERM is propagated, but a server that lacks that terminal's
    // terminfo entry (e.g. xterm-ghostty, xterm-kitty) breaks ncurses apps like
    // tmux/vim with "missing or unsuitable terminal". Keep the requested TERM
    // only when this host actually has its terminfo; otherwise fall back to a
    // universally available entry so remote apps always work.
    let term = match env
        .iter()
        .find(|(n, _)| n == "TERM")
        .map(|(_, v)| v.as_str())
    {
        Some(requested) if terminfo_available(requested) => requested.to_string(),
        _ => "xterm-256color".to_string(),
    };
    cmd.env("TERM", term);
    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) => {
                        for chunk in buf[..n].chunks(PTY_OUTPUT_CHUNK_BYTES) {
                            let _ = tx.send(PtyOutput {
                                session: reader_session.clone(),
                                bytes: chunk.to_vec(),
                                exited: false,
                            });
                        }
                    }
                    Err(_) => {
                        let _ = tx.send(PtyOutput {
                            session: reader_session.clone(),
                            bytes: Vec::new(),
                            exited: true,
                        });
                        break;
                    }
                }
            }
        })
        .context("spawn pty reader")?;
    Ok(())
}

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

    const _: () = assert!(PTY_OUTPUT_CHUNK_BYTES <= 1200);

    #[test]
    fn terminfo_available_detects_known_and_unknown() {
        // A near-universal entry should be present on any host with ncurses.
        assert!(terminfo_available("xterm") || terminfo_available("xterm-256color"));
        // Bogus / empty names must report missing so we fall back.
        assert!(!terminfo_available("definitely-not-a-real-terminal-xyz"));
        assert!(!terminfo_available(""));
    }
}