1#![allow(unused)]
27
28use std::{
30 cell::RefCell,
31 ffi::{
32 CStr,
33 CString,
34 OsStr,
35 },
36 fs::File,
37 io,
38 io::{
39 Read,
40 Result as IoResult,
41 Write,
42 },
43 mem,
44 os::{
45 fd::{
46 AsFd,
47 OwnedFd,
48 },
49 unix::{
50 ffi::{
51 OsStrExt,
52 OsStringExt,
53 },
54 io::{
55 AsRawFd,
56 FromRawFd,
57 RawFd,
58 },
59 process::CommandExt,
60 },
61 },
62 path::{
63 Path,
64 PathBuf,
65 },
66 process::Command,
67 ptr,
68};
69
70use color_eyre::eyre::{
71 Error,
72 bail,
73};
74use filedescriptor::FileDescriptor;
75use nix::{
76 libc::{
77 self,
78 pid_t,
79 winsize,
80 },
81 unistd::{
82 Pid,
83 dup2,
84 execv,
85 fork,
86 },
87};
88
89use crate::cmdbuilder::CommandBuilder;
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub struct PtySize {
94 pub rows: u16,
96 pub cols: u16,
98 pub pixel_width: u16,
101 pub pixel_height: u16,
104}
105
106impl Default for PtySize {
107 fn default() -> Self {
108 Self {
109 rows: 24,
110 cols: 80,
111 pixel_width: 0,
112 pixel_height: 0,
113 }
114 }
115}
116
117pub trait MasterPty: Send {
119 fn resize(&self, size: PtySize) -> Result<(), Error>;
123 fn get_size(&self) -> Result<PtySize, Error>;
125 fn try_clone_reader(&self) -> Result<Box<dyn std::io::Read + Send>, Error>;
128 fn take_writer(&self) -> Result<Box<dyn std::io::Write + Send>, Error>;
133
134 #[cfg(unix)]
137 fn process_group_leader(&self) -> Option<libc::pid_t>;
138
139 #[cfg(unix)]
144 fn as_raw_fd(&self) -> Option<RawFd>;
145
146 #[cfg(unix)]
147 fn tty_name(&self) -> Option<std::path::PathBuf>;
148
149 #[cfg(unix)]
152 fn get_termios(&self) -> Option<nix::sys::termios::Termios> {
153 None
154 }
155}
156
157pub trait Child: std::fmt::Debug + ChildKiller + Send {
160 fn try_wait(&mut self) -> IoResult<Option<ExitStatus>>;
165 fn wait(&mut self) -> IoResult<ExitStatus>;
168 fn process_id(&self) -> Pid;
171}
172
173pub trait ChildKiller: std::fmt::Debug + Send {
175 fn kill(&mut self) -> IoResult<()>;
177
178 fn clone_killer(&self) -> Box<dyn ChildKiller + Send + Sync>;
182}
183
184#[derive(Debug, Clone)]
186pub struct ExitStatus {
187 code: u32,
188 signal: Option<String>,
189}
190
191impl ExitStatus {
192 pub fn with_exit_code(code: u32) -> Self {
194 Self { code, signal: None }
195 }
196
197 pub fn with_signal(signal: &str) -> Self {
199 Self {
200 code: 1,
201 signal: Some(signal.to_string()),
202 }
203 }
204
205 pub fn success(&self) -> bool {
207 match self.signal {
208 None => self.code == 0,
209 Some(_) => false,
210 }
211 }
212
213 pub fn exit_code(&self) -> u32 {
215 self.code
216 }
217
218 pub fn signal(&self) -> Option<&str> {
220 self.signal.as_deref()
221 }
222}
223
224impl From<std::process::ExitStatus> for ExitStatus {
225 fn from(status: std::process::ExitStatus) -> Self {
226 #[cfg(unix)]
227 {
228 use std::os::unix::process::ExitStatusExt;
229
230 if let Some(signal) = status.signal() {
231 let signame = unsafe { libc::strsignal(signal) };
232 let signal = if signame.is_null() {
233 format!("Signal {signal}")
234 } else {
235 let signame = unsafe { std::ffi::CStr::from_ptr(signame) };
236 signame.to_string_lossy().to_string()
237 };
238
239 return Self {
240 code: status.code().map(|c| c as u32).unwrap_or(1),
241 signal: Some(signal),
242 };
243 }
244 }
245
246 let code = status
247 .code()
248 .map(|c| c as u32)
249 .unwrap_or_else(|| if status.success() { 0 } else { 1 });
250
251 Self { code, signal: None }
252 }
253}
254
255impl std::fmt::Display for ExitStatus {
256 fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
257 if self.success() {
258 write!(fmt, "Success")
259 } else {
260 match &self.signal {
261 Some(sig) => write!(fmt, "Terminated by {sig}"),
262 None => write!(fmt, "Exited with code {}", self.code),
263 }
264 }
265 }
266}
267
268pub struct PtyPair {
269 pub slave: UnixSlavePty,
272 pub master: UnixMasterPty,
273}
274
275pub trait PtySystem {
279 fn openpty(&self, size: PtySize) -> color_eyre::Result<PtyPair>;
283}
284
285impl Child for std::process::Child {
286 fn try_wait(&mut self) -> IoResult<Option<ExitStatus>> {
287 Self::try_wait(self).map(|s| s.map(Into::into))
288 }
289
290 fn wait(&mut self) -> IoResult<ExitStatus> {
291 Self::wait(self).map(Into::into)
292 }
293
294 fn process_id(&self) -> Pid {
295 Pid::from_raw(self.id() as pid_t)
296 }
297}
298
299#[derive(Debug)]
300struct ProcessSignaller {
301 pid: Option<Pid>,
302}
303
304impl ChildKiller for ProcessSignaller {
305 fn kill(&mut self) -> IoResult<()> {
306 if let Some(pid) = self.pid {
307 let result = unsafe { libc::kill(pid.as_raw(), libc::SIGHUP) };
308 if result != 0 {
309 return Err(std::io::Error::last_os_error());
310 }
311 }
312 Ok(())
313 }
314
315 fn clone_killer(&self) -> Box<dyn ChildKiller + Send + Sync> {
316 Box::new(Self { pid: self.pid })
317 }
318}
319
320impl ChildKiller for std::process::Child {
321 fn kill(&mut self) -> IoResult<()> {
322 #[cfg(unix)]
323 {
324 let result = unsafe { libc::kill(self.id() as i32, libc::SIGHUP) };
328 if result != 0 {
329 return Err(std::io::Error::last_os_error());
330 }
331
332 for attempt in 0..5 {
339 if attempt > 0 {
340 std::thread::sleep(std::time::Duration::from_millis(50));
341 }
342
343 if let Ok(Some(_)) = self.try_wait() {
344 return Ok(());
346 }
347 }
348
349 }
351
352 Self::kill(self)
353 }
354
355 fn clone_killer(&self) -> Box<dyn ChildKiller + Send + Sync> {
356 Box::new(ProcessSignaller {
357 pid: Some(self.process_id()),
358 })
359 }
360}
361
362pub fn native_pty_system() -> NativePtySystem {
363 NativePtySystem::default()
364}
365
366pub type NativePtySystem = UnixPtySystem;
367
368#[derive(Default)]
369pub struct UnixPtySystem {}
370
371fn openpty(size: PtySize) -> color_eyre::Result<(UnixMasterPty, UnixSlavePty)> {
372 let mut master: RawFd = -1;
373 let mut slave: RawFd = -1;
374
375 let mut size = winsize {
376 ws_row: size.rows,
377 ws_col: size.cols,
378 ws_xpixel: size.pixel_width,
379 ws_ypixel: size.pixel_height,
380 };
381
382 let result = unsafe {
383 libc::openpty(
384 &mut master,
385 &mut slave,
386 ptr::null_mut(),
387 ptr::null_mut(),
388 &size,
389 )
390 };
391
392 if result != 0 {
393 bail!("failed to openpty: {:?}", io::Error::last_os_error());
394 }
395
396 let tty_name = tty_name(slave);
397
398 let master = UnixMasterPty {
399 fd: PtyFd(unsafe { FileDescriptor::from_raw_fd(master) }),
400 took_writer: RefCell::new(false),
401 tty_name,
402 };
403 let slave = UnixSlavePty {
404 fd: PtyFd(unsafe { FileDescriptor::from_raw_fd(slave) }),
405 };
406
407 cloexec(master.fd.as_raw_fd())?;
412 cloexec(slave.fd.as_raw_fd())?;
413
414 Ok((master, slave))
415}
416
417impl PtySystem for UnixPtySystem {
418 fn openpty(&self, size: PtySize) -> color_eyre::Result<PtyPair> {
419 let (master, slave) = openpty(size)?;
420 Ok(PtyPair { master, slave })
421 }
422}
423
424pub struct PtyFd(pub FileDescriptor);
425impl std::ops::Deref for PtyFd {
426 type Target = FileDescriptor;
427 fn deref(&self) -> &FileDescriptor {
428 &self.0
429 }
430}
431impl std::ops::DerefMut for PtyFd {
432 fn deref_mut(&mut self) -> &mut FileDescriptor {
433 &mut self.0
434 }
435}
436
437impl Read for PtyFd {
438 fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
439 match self.0.read(buf) {
440 Err(ref e) if e.raw_os_error() == Some(libc::EIO) => {
441 Ok(0)
446 }
447 x => x,
448 }
449 }
450}
451
452fn tty_name(fd: RawFd) -> Option<PathBuf> {
453 let mut buf = vec![0 as std::ffi::c_char; 128];
454
455 loop {
456 let res = unsafe { libc::ttyname_r(fd, buf.as_mut_ptr(), buf.len()) };
457
458 if res == libc::ERANGE {
459 if buf.len() > 64 * 1024 {
460 return None;
464 }
465 buf.resize(buf.len() * 2, 0 as std::ffi::c_char);
466 continue;
467 }
468
469 return if res == 0 {
470 let cstr = unsafe { std::ffi::CStr::from_ptr(buf.as_ptr()) };
471 let osstr = OsStr::from_bytes(cstr.to_bytes());
472 Some(PathBuf::from(osstr))
473 } else {
474 None
475 };
476 }
477}
478
479fn close_random_fds() {
497 if let Ok(dir) = std::fs::read_dir("/proc/self/fd").or_else(|_| std::fs::read_dir("/dev/fd")) {
502 let mut fds = vec![];
503 for entry in dir {
504 if let Some(num) = entry
505 .ok()
506 .map(|e| e.file_name())
507 .and_then(|s| s.into_string().ok())
508 .and_then(|n| n.parse::<libc::c_int>().ok())
509 && num > 2
510 {
511 fds.push(num);
512 }
513 }
514 for fd in fds {
515 let _ = nix::unistd::close(fd);
516 }
517 }
518}
519
520impl PtyFd {
521 fn resize(&self, size: PtySize) -> Result<(), Error> {
522 let ws_size = winsize {
523 ws_row: size.rows,
524 ws_col: size.cols,
525 ws_xpixel: size.pixel_width,
526 ws_ypixel: size.pixel_height,
527 };
528
529 if unsafe {
530 libc::ioctl(
531 self.0.as_raw_fd(),
532 libc::TIOCSWINSZ as _,
533 &ws_size as *const _,
534 )
535 } != 0
536 {
537 bail!(
538 "failed to ioctl(TIOCSWINSZ): {:?}",
539 io::Error::last_os_error()
540 );
541 }
542
543 Ok(())
544 }
545
546 fn get_size(&self) -> Result<PtySize, Error> {
547 let mut size: winsize = unsafe { mem::zeroed() };
548 if unsafe {
549 libc::ioctl(
550 self.0.as_raw_fd(),
551 libc::TIOCGWINSZ as _,
552 &mut size as *mut _,
553 )
554 } != 0
555 {
556 bail!(
557 "failed to ioctl(TIOCGWINSZ): {:?}",
558 io::Error::last_os_error()
559 );
560 }
561 Ok(PtySize {
562 rows: size.ws_row,
563 cols: size.ws_col,
564 pixel_width: size.ws_xpixel,
565 pixel_height: size.ws_ypixel,
566 })
567 }
568
569 fn spawn_command(
570 &self,
571 command: CommandBuilder,
572 pre_exec: impl FnOnce(&Path) -> color_eyre::Result<()> + Send + Sync + 'static,
573 ) -> color_eyre::Result<Pid> {
574 spawn_command_from_pty_fd(Some(self), command, pre_exec)
575 }
576}
577
578pub fn spawn_command(
579 pts: Option<&UnixSlavePty>,
580 command: CommandBuilder,
581 pre_exec: impl FnOnce(&Path) -> color_eyre::Result<()> + Send + Sync + 'static,
582) -> color_eyre::Result<Pid> {
583 if let Some(pts) = pts {
584 pts.spawn_command(command, pre_exec)
585 } else {
586 spawn_command_from_pty_fd(None, command, pre_exec)
587 }
588}
589
590fn spawn_command_from_pty_fd(
591 pty: Option<&PtyFd>,
592 command: CommandBuilder,
593 pre_exec: impl FnOnce(&Path) -> color_eyre::Result<()> + Send + Sync + 'static,
594) -> color_eyre::Result<Pid> {
595 let configured_umask = command.umask;
596
597 let mut cmd = command.build()?;
598
599 match unsafe { fork()? } {
600 nix::unistd::ForkResult::Parent { child } => Ok(child),
601 nix::unistd::ForkResult::Child => {
602 if let Some(pty) = pty {
603 let mut stdin = unsafe { OwnedFd::from_raw_fd(0) };
604 let mut stdout = unsafe { OwnedFd::from_raw_fd(1) };
605 let mut stderr = unsafe { OwnedFd::from_raw_fd(2) };
606 dup2(pty.as_fd(), &mut stdin).unwrap();
607 dup2(pty.as_fd(), &mut stdout).unwrap();
608 dup2(pty.as_fd(), &mut stderr).unwrap();
609 std::mem::forget(stdin);
610 std::mem::forget(stdout);
611 std::mem::forget(stderr);
612 }
613
614 for signo in &[
618 libc::SIGCHLD,
619 libc::SIGHUP,
620 libc::SIGINT,
621 libc::SIGQUIT,
622 libc::SIGTERM,
623 libc::SIGALRM,
624 ] {
625 unsafe {
626 _ = libc::signal(*signo, libc::SIG_DFL);
627 }
628 }
629
630 unsafe {
631 let empty_set: libc::sigset_t = std::mem::zeroed();
632 _ = libc::sigprocmask(libc::SIG_SETMASK, &empty_set, std::ptr::null_mut());
633 }
634
635 pre_exec(&cmd.program).unwrap();
636
637 close_random_fds();
638
639 if let Some(mask) = configured_umask {
640 _ = unsafe { libc::umask(mask) };
641 }
642
643 execv(
644 &CString::new(cmd.program.into_os_string().into_vec()).unwrap(),
645 &cmd.args,
646 )
647 .unwrap();
648 unreachable!()
649 }
650 }
651}
652
653pub struct UnixMasterPty {
656 fd: PtyFd,
657 took_writer: RefCell<bool>,
658 tty_name: Option<PathBuf>,
659}
660
661pub struct UnixSlavePty {
664 pub fd: PtyFd,
665}
666
667impl UnixSlavePty {
668 pub fn spawn_command(
669 &self,
670 command: CommandBuilder,
671 pre_exec: impl FnOnce(&Path) -> color_eyre::Result<()> + Send + Sync + 'static,
672 ) -> color_eyre::Result<Pid> {
673 self.fd.spawn_command(command, pre_exec)
674 }
675}
676
677fn cloexec(fd: RawFd) -> Result<(), Error> {
679 let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
680 if flags == -1 {
681 bail!(
682 "fcntl to read flags failed: {:?}",
683 io::Error::last_os_error()
684 );
685 }
686 let result = unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) };
687 if result == -1 {
688 bail!(
689 "fcntl to set CLOEXEC failed: {:?}",
690 io::Error::last_os_error()
691 );
692 }
693 Ok(())
694}
695
696impl MasterPty for UnixMasterPty {
697 fn resize(&self, size: PtySize) -> Result<(), Error> {
698 self.fd.resize(size)
699 }
700
701 fn get_size(&self) -> Result<PtySize, Error> {
702 self.fd.get_size()
703 }
704
705 fn try_clone_reader(&self) -> Result<Box<dyn Read + Send>, Error> {
706 let fd = PtyFd(self.fd.try_clone()?);
707 Ok(Box::new(fd))
708 }
709
710 fn take_writer(&self) -> Result<Box<dyn Write + Send>, Error> {
711 if *self.took_writer.borrow() {
712 bail!("cannot take writer more than once");
713 }
714 *self.took_writer.borrow_mut() = true;
715 let fd = PtyFd(self.fd.try_clone()?);
716 Ok(Box::new(UnixMasterWriter { fd }))
717 }
718
719 fn as_raw_fd(&self) -> Option<RawFd> {
720 Some(self.fd.0.as_raw_fd())
721 }
722
723 fn tty_name(&self) -> Option<PathBuf> {
724 self.tty_name.clone()
725 }
726
727 fn process_group_leader(&self) -> Option<libc::pid_t> {
728 match unsafe { libc::tcgetpgrp(self.fd.0.as_raw_fd()) } {
729 pid if pid > 0 => Some(pid),
730 _ => None,
731 }
732 }
733
734 fn get_termios(&self) -> Option<nix::sys::termios::Termios> {
735 nix::sys::termios::tcgetattr(unsafe { File::from_raw_fd(self.fd.0.as_raw_fd()) }).ok()
736 }
737}
738
739struct UnixMasterWriter {
743 fd: PtyFd,
744}
745
746impl Drop for UnixMasterWriter {
747 fn drop(&mut self) {
748 let mut t: libc::termios = unsafe { std::mem::MaybeUninit::zeroed().assume_init() };
749 if unsafe { libc::tcgetattr(self.fd.0.as_raw_fd(), &mut t) } == 0 {
750 let eot = t.c_cc[libc::VEOF];
753 if eot != 0 {
754 let _ = self.fd.0.write_all(&[b'\n', eot]);
755 }
756 }
757 }
758}
759
760impl Write for UnixMasterWriter {
761 fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
762 self.fd.write(buf)
763 }
764 fn flush(&mut self) -> Result<(), io::Error> {
765 self.fd.flush()
766 }
767}
768
769#[cfg(test)]
770mod tests {
771 use std::{
772 io::{
773 Read,
774 Write,
775 },
776 time::Duration,
777 };
778
779 use nix::sys::wait::waitpid;
780
781 use super::*;
782
783 fn system() -> UnixPtySystem {
784 UnixPtySystem::default()
785 }
786
787 #[test]
788 fn test_ptysize_default() {
789 let s = PtySize::default();
790 assert_eq!(s.rows, 24);
791 assert_eq!(s.cols, 80);
792 assert_eq!(s.pixel_width, 0);
793 assert_eq!(s.pixel_height, 0);
794 }
795
796 #[test]
797 fn test_openpty_basic() {
798 let pty = system().openpty(PtySize::default()).unwrap();
799 assert!(pty.master.as_raw_fd().is_some());
800 }
801
802 #[test]
803 fn test_resize_and_get_size() {
804 let pty = system().openpty(PtySize::default()).unwrap();
805
806 let new_size = PtySize {
807 rows: 40,
808 cols: 100,
809 pixel_width: 0,
810 pixel_height: 0,
811 };
812
813 pty.master.resize(new_size).unwrap();
814 let got = pty.master.get_size().unwrap();
815
816 assert_eq!(got.rows, 40);
817 assert_eq!(got.cols, 100);
818 }
819
820 #[test]
821 fn test_master_slave_io() {
822 let pty = system().openpty(PtySize::default()).unwrap();
823
824 let mut writer = pty.master.take_writer().unwrap();
825 let mut reader = pty.master.try_clone_reader().unwrap();
826
827 writer.write_all(b"hello\n").unwrap();
828 writer.flush().unwrap();
829
830 let mut buf = [0u8; 64];
831 let n = reader.read(&mut buf).unwrap();
832
833 assert!(n > 0);
834 assert!(std::str::from_utf8(&buf[..n]).unwrap().contains("hello"));
835 }
836
837 #[test]
838 fn test_take_writer_only_once() {
839 let pty = system().openpty(PtySize::default()).unwrap();
840
841 let _w1 = pty.master.take_writer().unwrap();
842 let w2 = pty.master.take_writer();
843
844 assert!(w2.is_err());
845 }
846
847 #[test]
848 fn test_tty_name_present() {
849 let pty = system().openpty(PtySize::default()).unwrap();
850 let name = pty.master.tty_name();
851
852 assert!(name.is_some());
854 }
855
856 #[test]
857 fn test_spawn_command_echo() {
858 let pty = system().openpty(PtySize::default()).unwrap();
859
860 let mut cmd = CommandBuilder::new("echo");
861 cmd.arg("hello");
862
863 let pid = pty.slave.spawn_command(cmd, |_| Ok(())).unwrap();
864
865 let mut reader = pty.master.try_clone_reader().unwrap();
867 let mut buf = [0; 256];
868
869 std::thread::sleep(Duration::from_millis(100));
871 reader.read(&mut buf);
872
873 eprintln!("buf: {}", String::from_utf8_lossy(&buf));
874
875 assert!(buf.windows(5).any(|w| w == b"hello".as_slice()));
876
877 waitpid(pid, None).unwrap();
879 }
880
881 #[test]
882 fn test_exitstatus_helpers() {
883 let ok = ExitStatus::with_exit_code(0);
884 assert!(ok.success());
885 assert_eq!(ok.exit_code(), 0);
886 assert!(ok.signal().is_none());
887
888 let sig = ExitStatus::with_signal("SIGTERM");
889 assert!(!sig.success());
890 assert_eq!(sig.signal(), Some("SIGTERM"));
891 }
892}