Skip to main content

max / shop

6.5 KB · 193 lines History Blame Raw
1 //! PTY spawning + I/O for shop.
2 //!
3 //! Opens a Linux pseudo-terminal, forks, execs a shell in the child with the
4 //! slave as controlling terminal, and hands the parent a nonblocking master
5 //! fd for read/write. The master fd implements [`AsFd`] so callers can plug
6 //! it into any Unix event loop (calloop, mio, epoll).
7 //!
8 //! Reference: rio's teletypewriter/src/unix/mod.rs (MIT). Not a line-for-line
9 //! port — we use nix's safe wrappers rather than raw libc, and skip the
10 //! corcovado (mio 0.6) coupling entirely.
11
12 use std::ffi::CString;
13 use std::fs::File;
14 use std::io::{self, Write};
15 use std::os::fd::{AsFd, AsRawFd, BorrowedFd, OwnedFd};
16 use std::path::Path;
17
18 use nix::fcntl::{FcntlArg, OFlag, fcntl};
19 use nix::libc;
20 use nix::pty::{Winsize, openpty};
21 use nix::sys::signal::{Signal, kill};
22 use nix::unistd::{ForkResult, Pid, execvp, fork, setsid};
23
24 /// Dimensions of the terminal surface reported to the kernel via TIOCSWINSZ.
25 #[derive(Debug, Clone, Copy)]
26 pub struct PtySize {
27 pub cols: u16,
28 pub rows: u16,
29 pub cell_width: u16,
30 pub cell_height: u16,
31 }
32
33 impl PtySize {
34 fn to_winsize(self) -> Winsize {
35 Winsize {
36 ws_row: self.rows,
37 ws_col: self.cols,
38 ws_xpixel: self.cols.saturating_mul(self.cell_width),
39 ws_ypixel: self.rows.saturating_mul(self.cell_height),
40 }
41 }
42 }
43
44 /// A running shell attached to a PTY. Dropping this sends SIGHUP to the
45 /// child.
46 pub struct Pty {
47 master: OwnedFd,
48 child: Pid,
49 recorder: Option<File>,
50 }
51
52 impl Pty {
53 /// Fork a child running `command`, hand the child an already-set-up
54 /// controlling terminal, and return the parent-side master fd.
55 ///
56 /// `term` controls the value of the `TERM` env var for the child; the
57 /// shell/programs use it to select escape sequences. Start with
58 /// `"xterm-256color"` until we have a terminfo entry for shop.
59 pub fn spawn(command: &str, args: &[&str], size: PtySize, term: &str) -> anyhow::Result<Self> {
60 let ws = size.to_winsize();
61 let pair = openpty(Some(&ws), None)?;
62
63 // SAFETY: post-fork, the child only calls async-signal-safe libc
64 // calls (setsid, dup2, ioctl, execvp) plus setenv (not strictly
65 // signal-safe, tolerated by GLIBC for terminal setup).
66 #[allow(unsafe_code)]
67 match unsafe { fork() }? {
68 ForkResult::Parent { child } => {
69 drop(pair.slave);
70 let flags = fcntl(pair.master.as_fd(), FcntlArg::F_GETFL)?;
71 let nb = OFlag::from_bits_truncate(flags) | OFlag::O_NONBLOCK;
72 fcntl(pair.master.as_fd(), FcntlArg::F_SETFL(nb))?;
73 Ok(Self {
74 master: pair.master,
75 child,
76 recorder: None,
77 })
78 }
79 ForkResult::Child => {
80 run_child(pair.slave, command, args, term);
81 }
82 }
83 }
84
85 /// Send SIGWINCH to the child after updating the master's window size.
86 pub fn resize(&self, size: PtySize) -> anyhow::Result<()> {
87 let ws = size.to_winsize();
88 // SAFETY: master.as_raw_fd() is a valid fd for our process; ioctl
89 // TIOCSWINSZ reads Winsize by pointer.
90 #[allow(unsafe_code)]
91 let rc = unsafe {
92 libc::ioctl(
93 self.master.as_raw_fd(),
94 libc::TIOCSWINSZ,
95 std::ptr::from_ref(&ws),
96 )
97 };
98 if rc != 0 {
99 anyhow::bail!(io::Error::last_os_error());
100 }
101 Ok(())
102 }
103
104 /// Nonblocking read. Returns `Ok(0)` on EOF, `WouldBlock` when there's
105 /// nothing to read yet.
106 ///
107 /// If a recorder is attached (see [`set_recorder`](Self::set_recorder)),
108 /// bytes read here are teed to it byte-for-byte before returning.
109 /// Recorder write errors are swallowed — the terminal must not die
110 /// because a debug capture failed.
111 pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
112 let n = nix::unistd::read(&self.master, buf).map_err(io::Error::from)?;
113 if n > 0 {
114 if let Some(f) = self.recorder.as_ref() {
115 let _ = (&*f).write_all(&buf[..n]);
116 }
117 }
118 Ok(n)
119 }
120
121 /// Tee every subsequent successful [`read`](Self::read) to `path`. The
122 /// file is created (or truncated if it exists) up front, so a bad path
123 /// surfaces here rather than mid-session.
124 ///
125 /// Meant for capturing the PTY byte stream for feeding into the
126 /// `kitty-graphics-testkit` corpus. See `capture-apc`.
127 pub fn set_recorder(&mut self, path: impl AsRef<Path>) -> io::Result<()> {
128 self.recorder = Some(File::create(path)?);
129 Ok(())
130 }
131
132 /// Write bytes to the shell (e.g. keystrokes).
133 pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
134 nix::unistd::write(&self.master, buf).map_err(io::Error::from)
135 }
136
137 pub fn child(&self) -> Pid {
138 self.child
139 }
140 }
141
142 impl AsFd for Pty {
143 fn as_fd(&self) -> BorrowedFd<'_> {
144 self.master.as_fd()
145 }
146 }
147
148 impl Drop for Pty {
149 fn drop(&mut self) {
150 let _ = kill(self.child, Signal::SIGHUP);
151 }
152 }
153
154 fn run_child(slave: OwnedFd, command: &str, args: &[&str], term: &str) -> ! {
155 // New session detaches us from the parent's controlling terminal.
156 let _ = setsid();
157
158 // Make the slave our controlling terminal.
159 #[allow(unsafe_code)]
160 unsafe {
161 libc::ioctl(slave.as_raw_fd(), libc::TIOCSCTTY as _, 0);
162 }
163
164 // Duplicate slave onto stdin/stdout/stderr. libc::dup2 rather than nix's
165 // typed version, which needs an owned destination (unusable for stdio).
166 // SAFETY: slave.as_raw_fd() is valid; STDIN/OUT/ERR_FILENO are constants.
167 let slave_raw = slave.as_raw_fd();
168 #[allow(unsafe_code)]
169 unsafe {
170 libc::dup2(slave_raw, libc::STDIN_FILENO);
171 libc::dup2(slave_raw, libc::STDOUT_FILENO);
172 libc::dup2(slave_raw, libc::STDERR_FILENO);
173 }
174 drop(slave);
175
176 // TERM is the classic escape-set selector; SHELL is informational.
177 // SAFETY: single-threaded child process, no other threads to race.
178 #[allow(unsafe_code)]
179 unsafe {
180 std::env::set_var("TERM", term);
181 }
182
183 let cmd_c = CString::new(command).expect("command has NUL");
184 let mut arg_cstrs: Vec<CString> = Vec::with_capacity(args.len() + 1);
185 arg_cstrs.push(cmd_c.clone());
186 for a in args {
187 arg_cstrs.push(CString::new(*a).expect("arg has NUL"));
188 }
189 let _ = execvp(&cmd_c, &arg_cstrs);
190 // execvp only returns on failure.
191 std::process::exit(127);
192 }
193