From 23484dd8d55847510ca6e9b2c7a5059bf60fab3b Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Mon, 6 Jul 2026 16:21:34 -0300 Subject: [PATCH] feat(pty): support child fd mappings --- codex-rs/utils/pty/src/pty.rs | 196 +++++++++++++++++++++++++++++++- codex-rs/utils/pty/src/tests.rs | 115 +++++++++++++++++++ 2 files changed, 306 insertions(+), 5 deletions(-) diff --git a/codex-rs/utils/pty/src/pty.rs b/codex-rs/utils/pty/src/pty.rs index f2690c4deb..8c4292f946 100644 --- a/codex-rs/utils/pty/src/pty.rs +++ b/codex-rs/utils/pty/src/pty.rs @@ -3,8 +3,12 @@ use std::collections::HashMap; use std::fs::File; use std::io::ErrorKind; #[cfg(unix)] +use std::os::fd::AsFd; +#[cfg(unix)] use std::os::fd::AsRawFd; #[cfg(unix)] +use std::os::fd::BorrowedFd; +#[cfg(unix)] use std::os::fd::FromRawFd; #[cfg(unix)] use std::os::fd::RawFd; @@ -122,6 +126,44 @@ fn platform_native_pty_system() -> Box { } } +/// A Unix file descriptor to duplicate into an exact descriptor number in a +/// spawned PTY child. +/// +/// The source descriptor is borrowed for the duration of the spawn. The +/// parent descriptor table is never modified: duplication and remapping take +/// place in the child immediately before `exec`. +#[cfg(unix)] +#[derive(Clone, Copy, Debug)] +pub struct ChildFdMapping<'fd> { + source: BorrowedFd<'fd>, + child_fd: RawFd, +} + +#[cfg(unix)] +struct ChildFdPlan { + preserved_fds: Vec, + mappings: Vec<(RawFd, RawFd)>, +} + +#[cfg(unix)] +impl<'fd> ChildFdMapping<'fd> { + /// Maps `source` to `child_fd` in the spawned child. + /// + /// `source` may be any owned or borrowed descriptor implementing [`AsFd`]. + /// Descriptor numbers for standard input, output, and error are rejected + /// when the process is spawned because the PTY owns those streams. + #[must_use] + pub fn new(source: &'fd T, child_fd: RawFd) -> Self + where + T: AsFd + ?Sized, + { + Self { + source: source.as_fd(), + child_fd, + } + } +} + /// Spawn a process attached to a PTY, returning handles for stdin, split output, and exit. pub async fn spawn_process( program: &str, @@ -154,13 +196,108 @@ pub async fn spawn_process_with_inherited_fds( #[cfg(unix)] if !inherited_fds.is_empty() { - return spawn_process_preserving_fds(program, args, cwd, env, arg0, size, inherited_fds) - .await; + return spawn_process_preserving_fds( + program, + args, + cwd, + env, + arg0, + size, + ChildFdPlan { + preserved_fds: inherited_fds.to_vec(), + mappings: Vec::new(), + }, + ) + .await; } spawn_process_portable(program, args, cwd, env, arg0, size).await } +/// Spawn a process attached to a PTY and map Unix file descriptors to exact +/// descriptor numbers in the child. +/// +/// All sources remain untouched in the parent. In the child, every source is +/// first duplicated to a close-on-exec temporary descriptor before any target +/// is replaced, so mappings remain correct when sources and targets overlap or +/// form cycles. The resulting target descriptors survive `exec`. +/// +/// # Errors +/// +/// Returns an error if a source or target refers to standard input, output, or +/// error, if two mappings use the same target, or if the PTY child cannot be +/// spawned or configured. +#[cfg(unix)] +pub async fn spawn_process_with_fd_mappings( + program: &str, + args: &[String], + cwd: &Path, + env: &HashMap, + arg0: &Option, + size: TerminalSize, + fd_mappings: &[ChildFdMapping<'_>], +) -> Result { + if program.is_empty() { + anyhow::bail!("missing program for PTY spawn"); + } + validate_child_fd_mappings(fd_mappings)?; + + if fd_mappings.is_empty() { + return spawn_process_portable(program, args, cwd, env, arg0, size).await; + } + + let mappings = fd_mappings + .iter() + .map(|mapping| (mapping.source.as_raw_fd(), mapping.child_fd)) + .collect::>(); + let preserved_fds = mappings.iter().map(|(_, child_fd)| *child_fd).collect(); + spawn_process_preserving_fds( + program, + args, + cwd, + env, + arg0, + size, + ChildFdPlan { + preserved_fds, + mappings, + }, + ) + .await +} + +#[cfg(unix)] +fn validate_child_fd_mappings(fd_mappings: &[ChildFdMapping<'_>]) -> Result<()> { + for (index, mapping) in fd_mappings.iter().enumerate() { + let source_fd = mapping.source.as_raw_fd(); + if source_fd <= libc::STDERR_FILENO { + anyhow::bail!("PTY child fd mapping source must be greater than 2, got {source_fd}"); + } + if mapping.child_fd <= libc::STDERR_FILENO { + anyhow::bail!( + "PTY child fd mapping target must be greater than 2, got {}", + mapping.child_fd + ); + } + if mapping.child_fd == i32::MAX { + anyhow::bail!( + "PTY child fd mapping target is too large: {}", + mapping.child_fd + ); + } + if fd_mappings[..index] + .iter() + .any(|candidate| candidate.child_fd == mapping.child_fd) + { + anyhow::bail!( + "PTY child fd mapping target {} is specified more than once", + mapping.child_fd + ); + } + } + Ok(()) +} + async fn spawn_process_portable( program: &str, args: &[String], @@ -289,7 +426,7 @@ async fn spawn_process_preserving_fds( env: &HashMap, arg0: &Option, size: TerminalSize, - inherited_fds: &[RawFd], + fd_plan: ChildFdPlan, ) -> Result { let (master, slave) = open_unix_pty(size)?; let mut command = StdCommand::new(program); @@ -311,7 +448,10 @@ async fn spawn_process_preserving_fds( let stdin = slave.try_clone()?; let stdout = slave.try_clone()?; let stderr = slave.try_clone()?; - let inherited_fds = inherited_fds.to_vec(); + let ChildFdPlan { + preserved_fds, + mut mappings, + } = fd_plan; unsafe { command @@ -345,7 +485,8 @@ async fn spawn_process_preserving_fds( return Err(std::io::Error::last_os_error()); } - close_inherited_fds_except(&inherited_fds); + remap_child_fds(&mut mappings)?; + close_inherited_fds_except(&preserved_fds); Ok(()) }); } @@ -435,6 +576,51 @@ async fn spawn_process_preserving_fds( }) } +#[cfg(unix)] +fn remap_child_fds(fd_mappings: &mut [(RawFd, RawFd)]) -> std::io::Result<()> { + let Some(highest_child_fd) = fd_mappings.iter().map(|(_, child_fd)| *child_fd).max() else { + return Ok(()); + }; + let minimum_temporary_fd = highest_child_fd + .checked_add(/*rhs*/ 1) + .ok_or_else(|| std::io::Error::new(ErrorKind::InvalidInput, "child fd is too large"))?; + + // Duplicate every source before replacing any target. This makes swaps and + // longer source/target cycles safe without changing the parent fd table. + let mut duplicated_count = 0; + while duplicated_count < fd_mappings.len() { + let source_fd = fd_mappings[duplicated_count].0; + let temporary_fd = + unsafe { libc::fcntl(source_fd, libc::F_DUPFD_CLOEXEC, minimum_temporary_fd) }; + if temporary_fd == -1 { + let error = std::io::Error::last_os_error(); + close_mapping_sources(&fd_mappings[..duplicated_count]); + return Err(error); + } + fd_mappings[duplicated_count].0 = temporary_fd; + duplicated_count += 1; + } + + for &(temporary_fd, child_fd) in fd_mappings.iter() { + if unsafe { libc::dup2(temporary_fd, child_fd) } == -1 { + let error = std::io::Error::last_os_error(); + close_mapping_sources(fd_mappings); + return Err(error); + } + } + close_mapping_sources(fd_mappings); + Ok(()) +} + +#[cfg(unix)] +fn close_mapping_sources(fd_mappings: &[(RawFd, RawFd)]) { + for &(source_fd, _) in fd_mappings { + unsafe { + libc::close(source_fd); + } + } +} + #[cfg(unix)] fn open_unix_pty(size: TerminalSize) -> Result<(File, File)> { let mut master: RawFd = -1; diff --git a/codex-rs/utils/pty/src/tests.rs b/codex-rs/utils/pty/src/tests.rs index f9255176f4..452edf234f 100644 --- a/codex-rs/utils/pty/src/tests.rs +++ b/codex-rs/utils/pty/src/tests.rs @@ -10,6 +10,10 @@ use crate::combine_output_receivers; #[cfg(unix)] use crate::pipe::spawn_process_no_stdin_with_inherited_fds; #[cfg(unix)] +use crate::pty::ChildFdMapping; +#[cfg(unix)] +use crate::pty::spawn_process_with_fd_mappings; +#[cfg(unix)] use crate::pty::spawn_process_with_inherited_fds; use crate::spawn_from_driver; use crate::spawn_pipe_process; @@ -865,6 +869,117 @@ async fn pty_spawn_can_preserve_inherited_fds() -> anyhow::Result<()> { Ok(()) } +#[cfg(unix)] +fn set_close_on_exec(fd: std::os::fd::BorrowedFd<'_>) -> std::io::Result<()> { + use std::os::fd::AsRawFd; + + let flags = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFD) }; + if flags == -1 { + return Err(std::io::Error::last_os_error()); + } + if unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFD, flags | libc::FD_CLOEXEC) } == -1 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pty_spawn_maps_cloexec_fds_to_exact_child_numbers() -> anyhow::Result<()> { + use std::io::Read; + use std::io::Write; + use std::os::fd::AsFd; + use std::os::unix::net::UnixStream; + + let (mut parent_writer, child_reader) = UnixStream::pair()?; + let (child_writer, mut parent_reader) = UnixStream::pair()?; + set_close_on_exec(child_reader.as_fd())?; + set_close_on_exec(child_writer.as_fd())?; + + let mappings = [ + ChildFdMapping::new(&child_reader, /*child_fd*/ 3), + ChildFdMapping::new(&child_writer, /*child_fd*/ 4), + ]; + let script = "IFS= read -r message <&3; printf 'reply:%s' \"$message\" >&4"; + let env_map: HashMap = std::env::vars().collect(); + let spawned = spawn_process_with_fd_mappings( + "/bin/sh", + &["-c".to_string(), script.to_string()], + Path::new("."), + &env_map, + /*arg0*/ &None, + TerminalSize::default(), + &mappings, + ) + .await?; + drop(child_reader); + drop(child_writer); + + parent_writer.write_all(b"ping\n")?; + drop(parent_writer); + + let (_session, output_rx, exit_rx) = combine_spawned_output(spawned); + let (_, code) = collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 2_000).await; + assert_eq!(code, 0, "expected mapped-fd PTY child to exit cleanly"); + + let mut response = String::new(); + parent_reader.read_to_string(&mut response)?; + assert_eq!(response, "reply:ping"); + + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pty_spawn_fd_mappings_are_safe_when_sources_and_targets_swap() -> anyhow::Result<()> { + use std::io::Read; + use std::io::Write; + use std::os::fd::AsFd; + use std::os::fd::AsRawFd; + use std::os::unix::net::UnixStream; + + let (mut parent_writer, child_reader) = UnixStream::pair()?; + let (child_writer, mut parent_reader) = UnixStream::pair()?; + set_close_on_exec(child_reader.as_fd())?; + set_close_on_exec(child_writer.as_fd())?; + + let child_reader_fd = child_reader.as_raw_fd(); + let child_writer_fd = child_writer.as_raw_fd(); + let mappings = [ + ChildFdMapping::new(&child_reader, child_writer_fd), + ChildFdMapping::new(&child_writer, child_reader_fd), + ]; + let script = format!( + "IFS= read -r message <&{child_writer_fd}; printf 'reply:%s' \"$message\" >&{child_reader_fd}" + ); + let env_map: HashMap = std::env::vars().collect(); + let spawned = spawn_process_with_fd_mappings( + "/bin/sh", + &["-c".to_string(), script], + Path::new("."), + &env_map, + /*arg0*/ &None, + TerminalSize::default(), + &mappings, + ) + .await?; + drop(child_reader); + drop(child_writer); + + parent_writer.write_all(b"swap\n")?; + drop(parent_writer); + + let (_session, output_rx, exit_rx) = combine_spawned_output(spawned); + let (_, code) = collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 2_000).await; + assert_eq!(code, 0, "expected swapped-fd PTY child to exit cleanly"); + + let mut response = String::new(); + parent_reader.read_to_string(&mut response)?; + assert_eq!(response, "reply:swap"); + + Ok(()) +} + #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn pty_preserving_inherited_fds_keeps_python_repl_running() -> anyhow::Result<()> {