Skip to main content

codex_utils_pty/
pty.rs

1use std::collections::HashMap;
2#[cfg(unix)]
3use std::fs::File;
4use std::io::ErrorKind;
5#[cfg(unix)]
6use std::os::fd::AsRawFd;
7#[cfg(unix)]
8use std::os::fd::FromRawFd;
9#[cfg(unix)]
10use std::os::fd::RawFd;
11#[cfg(unix)]
12use std::os::unix::process::CommandExt;
13use std::path::Path;
14#[cfg(unix)]
15use std::process::Command as StdCommand;
16#[cfg(unix)]
17use std::process::Stdio;
18use std::sync::Arc;
19use std::sync::Mutex as StdMutex;
20use std::sync::atomic::AtomicBool;
21use std::time::Duration;
22
23use anyhow::Result;
24use portable_pty::CommandBuilder;
25#[cfg(not(windows))]
26use portable_pty::native_pty_system;
27use tokio::sync::mpsc;
28use tokio::sync::oneshot;
29use tokio::task::JoinHandle;
30
31use crate::process::ChildTerminator;
32use crate::process::ProcessHandle;
33use crate::process::ProcessSignal;
34use crate::process::PtyHandles;
35use crate::process::PtyMasterHandle;
36use crate::process::SpawnedProcess;
37use crate::process::TerminalSize;
38#[cfg(unix)]
39use crate::process::exit_code_from_status;
40
41/// Returns true when ConPTY support is available (Windows only).
42#[cfg(windows)]
43pub fn conpty_supported() -> bool {
44    crate::win::conpty_supported()
45}
46
47/// Returns true when ConPTY support is available (non-Windows always true).
48#[cfg(not(windows))]
49pub fn conpty_supported() -> bool {
50    true
51}
52
53struct PtyChildTerminator {
54    killer: Box<dyn portable_pty::ChildKiller + Send + Sync>,
55    #[cfg(unix)]
56    process_group_id: Option<u32>,
57}
58
59impl ChildTerminator for PtyChildTerminator {
60    fn signal(&mut self, signal: ProcessSignal) -> std::io::Result<()> {
61        match signal {
62            ProcessSignal::Interrupt => {
63                #[cfg(unix)]
64                if let Some(process_group_id) = self.process_group_id {
65                    return crate::process_group::interrupt_process_group(process_group_id);
66                }
67
68                Err(crate::process::unsupported_signal(signal))
69            }
70        }
71    }
72
73    fn kill(&mut self) -> std::io::Result<()> {
74        #[cfg(unix)]
75        if let Some(process_group_id) = self.process_group_id {
76            // Match the pipe backend's hard-kill behavior so descendant
77            // processes from interactive shells/REPLs do not survive shutdown.
78            // Also try the direct child killer in case the cached PGID is stale.
79            let process_group_kill_result =
80                crate::process_group::kill_process_group(process_group_id);
81            let child_kill_result = self.killer.kill();
82            return match child_kill_result {
83                Ok(()) => Ok(()),
84                Err(err) if err.kind() == ErrorKind::NotFound => process_group_kill_result,
85                Err(err) => process_group_kill_result.or(Err(err)),
86            };
87        }
88
89        self.killer.kill()
90    }
91}
92
93#[cfg(unix)]
94struct RawPidTerminator {
95    process_group_id: u32,
96}
97
98#[cfg(unix)]
99impl ChildTerminator for RawPidTerminator {
100    fn signal(&mut self, signal: ProcessSignal) -> std::io::Result<()> {
101        match signal {
102            ProcessSignal::Interrupt => {
103                crate::process_group::interrupt_process_group(self.process_group_id)
104            }
105        }
106    }
107
108    fn kill(&mut self) -> std::io::Result<()> {
109        crate::process_group::kill_process_group(self.process_group_id)
110    }
111}
112
113fn platform_native_pty_system() -> Box<dyn portable_pty::PtySystem + Send> {
114    #[cfg(windows)]
115    {
116        Box::new(crate::win::ConPtySystem::default())
117    }
118
119    #[cfg(not(windows))]
120    {
121        native_pty_system()
122    }
123}
124
125/// Spawn a process attached to a PTY, preserving selected inherited file
126/// descriptors across exec on Unix.
127pub async fn spawn_process(
128    program: &str,
129    args: &[String],
130    cwd: &Path,
131    env: &HashMap<String, String>,
132    arg0: &Option<String>,
133    size: TerminalSize,
134    inherited_fds: &[i32],
135) -> Result<SpawnedProcess> {
136    if program.is_empty() {
137        anyhow::bail!("missing program for PTY spawn");
138    }
139
140    #[cfg(not(unix))]
141    let _ = inherited_fds;
142
143    #[cfg(unix)]
144    if !inherited_fds.is_empty() {
145        return spawn_process_preserving_fds(program, args, cwd, env, arg0, size, inherited_fds)
146            .await;
147    }
148
149    spawn_process_portable(program, args, cwd, env, arg0, size).await
150}
151
152async fn spawn_process_portable(
153    program: &str,
154    args: &[String],
155    cwd: &Path,
156    env: &HashMap<String, String>,
157    arg0: &Option<String>,
158    size: TerminalSize,
159) -> Result<SpawnedProcess> {
160    let pty_system = platform_native_pty_system();
161    let pair = pty_system.openpty(size.into())?;
162
163    let mut command_builder = CommandBuilder::new(arg0.as_ref().unwrap_or(&program.to_string()));
164    command_builder.cwd(cwd);
165    command_builder.env_clear();
166    for arg in args {
167        command_builder.arg(arg);
168    }
169    for (key, value) in env {
170        command_builder.env(key, value);
171    }
172
173    let mut child = pair.slave.spawn_command(command_builder)?;
174    #[cfg(unix)]
175    // portable-pty establishes the spawned PTY child as a new session leader on
176    // Unix, so PID == PGID and we can reuse the pipe backend's process-group
177    // hard-kill semantics for descendants.
178    let process_group_id = child.process_id();
179    let killer = child.clone_killer();
180
181    let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
182    let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>(128);
183    let (_stderr_tx, stderr_rx) = mpsc::channel::<Vec<u8>>(1);
184    let mut reader = pair.master.try_clone_reader()?;
185    let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
186        let mut buf = [0u8; 8_192];
187        loop {
188            match reader.read(&mut buf) {
189                Ok(0) => break,
190                Ok(n) => {
191                    let _ = stdout_tx.blocking_send(buf[..n].to_vec());
192                }
193                Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
194                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
195                    std::thread::sleep(Duration::from_millis(5));
196                    continue;
197                }
198                Err(_) => break,
199            }
200        }
201    });
202
203    let writer = pair.master.take_writer()?;
204    let writer = Arc::new(tokio::sync::Mutex::new(writer));
205    let writer_handle: JoinHandle<()> = tokio::spawn({
206        let writer = Arc::clone(&writer);
207        async move {
208            #[cfg(windows)]
209            let mut windows_input = crate::WindowsTtyInputNormalizer::default();
210            while let Some(bytes) = writer_rx.recv().await {
211                #[cfg(windows)]
212                let bytes = windows_input.normalize(&bytes);
213                let mut guard = writer.lock().await;
214                use std::io::Write;
215                let _ = guard.write_all(&bytes);
216                let _ = guard.flush();
217            }
218        }
219    });
220
221    let (exit_tx, exit_rx) = oneshot::channel::<i32>();
222    let exit_status = Arc::new(AtomicBool::new(false));
223    let wait_exit_status = Arc::clone(&exit_status);
224    let exit_code = Arc::new(StdMutex::new(None));
225    let wait_exit_code = Arc::clone(&exit_code);
226    let wait_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
227        let code = match child.wait() {
228            Ok(status) => status.exit_code() as i32,
229            Err(_) => -1,
230        };
231        wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst);
232        if let Ok(mut guard) = wait_exit_code.lock() {
233            *guard = Some(code);
234        }
235        let _ = exit_tx.send(code);
236    });
237
238    let handles = PtyHandles {
239        _slave: if cfg!(windows) {
240            Some(pair.slave)
241        } else {
242            None
243        },
244        _master: PtyMasterHandle::Resizable(pair.master),
245    };
246
247    let handle = ProcessHandle::new(
248        writer_tx,
249        Box::new(PtyChildTerminator {
250            killer,
251            #[cfg(unix)]
252            process_group_id,
253        }),
254        reader_handle,
255        Vec::new(),
256        writer_handle,
257        wait_handle,
258        exit_status,
259        exit_code,
260        Some(handles),
261        /*resizer*/ None,
262    );
263
264    Ok(SpawnedProcess {
265        session: handle,
266        stdout_rx,
267        stderr_rx,
268        exit_rx,
269    })
270}
271
272#[cfg(unix)]
273async fn spawn_process_preserving_fds(
274    program: &str,
275    args: &[String],
276    cwd: &Path,
277    env: &HashMap<String, String>,
278    arg0: &Option<String>,
279    size: TerminalSize,
280    inherited_fds: &[RawFd],
281) -> Result<SpawnedProcess> {
282    let (master, slave) = open_unix_pty(size)?;
283    let mut command = StdCommand::new(program);
284    if let Some(arg0) = arg0 {
285        command.arg0(arg0);
286    }
287    command.current_dir(cwd);
288    command.env_clear();
289    for arg in args {
290        command.arg(arg);
291    }
292    for (key, value) in env {
293        command.env(key, value);
294    }
295
296    // The child should see one terminal on all three stdio streams. Cloning
297    // the slave fd gives us three owned handles to the same PTY slave device
298    // so Command can wire them up independently as stdin/stdout/stderr.
299    let stdin = slave.try_clone()?;
300    let stdout = slave.try_clone()?;
301    let stderr = slave.try_clone()?;
302    let inherited_fds = inherited_fds.to_vec();
303
304    unsafe {
305        command
306            .stdin(Stdio::from(stdin))
307            .stdout(Stdio::from(stdout))
308            .stderr(Stdio::from(stderr))
309            .pre_exec(move || {
310                for signo in &[
311                    libc::SIGCHLD,
312                    libc::SIGHUP,
313                    libc::SIGINT,
314                    libc::SIGQUIT,
315                    libc::SIGTERM,
316                    libc::SIGALRM,
317                ] {
318                    libc::signal(*signo, libc::SIG_DFL);
319                }
320
321                let empty_set: libc::sigset_t = std::mem::zeroed();
322                libc::sigprocmask(libc::SIG_SETMASK, &empty_set, std::ptr::null_mut());
323
324                if libc::setsid() == -1 {
325                    return Err(std::io::Error::last_os_error());
326                }
327
328                // stdin now refers to the PTY slave, so make that fd the
329                // controlling terminal for the child's new session. stdout and
330                // stderr point at clones of the same slave device.
331                #[allow(clippy::cast_lossless)]
332                if libc::ioctl(0, libc::TIOCSCTTY as _, 0) == -1 {
333                    return Err(std::io::Error::last_os_error());
334                }
335
336                close_inherited_fds_except(&inherited_fds);
337                Ok(())
338            });
339    }
340
341    let mut child = command.spawn()?;
342    drop(slave);
343    let process_group_id = child.id();
344
345    let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
346    let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>(128);
347    let (_stderr_tx, stderr_rx) = mpsc::channel::<Vec<u8>>(1);
348    let mut reader = master.try_clone()?;
349    let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
350        let mut buf = [0u8; 8_192];
351        loop {
352            match std::io::Read::read(&mut reader, &mut buf) {
353                Ok(0) => break,
354                Ok(n) => {
355                    let _ = stdout_tx.blocking_send(buf[..n].to_vec());
356                }
357                Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
358                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
359                    std::thread::sleep(Duration::from_millis(5));
360                    continue;
361                }
362                Err(_) => break,
363            }
364        }
365    });
366
367    let writer = Arc::new(tokio::sync::Mutex::new(master.try_clone()?));
368    let writer_handle: JoinHandle<()> = tokio::spawn({
369        let writer = Arc::clone(&writer);
370        async move {
371            while let Some(bytes) = writer_rx.recv().await {
372                let mut guard = writer.lock().await;
373                use std::io::Write;
374                let _ = guard.write_all(&bytes);
375                let _ = guard.flush();
376            }
377        }
378    });
379
380    let (exit_tx, exit_rx) = oneshot::channel::<i32>();
381    let exit_status = Arc::new(AtomicBool::new(false));
382    let wait_exit_status = Arc::clone(&exit_status);
383    let exit_code = Arc::new(StdMutex::new(None));
384    let wait_exit_code = Arc::clone(&exit_code);
385    let wait_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
386        let code = match child.wait() {
387            Ok(status) => exit_code_from_status(status),
388            Err(_) => -1,
389        };
390        wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst);
391        if let Ok(mut guard) = wait_exit_code.lock() {
392            *guard = Some(code);
393        }
394        let _ = exit_tx.send(code);
395    });
396
397    let handles = PtyHandles {
398        _slave: None,
399        _master: PtyMasterHandle::Opaque {
400            raw_fd: master.as_raw_fd(),
401            _handle: Box::new(master),
402        },
403    };
404
405    let handle = ProcessHandle::new(
406        writer_tx,
407        Box::new(RawPidTerminator { process_group_id }),
408        reader_handle,
409        Vec::new(),
410        writer_handle,
411        wait_handle,
412        exit_status,
413        exit_code,
414        Some(handles),
415        /*resizer*/ None,
416    );
417
418    Ok(SpawnedProcess {
419        session: handle,
420        stdout_rx,
421        stderr_rx,
422        exit_rx,
423    })
424}
425
426#[cfg(unix)]
427fn open_unix_pty(size: TerminalSize) -> Result<(File, File)> {
428    let mut master: RawFd = -1;
429    let mut slave: RawFd = -1;
430    let mut size = libc::winsize {
431        ws_row: size.rows,
432        ws_col: size.cols,
433        ws_xpixel: 0,
434        ws_ypixel: 0,
435    };
436    let winp = std::ptr::addr_of_mut!(size);
437
438    let result = unsafe {
439        libc::openpty(
440            &mut master,
441            &mut slave,
442            std::ptr::null_mut(),
443            std::ptr::null_mut(),
444            winp,
445        )
446    };
447    if result != 0 {
448        anyhow::bail!("failed to openpty: {:?}", std::io::Error::last_os_error());
449    }
450
451    set_cloexec(master)?;
452    set_cloexec(slave)?;
453
454    Ok(unsafe { (File::from_raw_fd(master), File::from_raw_fd(slave)) })
455}
456
457#[cfg(unix)]
458fn set_cloexec(fd: RawFd) -> std::io::Result<()> {
459    let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
460    if flags == -1 {
461        return Err(std::io::Error::last_os_error());
462    }
463    let result = unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) };
464    if result == -1 {
465        return Err(std::io::Error::last_os_error());
466    }
467    Ok(())
468}
469
470#[cfg(unix)]
471pub(crate) fn close_inherited_fds_except(preserved_fds: &[RawFd]) {
472    if let Ok(dir) = std::fs::read_dir("/dev/fd") {
473        let mut fds = Vec::new();
474        for entry in dir {
475            let num = entry
476                .ok()
477                .map(|entry| entry.file_name())
478                .and_then(|name| name.into_string().ok())
479                .and_then(|name| name.parse::<RawFd>().ok());
480            if let Some(num) = num {
481                if num <= 2 || preserved_fds.contains(&num) {
482                    continue;
483                }
484                // Keep CLOEXEC descriptors open so std::process can still use
485                // its internal exec-error pipe to report spawn failures.
486                let flags = unsafe { libc::fcntl(num, libc::F_GETFD) };
487                if flags == -1 || flags & libc::FD_CLOEXEC != 0 {
488                    continue;
489                }
490                fds.push(num);
491            }
492        }
493        for fd in fds {
494            unsafe {
495                libc::close(fd);
496            }
497        }
498    }
499}