beer

//! 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");

    }

}