diff --git a/crates/rattler_pty/CHANGELOG.md b/crates/rattler_pty/CHANGELOG.md new file mode 100644 index 000000000..d6637e049 --- /dev/null +++ b/crates/rattler_pty/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/crates/rattler_pty/Cargo.toml b/crates/rattler_pty/Cargo.toml new file mode 100644 index 000000000..843b4067f --- /dev/null +++ b/crates/rattler_pty/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "rattler_pty" +version = "0.1.0" +description = "A crate to create pty" +categories.workspace = true +homepage.workspace = true +repository.workspace = true +license.workspace = true +edition.workspace = true +readme.workspace = true + +[dependencies] + +[target.'cfg(unix)'.dependencies] +libc = { workspace = true } +nix = { version = "0.29.0", features = ["fs", "signal", "term", "poll"] } +signal-hook = "0.3.17" diff --git a/crates/rattler_pty/src/lib.rs b/crates/rattler_pty/src/lib.rs new file mode 100644 index 000000000..b13064bf7 --- /dev/null +++ b/crates/rattler_pty/src/lib.rs @@ -0,0 +1,2 @@ +#[cfg(unix)] +pub mod unix; diff --git a/crates/rattler_pty/src/unix/mod.rs b/crates/rattler_pty/src/unix/mod.rs new file mode 100644 index 000000000..39156dae5 --- /dev/null +++ b/crates/rattler_pty/src/unix/mod.rs @@ -0,0 +1,6 @@ +mod pty_process; +mod pty_session; + +pub use pty_process::PtyProcess; +pub use pty_process::PtyProcessOptions; +pub use pty_session::PtySession; diff --git a/crates/rattler_pty/src/unix/pty_process.rs b/crates/rattler_pty/src/unix/pty_process.rs new file mode 100644 index 000000000..e6c2b88bd --- /dev/null +++ b/crates/rattler_pty/src/unix/pty_process.rs @@ -0,0 +1,249 @@ +pub use nix::sys::{signal, wait}; +use nix::{ + self, + fcntl::{open, OFlag}, + libc::{STDERR_FILENO, STDIN_FILENO, STDOUT_FILENO}, + pty::{grantpt, posix_openpt, unlockpt, PtyMaster, Winsize}, + sys::termios::{InputFlags, Termios}, + sys::{stat, termios}, + unistd::{close, dup, dup2, fork, setsid, ForkResult, Pid}, +}; +use std::os::fd::AsFd; +use std::{ + self, + fs::File, + io, + os::unix::{ + io::{AsRawFd, FromRawFd}, + process::CommandExt, + }, + process::Command, + thread, time, +}; + +#[cfg(target_os = "linux")] +use nix::pty::ptsname_r; + +/// Start a process in a forked tty so you can interact with it the same as you would +/// within a terminal +/// +/// The process and pty session are killed upon dropping PtyProcess +pub struct PtyProcess { + pub pty: PtyMaster, + pub child_pid: Pid, + kill_timeout: Option, +} + +#[cfg(target_os = "macos")] +/// ptsname_r is a linux extension but ptsname isn't thread-safe +/// instead of using a static mutex this calls ioctl with TIOCPTYGNAME directly +/// based on https://blog.tarq.io/ptsname-on-osx-with-rust/ +fn ptsname_r(fd: &PtyMaster) -> nix::Result { + use nix::libc::{ioctl, TIOCPTYGNAME}; + use std::ffi::CStr; + + // the buffer size on OSX is 128, defined by sys/ttycom.h + 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()), + } + } +} + +#[derive(Default)] +pub struct PtyProcessOptions { + pub echo: bool, + pub window_size: Option, +} + +impl PtyProcess { + /// Start a process in a forked pty + pub fn new(mut command: Command, opts: PtyProcessOptions) -> nix::Result { + // Open a new PTY master + let master_fd = posix_openpt(OFlag::O_RDWR)?; + + // Allow a slave to be generated for it + grantpt(&master_fd)?; + unlockpt(&master_fd)?; + + // on Linux this is the libc function, on OSX this is our implementation of ptsname_r + let slave_name = ptsname_r(&master_fd)?; + + // Get the current window size if it was not specified + let window_size = opts.window_size.unwrap_or_else(|| { + // find current window size with ioctl + let mut size: libc::winsize = unsafe { std::mem::zeroed() }; + // Query the terminal dimensions + unsafe { libc::ioctl(io::stdout().as_raw_fd(), libc::TIOCGWINSZ, &mut size) }; + size + }); + + match unsafe { fork()? } { + ForkResult::Child => { + // Avoid leaking master fd + close(master_fd.as_raw_fd())?; + + setsid()?; // create new session with child as session leader + let slave_fd = open( + std::path::Path::new(&slave_name), + OFlag::O_RDWR, + stat::Mode::empty(), + )?; + + // assign stdin, stdout, stderr to the tty, just like a terminal does + dup2(slave_fd, STDIN_FILENO)?; + dup2(slave_fd, STDOUT_FILENO)?; + dup2(slave_fd, STDERR_FILENO)?; + + // Avoid leaking slave fd + if slave_fd > STDERR_FILENO { + close(slave_fd)?; + } + + // set echo off + set_echo(io::stdin(), opts.echo)?; + set_window_size(io::stdout().as_raw_fd(), window_size)?; + + // let mut flags = termios::tcgetattr(io::stdin())?; + // flags.local_flags |= termios::LocalFlags::ECHO; + // termios::tcsetattr(io::stdin(), termios::SetArg::TCSANOW, &flags)?; + + let _ = command.exec(); + Err(nix::Error::last()) + } + ForkResult::Parent { child: child_pid } => Ok(PtyProcess { + pty: master_fd, + child_pid, + kill_timeout: None, + }), + } + } + + /// Get handle to pty fork for reading/writing + pub fn get_file_handle(&self) -> nix::Result { + // needed because otherwise fd is closed both by dropping process and reader/writer + let fd = dup(self.pty.as_raw_fd())?; + unsafe { Ok(File::from_raw_fd(fd)) } + } + + /// Get status of child process, non-blocking. + /// + /// This method runs waitpid on the process. + /// This means: If you ran `exit()` before or `status()` this method will + /// return `None` + pub fn status(&self) -> Option { + if let Ok(status) = wait::waitpid(self.child_pid, Some(wait::WaitPidFlag::WNOHANG)) { + Some(status) + } else { + None + } + } + + /// Regularly exit the process, this method is blocking until the process is dead + pub fn exit(&mut self) -> nix::Result { + self.kill(signal::SIGTERM) + } + + /// Kill the process with a specific signal. This method blocks, until the process is dead + /// + /// repeatedly sends SIGTERM to the process until it died, + /// the pty session is closed upon dropping PtyMaster, + /// so we don't need to explicitly do that here. + /// + /// if `kill_timeout` is set and a repeated sending of signal does not result in the process + /// being killed, then `kill -9` is sent after the `kill_timeout` duration has elapsed. + pub fn kill(&mut self, sig: signal::Signal) -> nix::Result { + let start = time::Instant::now(); + loop { + match signal::kill(self.child_pid, sig) { + Ok(_) => {} + // process was already killed before -> ignore + Err(nix::errno::Errno::ESRCH) => { + return Ok(wait::WaitStatus::Exited(Pid::from_raw(0), 0)); + } + Err(e) => return Err(e), + } + + match self.status() { + Some(status) if status != wait::WaitStatus::StillAlive => return Ok(status), + Some(_) | None => thread::sleep(time::Duration::from_millis(100)), + } + // kill -9 if timeout is reached + if let Some(timeout) = self.kill_timeout { + if start.elapsed() > timeout { + signal::kill(self.child_pid, signal::Signal::SIGKILL)? + } + } + } + } + + /// Set raw mode on stdin and return the original mode + pub fn set_raw(&self) -> nix::Result { + let original_mode = termios::tcgetattr(io::stdin())?; + let mut raw_mode = original_mode.clone(); + raw_mode.input_flags.remove( + InputFlags::BRKINT + | InputFlags::ICRNL + | InputFlags::INPCK + | InputFlags::ISTRIP + | InputFlags::IXON, + ); + raw_mode.output_flags.remove(termios::OutputFlags::OPOST); + raw_mode + .control_flags + .remove(termios::ControlFlags::CSIZE | termios::ControlFlags::PARENB); + raw_mode.control_flags.insert(termios::ControlFlags::CS8); + raw_mode.local_flags.remove( + termios::LocalFlags::ECHO + | termios::LocalFlags::ICANON + | termios::LocalFlags::IEXTEN + | termios::LocalFlags::ISIG, + ); + + raw_mode.control_chars[termios::SpecialCharacterIndices::VMIN as usize] = 1; + raw_mode.control_chars[termios::SpecialCharacterIndices::VTIME as usize] = 0; + + termios::tcsetattr(io::stdin(), termios::SetArg::TCSAFLUSH, &raw_mode)?; + + Ok(original_mode) + } + + pub fn set_mode(&self, original_mode: Termios) -> nix::Result<()> { + termios::tcsetattr(io::stdin(), termios::SetArg::TCSAFLUSH, &original_mode)?; + Ok(()) + } + + pub fn set_window_size(&self, window_size: Winsize) -> nix::Result<()> { + set_window_size(self.pty.as_raw_fd(), window_size) + } +} + +pub fn set_window_size(raw_fd: i32, window_size: Winsize) -> nix::Result<()> { + unsafe { libc::ioctl(raw_fd, nix::libc::TIOCSWINSZ, &window_size) }; + Ok(()) +} + +pub fn set_echo(fd: Fd, echo: bool) -> nix::Result<()> { + let mut flags = termios::tcgetattr(&fd)?; + if echo { + flags.local_flags.insert(termios::LocalFlags::ECHO); + } else { + flags.local_flags.remove(termios::LocalFlags::ECHO); + } + termios::tcsetattr(&fd, termios::SetArg::TCSANOW, &flags)?; + Ok(()) +} + +impl Drop for PtyProcess { + fn drop(&mut self) { + if let Some(wait::WaitStatus::StillAlive) = self.status() { + self.exit().expect("cannot exit"); + } + } +} diff --git a/crates/rattler_pty/src/unix/pty_session.rs b/crates/rattler_pty/src/unix/pty_session.rs new file mode 100644 index 000000000..057fd579d --- /dev/null +++ b/crates/rattler_pty/src/unix/pty_session.rs @@ -0,0 +1,177 @@ +use super::PtyProcess; +use crate::unix::pty_process::PtyProcessOptions; +use libc::SIGWINCH; +use nix::sys::select::FdSet; +use nix::{ + errno::Errno, + sys::{select, time::TimeVal, wait::WaitStatus}, +}; +use signal_hook::iterator::Signals; +use std::{ + fs::File, + io::{self, Read, Write}, + os::fd::AsFd, + process::Command, +}; + +pub struct PtySession { + pub process: PtyProcess, + + /// A file handle of the stdout of the pty process + pub process_stdout: File, + + /// A file handle of the stdin of the pty process + pub process_stdin: File, +} + +/// ``` +/// use std::process::Command; +/// use rattler_pty::unix::PtySession; +/// +/// let process = PtySession::new(Command::new("bash")).unwrap(); +/// ``` +impl PtySession { + /// Constructs a new session + pub fn new(command: Command) -> io::Result { + let process = PtyProcess::new( + command, + PtyProcessOptions { + echo: true, + ..Default::default() + }, + )?; + + let process_stdin = process.get_file_handle()?; + let process_stdout = process.get_file_handle()?; + + Ok(Self { + process, + process_stdout, + process_stdin, + }) + } + + /// Send string to process. As stdin of the process is most likely buffered, you'd + /// need to call `flush()` after `send()` to make the process actually see your input. + /// + /// Returns number of written bytes + pub fn send>(&mut self, s: B) -> io::Result { + self.process_stdin.write(s.as_ref()) + } + + /// Sends string and a newline to process. This is guaranteed to be flushed to the process. + /// Returns number of written bytes. + pub fn send_line(&mut self, line: &str) -> io::Result { + let mut len = self.send(line)?; + len += self.process_stdin.write(b"\n")?; + Ok(len) + } + + /// Make sure all bytes written via `send()` are sent to the process + pub fn flush(&mut self) -> io::Result<()> { + self.process_stdin.flush() + } + + /// Interact with the process. This will put the current process into raw mode and + /// forward all input from stdin to the process and all output from the process to stdout. + /// This will block until the process exits. + pub fn interact(&mut self, wait_until: Option<&str>) -> io::Result> { + // Make sure anything we have written so far has been flushed. + self.flush()?; + + // Put the process into raw mode + let original_mode = self.process.set_raw()?; + + let process_stdout_clone = self.process_stdout.try_clone()?; + let process_stdout_fd = process_stdout_clone.as_fd(); + let stdin = std::io::stdin(); + let stdin_fd = stdin.as_fd(); + + // Create a FDSet for the select call + let mut fd_set = FdSet::new(); + fd_set.insert(process_stdout_fd); + fd_set.insert(stdin_fd); + + // Create a buffer for reading from the process + let mut buf = [0u8; 2048]; + + // Catch the SIGWINCH signal to handle window resizing + // and forward the new terminal size to the process + let mut signals = Signals::new([SIGWINCH])?; + + let mut write_stdout = wait_until.is_none(); + // Call select in a loop and handle incoming data + let exit_status = loop { + // Make sure that the process is still alive + let status = self.process.status(); + if status != Some(WaitStatus::StillAlive) { + break status; + } + + // Handle window resizing + for signal in signals.pending() { + match signal { + SIGWINCH => { + // get window size + let mut size: libc::winsize = unsafe { std::mem::zeroed() }; + let res = unsafe { + libc::ioctl(libc::STDOUT_FILENO, libc::TIOCGWINSZ, &mut size) + }; + if res == 0 { + self.process.set_window_size(size)?; + } + } + _ => unreachable!(), + } + } + + let mut select_timeout = TimeVal::new(4, 0); + let mut select_set = fd_set; + + let res = select::select(None, &mut select_set, None, None, &mut select_timeout); + if let Err(error) = res { + if error == Errno::EINTR { + // EINTR is not an error, it just means that we got interrupted by a signal (e.g. SIGWINCH) + continue; + } else { + self.process.set_mode(original_mode)?; + return Err(std::io::Error::from(error)); + } + } else { + // We have new data coming from the process + if select_set.contains(process_stdout_fd) { + let bytes_read = self.process_stdout.read(&mut buf).unwrap_or(0); + if !write_stdout { + if let Some(wait_until) = wait_until { + if buf[..bytes_read] + .windows(wait_until.len()) + .any(|window| window == wait_until.as_bytes()) + { + write_stdout = true; + } + } + } else if bytes_read > 0 { + io::stdout().write_all(&buf[..bytes_read])?; + io::stdout().flush()?; + } + } + + // or from stdin + if select_set.contains(stdin_fd) { + let bytes_read = io::stdin().read(&mut buf)?; + self.process_stdin.write_all(&buf[..bytes_read])?; + self.process_stdin.flush()?; + } + } + }; + + // Restore the original terminal mode + self.process.set_mode(original_mode)?; + + match exit_status { + Some(WaitStatus::Exited(_, code)) => Ok(Some(code)), + Some(WaitStatus::Signaled(_, signal, _)) => Ok(Some(128 + signal as i32)), + _ => Ok(None), + } + } +} diff --git a/crates/rattler_pty/tests/integration_test.rs b/crates/rattler_pty/tests/integration_test.rs new file mode 100644 index 000000000..c77db4da3 --- /dev/null +++ b/crates/rattler_pty/tests/integration_test.rs @@ -0,0 +1,40 @@ +#[cfg(test)] +#[cfg(unix)] +mod tests { + + use rattler_pty::unix::PtyProcess; + use rattler_pty::unix::PtyProcessOptions; + + use nix::sys::{signal, wait}; + use std::io::{BufRead, BufReader, LineWriter, Write}; + use std::{self, process::Command, thread, time}; + + #[test] + /// Open cat, write string, read back string twice, send Ctrl^C and check that cat exited + fn test_cat() -> std::io::Result<()> { + let process = PtyProcess::new( + Command::new("cat"), + PtyProcessOptions { + echo: false, + window_size: Default::default(), + }, + ) + .expect("could not execute cat"); + let f = process.get_file_handle().unwrap(); + 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"); + + // this sleep solves an edge case of some cases when cat is somehow not "ready" + // to take the ^C (occasional test hangs) + thread::sleep(time::Duration::from_millis(100)); + writer.write_all(&[3])?; // send ^C + 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(()) + } +}