|
| 1 | +pub use nix::sys::{signal, wait}; |
| 2 | +use nix::{ |
| 3 | + self, |
| 4 | + fcntl::{open, OFlag}, |
| 5 | + libc::{STDERR_FILENO, STDIN_FILENO, STDOUT_FILENO}, |
| 6 | + pty::{grantpt, posix_openpt, unlockpt, PtyMaster, Winsize}, |
| 7 | + sys::termios::{InputFlags, Termios}, |
| 8 | + sys::{stat, termios}, |
| 9 | + unistd::{close, dup, dup2, fork, setsid, ForkResult, Pid}, |
| 10 | +}; |
| 11 | +use std::os::fd::AsFd; |
| 12 | +use std::{ |
| 13 | + self, |
| 14 | + fs::File, |
| 15 | + io, |
| 16 | + os::unix::{ |
| 17 | + io::{AsRawFd, FromRawFd}, |
| 18 | + process::CommandExt, |
| 19 | + }, |
| 20 | + process::Command, |
| 21 | + thread, time, |
| 22 | +}; |
| 23 | + |
| 24 | +#[cfg(target_os = "linux")] |
| 25 | +use nix::pty::ptsname_r; |
| 26 | + |
| 27 | +/// Start a process in a forked tty so you can interact with it the same as you would |
| 28 | +/// within a terminal |
| 29 | +/// |
| 30 | +/// The process and pty session are killed upon dropping PtyProcess |
| 31 | +pub struct PtyProcess { |
| 32 | + pub pty: PtyMaster, |
| 33 | + pub child_pid: Pid, |
| 34 | + kill_timeout: Option<time::Duration>, |
| 35 | +} |
| 36 | + |
| 37 | +#[cfg(target_os = "macos")] |
| 38 | +/// ptsname_r is a linux extension but ptsname isn't thread-safe |
| 39 | +/// instead of using a static mutex this calls ioctl with TIOCPTYGNAME directly |
| 40 | +/// based on https://blog.tarq.io/ptsname-on-osx-with-rust/ |
| 41 | +fn ptsname_r(fd: &PtyMaster) -> nix::Result<String> { |
| 42 | + use nix::libc::{ioctl, TIOCPTYGNAME}; |
| 43 | + use std::ffi::CStr; |
| 44 | + |
| 45 | + // the buffer size on OSX is 128, defined by sys/ttycom.h |
| 46 | + let mut buf: [i8; 128] = [0; 128]; |
| 47 | + |
| 48 | + unsafe { |
| 49 | + match ioctl(fd.as_raw_fd(), TIOCPTYGNAME as u64, &mut buf) { |
| 50 | + 0 => { |
| 51 | + let res = CStr::from_ptr(buf.as_ptr()).to_string_lossy().into_owned(); |
| 52 | + Ok(res) |
| 53 | + } |
| 54 | + _ => Err(nix::Error::last()), |
| 55 | + } |
| 56 | + } |
| 57 | +} |
| 58 | + |
| 59 | +#[derive(Default)] |
| 60 | +pub struct PtyProcessOptions { |
| 61 | + pub echo: bool, |
| 62 | + pub window_size: Option<Winsize>, |
| 63 | +} |
| 64 | + |
| 65 | +impl PtyProcess { |
| 66 | + /// Start a process in a forked pty |
| 67 | + pub fn new(mut command: Command, opts: PtyProcessOptions) -> nix::Result<Self> { |
| 68 | + // Open a new PTY master |
| 69 | + let master_fd = posix_openpt(OFlag::O_RDWR)?; |
| 70 | + |
| 71 | + // Allow a slave to be generated for it |
| 72 | + grantpt(&master_fd)?; |
| 73 | + unlockpt(&master_fd)?; |
| 74 | + |
| 75 | + // on Linux this is the libc function, on OSX this is our implementation of ptsname_r |
| 76 | + let slave_name = ptsname_r(&master_fd)?; |
| 77 | + |
| 78 | + // Get the current window size if it was not specified |
| 79 | + let window_size = opts.window_size.unwrap_or_else(|| { |
| 80 | + // find current window size with ioctl |
| 81 | + let mut size: libc::winsize = unsafe { std::mem::zeroed() }; |
| 82 | + // Query the terminal dimensions |
| 83 | + unsafe { libc::ioctl(io::stdout().as_raw_fd(), libc::TIOCGWINSZ, &mut size) }; |
| 84 | + size |
| 85 | + }); |
| 86 | + |
| 87 | + match unsafe { fork()? } { |
| 88 | + ForkResult::Child => { |
| 89 | + // Avoid leaking master fd |
| 90 | + close(master_fd.as_raw_fd())?; |
| 91 | + |
| 92 | + setsid()?; // create new session with child as session leader |
| 93 | + let slave_fd = open( |
| 94 | + std::path::Path::new(&slave_name), |
| 95 | + OFlag::O_RDWR, |
| 96 | + stat::Mode::empty(), |
| 97 | + )?; |
| 98 | + |
| 99 | + // assign stdin, stdout, stderr to the tty, just like a terminal does |
| 100 | + dup2(slave_fd, STDIN_FILENO)?; |
| 101 | + dup2(slave_fd, STDOUT_FILENO)?; |
| 102 | + dup2(slave_fd, STDERR_FILENO)?; |
| 103 | + |
| 104 | + // Avoid leaking slave fd |
| 105 | + if slave_fd > STDERR_FILENO { |
| 106 | + close(slave_fd)?; |
| 107 | + } |
| 108 | + |
| 109 | + // set echo off |
| 110 | + set_echo(io::stdin(), opts.echo)?; |
| 111 | + set_window_size(io::stdout().as_raw_fd(), window_size)?; |
| 112 | + |
| 113 | + // let mut flags = termios::tcgetattr(io::stdin())?; |
| 114 | + // flags.local_flags |= termios::LocalFlags::ECHO; |
| 115 | + // termios::tcsetattr(io::stdin(), termios::SetArg::TCSANOW, &flags)?; |
| 116 | + |
| 117 | + let _ = command.exec(); |
| 118 | + Err(nix::Error::last()) |
| 119 | + } |
| 120 | + ForkResult::Parent { child: child_pid } => Ok(PtyProcess { |
| 121 | + pty: master_fd, |
| 122 | + child_pid, |
| 123 | + kill_timeout: None, |
| 124 | + }), |
| 125 | + } |
| 126 | + } |
| 127 | + |
| 128 | + /// Get handle to pty fork for reading/writing |
| 129 | + pub fn get_file_handle(&self) -> nix::Result<File> { |
| 130 | + // needed because otherwise fd is closed both by dropping process and reader/writer |
| 131 | + let fd = dup(self.pty.as_raw_fd())?; |
| 132 | + unsafe { Ok(File::from_raw_fd(fd)) } |
| 133 | + } |
| 134 | + |
| 135 | + /// Get status of child process, non-blocking. |
| 136 | + /// |
| 137 | + /// This method runs waitpid on the process. |
| 138 | + /// This means: If you ran `exit()` before or `status()` this method will |
| 139 | + /// return `None` |
| 140 | + pub fn status(&self) -> Option<wait::WaitStatus> { |
| 141 | + if let Ok(status) = wait::waitpid(self.child_pid, Some(wait::WaitPidFlag::WNOHANG)) { |
| 142 | + Some(status) |
| 143 | + } else { |
| 144 | + None |
| 145 | + } |
| 146 | + } |
| 147 | + |
| 148 | + /// Regularly exit the process, this method is blocking until the process is dead |
| 149 | + pub fn exit(&mut self) -> nix::Result<wait::WaitStatus> { |
| 150 | + self.kill(signal::SIGTERM) |
| 151 | + } |
| 152 | + |
| 153 | + /// Kill the process with a specific signal. This method blocks, until the process is dead |
| 154 | + /// |
| 155 | + /// repeatedly sends SIGTERM to the process until it died, |
| 156 | + /// the pty session is closed upon dropping PtyMaster, |
| 157 | + /// so we don't need to explicitly do that here. |
| 158 | + /// |
| 159 | + /// if `kill_timeout` is set and a repeated sending of signal does not result in the process |
| 160 | + /// being killed, then `kill -9` is sent after the `kill_timeout` duration has elapsed. |
| 161 | + pub fn kill(&mut self, sig: signal::Signal) -> nix::Result<wait::WaitStatus> { |
| 162 | + let start = time::Instant::now(); |
| 163 | + loop { |
| 164 | + match signal::kill(self.child_pid, sig) { |
| 165 | + Ok(_) => {} |
| 166 | + // process was already killed before -> ignore |
| 167 | + Err(nix::errno::Errno::ESRCH) => { |
| 168 | + return Ok(wait::WaitStatus::Exited(Pid::from_raw(0), 0)); |
| 169 | + } |
| 170 | + Err(e) => return Err(e), |
| 171 | + } |
| 172 | + |
| 173 | + match self.status() { |
| 174 | + Some(status) if status != wait::WaitStatus::StillAlive => return Ok(status), |
| 175 | + Some(_) | None => thread::sleep(time::Duration::from_millis(100)), |
| 176 | + } |
| 177 | + // kill -9 if timeout is reached |
| 178 | + if let Some(timeout) = self.kill_timeout { |
| 179 | + if start.elapsed() > timeout { |
| 180 | + signal::kill(self.child_pid, signal::Signal::SIGKILL)? |
| 181 | + } |
| 182 | + } |
| 183 | + } |
| 184 | + } |
| 185 | + |
| 186 | + /// Set raw mode on stdin and return the original mode |
| 187 | + pub fn set_raw(&self) -> nix::Result<Termios> { |
| 188 | + let original_mode = termios::tcgetattr(io::stdin())?; |
| 189 | + let mut raw_mode = original_mode.clone(); |
| 190 | + raw_mode.input_flags.remove( |
| 191 | + InputFlags::BRKINT |
| 192 | + | InputFlags::ICRNL |
| 193 | + | InputFlags::INPCK |
| 194 | + | InputFlags::ISTRIP |
| 195 | + | InputFlags::IXON, |
| 196 | + ); |
| 197 | + raw_mode.output_flags.remove(termios::OutputFlags::OPOST); |
| 198 | + raw_mode |
| 199 | + .control_flags |
| 200 | + .remove(termios::ControlFlags::CSIZE | termios::ControlFlags::PARENB); |
| 201 | + raw_mode.control_flags.insert(termios::ControlFlags::CS8); |
| 202 | + raw_mode.local_flags.remove( |
| 203 | + termios::LocalFlags::ECHO |
| 204 | + | termios::LocalFlags::ICANON |
| 205 | + | termios::LocalFlags::IEXTEN |
| 206 | + | termios::LocalFlags::ISIG, |
| 207 | + ); |
| 208 | + |
| 209 | + raw_mode.control_chars[termios::SpecialCharacterIndices::VMIN as usize] = 1; |
| 210 | + raw_mode.control_chars[termios::SpecialCharacterIndices::VTIME as usize] = 0; |
| 211 | + |
| 212 | + termios::tcsetattr(io::stdin(), termios::SetArg::TCSAFLUSH, &raw_mode)?; |
| 213 | + |
| 214 | + Ok(original_mode) |
| 215 | + } |
| 216 | + |
| 217 | + pub fn set_mode(&self, original_mode: Termios) -> nix::Result<()> { |
| 218 | + termios::tcsetattr(io::stdin(), termios::SetArg::TCSAFLUSH, &original_mode)?; |
| 219 | + Ok(()) |
| 220 | + } |
| 221 | + |
| 222 | + pub fn set_window_size(&self, window_size: Winsize) -> nix::Result<()> { |
| 223 | + set_window_size(self.pty.as_raw_fd(), window_size) |
| 224 | + } |
| 225 | +} |
| 226 | + |
| 227 | +pub fn set_window_size(raw_fd: i32, window_size: Winsize) -> nix::Result<()> { |
| 228 | + unsafe { libc::ioctl(raw_fd, nix::libc::TIOCSWINSZ, &window_size) }; |
| 229 | + Ok(()) |
| 230 | +} |
| 231 | + |
| 232 | +pub fn set_echo<Fd: AsFd>(fd: Fd, echo: bool) -> nix::Result<()> { |
| 233 | + let mut flags = termios::tcgetattr(&fd)?; |
| 234 | + if echo { |
| 235 | + flags.local_flags.insert(termios::LocalFlags::ECHO); |
| 236 | + } else { |
| 237 | + flags.local_flags.remove(termios::LocalFlags::ECHO); |
| 238 | + } |
| 239 | + termios::tcsetattr(&fd, termios::SetArg::TCSANOW, &flags)?; |
| 240 | + Ok(()) |
| 241 | +} |
| 242 | + |
| 243 | +impl Drop for PtyProcess { |
| 244 | + fn drop(&mut self) { |
| 245 | + if let Some(wait::WaitStatus::StillAlive) = self.status() { |
| 246 | + self.exit().expect("cannot exit"); |
| 247 | + } |
| 248 | + } |
| 249 | +} |
0 commit comments