use crate::error::Error;
use nix;
use nix::fcntl::{open, OFlag};
use nix::libc::{STDERR_FILENO, STDIN_FILENO, STDOUT_FILENO};
use nix::pty::{grantpt, posix_openpt, unlockpt, PtyMaster};
pub use nix::sys::{signal, wait};
use nix::sys::{stat, termios};
use nix::unistd::{dup, dup2, fork, setsid, ForkResult, Pid};
use std;
use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::{thread, time};
pub struct PtyProcess {
pub pty: PtyMaster,
pub child_pid: Pid,
kill_timeout: Option<time::Duration>,
}
#[cfg(target_os = "linux")]
use nix::pty::ptsname_r;
#[cfg(target_os = "macos")]
fn ptsname_r(fd: &PtyMaster) -> nix::Result<String> {
use nix::libc::{ioctl, TIOCPTYGNAME};
use std::ffi::CStr;
let mut buf: [i8; 128] = [0; 128];
unsafe {
match ioctl(fd.as_raw_fd(), TIOCPTYGNAME as u64, &mut buf) {
0 => {
let res = CStr::from_ptr(buf.as_ptr()).to_string_lossy().into_owned();
Ok(res)
}
_ => Err(nix::Error::last()),
}
}
}
impl PtyProcess {
pub fn new(mut command: Command) -> Result<Self, Error> {
let master_fd = posix_openpt(OFlag::O_RDWR)?;
grantpt(&master_fd)?;
unlockpt(&master_fd)?;
let slave_name = ptsname_r(&master_fd)?;
match unsafe { fork()? } {
ForkResult::Child => {
setsid()?; let slave_fd = open(
std::path::Path::new(&slave_name),
OFlag::O_RDWR,
stat::Mode::empty(),
)?;
dup2(slave_fd, STDIN_FILENO)?;
dup2(slave_fd, STDOUT_FILENO)?;
dup2(slave_fd, STDERR_FILENO)?;
let mut flags = termios::tcgetattr(STDIN_FILENO)?;
flags.local_flags &= !termios::LocalFlags::ECHO;
termios::tcsetattr(STDIN_FILENO, termios::SetArg::TCSANOW, &flags)?;
command.exec();
Err(Error::Nix(nix::Error::last()))
}
ForkResult::Parent { child: child_pid } => Ok(PtyProcess {
pty: master_fd,
child_pid,
kill_timeout: None,
}),
}
}
pub fn get_file_handle(&self) -> File {
let fd = dup(self.pty.as_raw_fd()).unwrap();
unsafe { File::from_raw_fd(fd) }
}
pub fn set_kill_timeout(&mut self, timeout_ms: Option<u64>) {
self.kill_timeout = timeout_ms.map(time::Duration::from_millis);
}
pub fn status(&self) -> Option<wait::WaitStatus> {
if let Ok(status) = wait::waitpid(self.child_pid, Some(wait::WaitPidFlag::WNOHANG)) {
Some(status)
} else {
None
}
}
pub fn wait(&self) -> Result<wait::WaitStatus, Error> {
wait::waitpid(self.child_pid, None).map_err(Error::from)
}
pub fn exit(&mut self) -> Result<wait::WaitStatus, Error> {
self.kill(signal::SIGTERM).map_err(Error::from)
}
pub fn signal(&mut self, sig: signal::Signal) -> Result<(), Error> {
signal::kill(self.child_pid, sig).map_err(Error::from)
}
pub fn kill(&mut self, sig: signal::Signal) -> Result<wait::WaitStatus, Error> {
let start = time::Instant::now();
loop {
match signal::kill(self.child_pid, sig) {
Ok(_) => {}
Err(nix::errno::Errno::ESRCH) => {
return Ok(wait::WaitStatus::Exited(Pid::from_raw(0), 0))
}
Err(e) => return Err(Error::from(e)),
}
match self.status() {
Some(status) if status != wait::WaitStatus::StillAlive => return Ok(status),
Some(_) | None => thread::sleep(time::Duration::from_millis(100)),
}
if let Some(timeout) = self.kill_timeout {
if start.elapsed() > timeout {
signal::kill(self.child_pid, signal::Signal::SIGKILL).map_err(Error::from)?
}
}
}
}
}
impl Drop for PtyProcess {
fn drop(&mut self) {
if let Some(wait::WaitStatus::StillAlive) = self.status() {
self.exit().expect("cannot exit");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use nix::sys::{signal, wait};
use std::io::prelude::*;
use std::io::{BufReader, LineWriter};
#[test]
fn test_cat() -> std::io::Result<()> {
let process = PtyProcess::new(Command::new("cat")).expect("could not execute cat");
let f = process.get_file_handle();
let mut writer = LineWriter::new(&f);
let mut reader = BufReader::new(&f);
let _ = writer.write(b"hello cat\n")?;
let mut buf = String::new();
reader.read_line(&mut buf)?;
assert_eq!(buf, "hello cat\r\n");
thread::sleep(time::Duration::from_millis(100));
writer.write_all(&[3])?; writer.flush()?;
let should = wait::WaitStatus::Signaled(process.child_pid, signal::Signal::SIGINT, false);
assert_eq!(should, wait::waitpid(process.child_pid, None).unwrap());
Ok(())
}
}