beer/src/pty.rs

//! Pseudo-terminal: open a master/slave pair and run the user's shell on it.

use std::ffi::OsStr;
use std::io;
use std::os::fd::{AsFd, OwnedFd};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Child, Command, ExitStatus, Stdio};

use anyhow::Context;
use rustix::pty::{OpenptFlags, grantpt, ioctl_tiocgptpeer, openpt, unlockpt};
use rustix::termios::{Winsize, tcsetwinsize};

/// A running child process attached to a PTY master.
#[derive(Debug)]
pub struct Pty {
    master: OwnedFd,
    child: Child,
}

impl Pty {
    /// Open a PTY, size it to `cols`x`rows`, and exec the user's login shell on
    /// the slave end with `TERM=beer`.
    pub fn spawn(cols: u16, rows: u16) -> anyhow::Result<Self> {
        let master = openpt(OpenptFlags::RDWR | OpenptFlags::NOCTTY | OpenptFlags::CLOEXEC)
            .context("open pty master")?;
        grantpt(&master).context("grantpt")?;
        unlockpt(&master).context("unlockpt")?;
        let slave = ioctl_tiocgptpeer(&master, OpenptFlags::RDWR | OpenptFlags::NOCTTY)
            .context("open pty slave")?;

        set_winsize(&master, cols, rows)?;

        let shell = std::env::var_os("SHELL").unwrap_or_else(|| "/bin/sh".into());
        let argv0 = login_argv0(&shell);

        // Hand the slave to the child's stdio. try_clone gives O_CLOEXEC dups, so
        // the parent's copies and the controlling-tty handle vanish at exec.
        let ctty = slave.try_clone().context("dup slave")?;
        let (stdin, stdout, stderr) = (
            slave.try_clone().context("dup slave")?,
            slave.try_clone().context("dup slave")?,
            slave,
        );

        let mut cmd = Command::new(&shell);
        cmd.arg0(&argv0)
            .env("TERM", "beer")
            .env_remove("COLUMNS")
            .env_remove("LINES")
            .env_remove("TERMCAP")
            .stdin(Stdio::from(stdin))
            .stdout(Stdio::from(stdout))
            .stderr(Stdio::from(stderr));

        // SAFETY: setsid and the TIOCSCTTY ioctl are async-signal-safe raw
        // syscalls; ctty is a valid fd captured by move. We touch no parent
        // heap state, satisfying pre_exec's contract.
        unsafe {
            cmd.pre_exec(move || {
                rustix::process::setsid()?;
                rustix::process::ioctl_tiocsctty(&ctty)?;
                Ok(())
            });
        }

        let child = cmd.spawn().context("spawn shell")?;
        Ok(Self { master, child })
    }

    /// The PTY master, for reading child output and writing input.
    pub fn master(&self) -> &OwnedFd {
        &self.master
    }

    /// Reap the child if it has exited.
    pub fn wait(&mut self) -> io::Result<ExitStatus> {
        self.child.wait()
    }
}

fn set_winsize(master: &OwnedFd, cols: u16, rows: u16) -> anyhow::Result<()> {
    let ws = Winsize {
        ws_row: rows,
        ws_col: cols,
        ws_xpixel: 0,
        ws_ypixel: 0,
    };
    tcsetwinsize(master.as_fd(), ws).context("set pty winsize")
}

/// Login-shell argv[0] is the shell's basename with a leading '-'.
fn login_argv0(shell: &OsStr) -> std::ffi::OsString {
    let name = Path::new(shell)
        .file_name()
        .unwrap_or_else(|| OsStr::new("sh"));
    let mut argv0 = std::ffi::OsString::from("-");
    argv0.push(name);
    argv0
}

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

    #[test]
    fn argv0_is_dash_prefixed_basename() {
        assert_eq!(login_argv0(OsStr::new("/usr/bin/bash")), "-bash");
        assert_eq!(login_argv0(OsStr::new("zsh")), "-zsh");
        assert_eq!(login_argv0(OsStr::new("")), "-sh");
    }
}