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