mprocs-pty 0.7.0

//! Working with pseudo-terminals

use crate::{
  Child, CommandBuilder, MasterPty, PtyPair, PtySize, PtySystem, SlavePty,
};
use anyhow::{bail, Error};
use filedescriptor::FileDescriptor;
use libc::{self, winsize};
use std::io::{Read, Write};
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::os::unix::process::CommandExt;
use std::{io, mem, ptr};

#[derive(Default)]
pub struct UnixPtySystem {}

fn openpty(size: PtySize) -> anyhow::Result<(UnixMasterPty, UnixSlavePty)> {
  let mut master: RawFd = -1;
  let mut slave: RawFd = -1;

  let mut size = winsize {
    ws_row: size.rows,
    ws_col: size.cols,
    ws_xpixel: size.pixel_width,
    ws_ypixel: size.pixel_height,
  };

  let result = unsafe {
    // BSDish systems may require mut pointers to some args
    #[cfg_attr(feature = "cargo-clippy", allow(clippy::unnecessary_mut_passed))]
    libc::openpty(
      &mut master,
      &mut slave,
      ptr::null_mut(),
      ptr::null_mut(),
      &mut size,
    )
  };

  if result != 0 {
    bail!("failed to openpty: {:?}", io::Error::last_os_error());
  }

  let master = UnixMasterPty {
    fd: PtyFd(unsafe { FileDescriptor::from_raw_fd(master) }),
  };
  let slave = UnixSlavePty {
    fd: PtyFd(unsafe { FileDescriptor::from_raw_fd(slave) }),
  };

  // Ensure that these descriptors will get closed when we execute
  // the child process.  This is done after constructing the Pty
  // instances so that we ensure that the Ptys get drop()'d if
  // the cloexec() functions fail (unlikely!).
  cloexec(master.fd.as_raw_fd())?;
  cloexec(slave.fd.as_raw_fd())?;

  Ok((master, slave))
}

impl PtySystem for UnixPtySystem {
  fn openpty(&self, size: PtySize) -> anyhow::Result<PtyPair> {
    let (master, slave) = openpty(size)?;
    Ok(PtyPair {
      master: Box::new(master),
      slave: Box::new(slave),
    })
  }
}

struct PtyFd(pub FileDescriptor);
impl std::ops::Deref for PtyFd {
  type Target = FileDescriptor;
  fn deref(&self) -> &FileDescriptor {
    &self.0
  }
}
impl std::ops::DerefMut for PtyFd {
  fn deref_mut(&mut self) -> &mut FileDescriptor {
    &mut self.0
  }
}

impl Read for PtyFd {
  fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
    match self.0.read(buf) {
      Err(ref e) if e.raw_os_error() == Some(libc::EIO) => {
        // EIO indicates that the slave pty has been closed.
        // Treat this as EOF so that std::io::Read::read_to_string
        // and similar functions gracefully terminate when they
        // encounter this condition
        Ok(0)
      }
      x => x,
    }
  }
}

/// On Big Sur, Cocoa leaks various file descriptors to child processes,
/// so we need to make a pass through the open descriptors beyond just the
/// stdio descriptors and close them all out.
/// This is approximately equivalent to the darwin `posix_spawnattr_setflags`
/// option POSIX_SPAWN_CLOEXEC_DEFAULT which is used as a bit of a cheat
/// on macOS.
/// On Linux, gnome/mutter leak shell extension fds to wezterm too, so we
/// also need to make an effort to clean up the mess.
///
/// This function enumerates the open filedescriptors in the current process
/// and then will forcibly call close(2) on each open fd that is numbered
/// 3 or higher, effectively closing all descriptors except for the stdio
/// streams.
///
/// The implementation of this function relies on `/dev/fd` being available
/// to provide the list of open fds.  Any errors in enumerating or closing
/// the fds are silently ignored.
pub fn close_random_fds() {
  // FreeBSD, macOS and presumably other BSDish systems have /dev/fd as
  // a directory listing the current fd numbers for the process.
  //
  // On Linux, /dev/fd is a symlink to /proc/self/fd
  if let Ok(dir) = std::fs::read_dir("/dev/fd") {
    let mut fds = vec![];
    for entry in dir {
      if let Some(num) = entry
        .ok()
        .map(|e| e.file_name())
        .and_then(|s| s.into_string().ok())
        .and_then(|n| n.parse::<libc::c_int>().ok())
      {
        if num > 2 {
          fds.push(num);
        }
      }
    }
    for fd in fds {
      unsafe {
        libc::close(fd);
      }
    }
  }
}

impl PtyFd {
  fn resize(&self, size: PtySize) -> Result<(), Error> {
    let ws_size = winsize {
      ws_row: size.rows,
      ws_col: size.cols,
      ws_xpixel: size.pixel_width,
      ws_ypixel: size.pixel_height,
    };

    if unsafe {
      libc::ioctl(
        self.0.as_raw_fd(),
        libc::TIOCSWINSZ as _,
        &ws_size as *const _,
      )
    } != 0
    {
      bail!(
        "failed to ioctl(TIOCSWINSZ): {:?}",
        io::Error::last_os_error()
      );
    }

    Ok(())
  }

  fn get_size(&self) -> Result<PtySize, Error> {
    let mut size: winsize = unsafe { mem::zeroed() };
    if unsafe {
      libc::ioctl(
        self.0.as_raw_fd(),
        libc::TIOCGWINSZ as _,
        &mut size as *mut _,
      )
    } != 0
    {
      bail!(
        "failed to ioctl(TIOCGWINSZ): {:?}",
        io::Error::last_os_error()
      );
    }
    Ok(PtySize {
      rows: size.ws_row,
      cols: size.ws_col,
      pixel_width: size.ws_xpixel,
      pixel_height: size.ws_ypixel,
    })
  }

  fn spawn_command(
    &self,
    builder: CommandBuilder,
  ) -> anyhow::Result<std::process::Child> {
    let configured_umask = builder.umask;

    let mut cmd = builder.as_command()?;
    let controlling_tty = builder.get_controlling_tty();

    unsafe {
      cmd
        .stdin(self.as_stdio()?)
        .stdout(self.as_stdio()?)
        .stderr(self.as_stdio()?)
        .pre_exec(move || {
          // Clean up a few things before we exec the program
          // Clear out any potentially problematic signal
          // dispositions that we might have inherited
          for signo in &[
            libc::SIGCHLD,
            libc::SIGHUP,
            libc::SIGINT,
            libc::SIGQUIT,
            libc::SIGTERM,
            libc::SIGALRM,
          ] {
            libc::signal(*signo, libc::SIG_DFL);
          }

          // Establish ourselves as a session leader.
          if libc::setsid() == -1 {
            return Err(io::Error::last_os_error());
          }

          // Clippy wants us to explicitly cast TIOCSCTTY using
          // type::from(), but the size and potentially signedness
          // are system dependent, which is why we're using `as _`.
          // Suppress this lint for this section of code.
          #[cfg_attr(feature = "cargo-clippy", allow(clippy::cast_lossless))]
          if controlling_tty {
            // Set the pty as the controlling terminal.
            // Failure to do this means that delivery of
            // SIGWINCH won't happen when we resize the
            // terminal, among other undesirable effects.
            if libc::ioctl(0, libc::TIOCSCTTY as _, 0) == -1 {
              return Err(io::Error::last_os_error());
            }
          }

          close_random_fds();

          if let Some(mask) = configured_umask {
            libc::umask(mask);
          }

          Ok(())
        })
    };

    let mut child = cmd.spawn()?;

    // Ensure that we close out the slave fds that Child retains;
    // they are not what we need (we need the master side to reference
    // them) and won't work in the usual way anyway.
    // In practice these are None, but it seems best to be move them
    // out in case the behavior of Command changes in the future.
    child.stdin.take();
    child.stdout.take();
    child.stderr.take();

    Ok(child)
  }
}

/// Represents the master end of a pty.
/// The file descriptor will be closed when the Pty is dropped.
struct UnixMasterPty {
  fd: PtyFd,
}

/// Represents the slave end of a pty.
/// The file descriptor will be closed when the Pty is dropped.
struct UnixSlavePty {
  fd: PtyFd,
}

/// Helper function to set the close-on-exec flag for a raw descriptor
fn cloexec(fd: RawFd) -> Result<(), Error> {
  let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
  if flags == -1 {
    bail!(
      "fcntl to read flags failed: {:?}",
      io::Error::last_os_error()
    );
  }
  let result =
    unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) };
  if result == -1 {
    bail!(
      "fcntl to set CLOEXEC failed: {:?}",
      io::Error::last_os_error()
    );
  }
  Ok(())
}

impl SlavePty for UnixSlavePty {
  fn spawn_command(
    &self,
    builder: CommandBuilder,
  ) -> Result<Box<dyn Child + Send + Sync>, Error> {
    Ok(Box::new(self.fd.spawn_command(builder)?))
  }
}

impl MasterPty for UnixMasterPty {
  fn resize(&self, size: PtySize) -> Result<(), Error> {
    self.fd.resize(size)
  }

  fn get_size(&self) -> Result<PtySize, Error> {
    self.fd.get_size()
  }

  fn try_clone_reader(&self) -> Result<Box<dyn Read + Send>, Error> {
    let fd = PtyFd(self.fd.try_clone()?);
    Ok(Box::new(fd))
  }

  fn try_clone_writer(&self) -> Result<Box<dyn Write + Send>, Error> {
    let fd = PtyFd(self.fd.try_clone()?);
    Ok(Box::new(UnixMasterPty { fd }))
  }

  fn process_group_leader(&self) -> Option<libc::pid_t> {
    match unsafe { libc::tcgetpgrp(self.fd.0.as_raw_fd()) } {
      pid if pid > 0 => Some(pid),
      _ => None,
    }
  }
}

impl Write for UnixMasterPty {
  fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
    self.fd.write(buf)
  }
  fn flush(&mut self) -> Result<(), io::Error> {
    self.fd.flush()
  }
}