From cbe40fdfc8cbb3a4469c7e644e4712b4d0617cdf Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 25 Jun 2026 03:55:00 +0000 Subject: [PATCH] windows: add kill-on-close job process primitives --- codex-rs/utils/pty/Cargo.toml | 1 + codex-rs/utils/pty/src/lib.rs | 6 + codex-rs/utils/pty/src/win/job.rs | 285 +++++++++++++++++++ codex-rs/utils/pty/src/win/job_tests.rs | 351 ++++++++++++++++++++++++ codex-rs/utils/pty/src/win/mod.rs | 4 + 5 files changed, 647 insertions(+) create mode 100644 codex-rs/utils/pty/src/win/job.rs create mode 100644 codex-rs/utils/pty/src/win/job_tests.rs diff --git a/codex-rs/utils/pty/Cargo.toml b/codex-rs/utils/pty/Cargo.toml index f38e8f7a63..d03962edc7 100644 --- a/codex-rs/utils/pty/Cargo.toml +++ b/codex-rs/utils/pty/Cargo.toml @@ -22,6 +22,7 @@ log = { workspace = true } shared_library = "0.1.9" winapi = { version = "0.3.9", features = [ "handleapi", + "jobapi2", "minwinbase", "processthreadsapi", "synchapi", diff --git a/codex-rs/utils/pty/src/lib.rs b/codex-rs/utils/pty/src/lib.rs index ff0d168443..fb359ebd9f 100644 --- a/codex-rs/utils/pty/src/lib.rs +++ b/codex-rs/utils/pty/src/lib.rs @@ -38,8 +38,14 @@ pub use pty::conpty_supported; /// Spawn a process attached to a PTY for interactive use. pub use pty::spawn_process as spawn_pty_process; #[cfg(windows)] +pub use win::JobProcess; +#[cfg(windows)] +pub use win::KillOnCloseJob; +#[cfg(windows)] pub use win::PsuedoCon; #[cfg(windows)] +pub use win::SuspendedProcess; +#[cfg(windows)] pub use win::conpty::RawConPty; #[cfg(windows)] pub use windows_input::WindowsTtyInputNormalizer; diff --git a/codex-rs/utils/pty/src/win/job.rs b/codex-rs/utils/pty/src/win/job.rs new file mode 100644 index 0000000000..536c169807 --- /dev/null +++ b/codex-rs/utils/pty/src/win/job.rs @@ -0,0 +1,285 @@ +use std::io; +use std::mem; +use std::os::windows::io::AsRawHandle; +use std::os::windows::io::FromRawHandle; +use std::os::windows::io::IntoRawHandle; +use std::os::windows::io::OwnedHandle; +use std::os::windows::io::RawHandle; +use std::ptr; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; + +use anyhow::Context; +use winapi::shared::minwindef::FALSE; +use winapi::um::handleapi::CloseHandle; +use winapi::um::jobapi2::AssignProcessToJobObject; +use winapi::um::jobapi2::CreateJobObjectW; +use winapi::um::jobapi2::SetInformationJobObject; +use winapi::um::jobapi2::TerminateJobObject; +use winapi::um::processthreadsapi::ResumeThread; +use winapi::um::processthreadsapi::TerminateProcess; +use winapi::um::synchapi::WaitForSingleObject; +use winapi::um::winbase::INFINITE; +use winapi::um::winnt::HANDLE; +use winapi::um::winnt::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; +use winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION; +use winapi::um::winnt::JobObjectExtendedLimitInformation; + +/// Injectable Win32 operations for the create/configure/assign/resume state +/// machine. Production uses [`Win32JobObjectApi`]; tests replace individual +/// stages so every failure path can be exercised deterministically. +trait JobObjectApi { + fn create_job(&self) -> io::Result; + fn configure_kill_on_close(&self, job: RawHandle) -> io::Result<()>; + fn assign_process(&self, job: RawHandle, process: RawHandle) -> io::Result<()>; + fn resume_thread(&self, primary_thread: RawHandle) -> io::Result; +} + +struct Win32JobObjectApi; + +impl JobObjectApi for Win32JobObjectApi { + fn create_job(&self) -> io::Result { + let raw_handle = unsafe { CreateJobObjectW(ptr::null_mut(), ptr::null()) }; + if raw_handle.is_null() { + Err(io::Error::last_os_error()) + } else { + Ok(unsafe { OwnedHandle::from_raw_handle(raw_handle.cast()) }) + } + } + + fn configure_kill_on_close(&self, job: RawHandle) -> io::Result<()> { + let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { mem::zeroed() }; + limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + let result = unsafe { + SetInformationJobObject( + job as HANDLE, + JobObjectExtendedLimitInformation, + ptr::addr_of_mut!(limits).cast(), + mem::size_of::() as u32, + ) + }; + if result == FALSE { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } + } + + fn assign_process(&self, job: RawHandle, process: RawHandle) -> io::Result<()> { + let result = unsafe { AssignProcessToJobObject(job as HANDLE, process as HANDLE) }; + if result == FALSE { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } + } + + fn resume_thread(&self, primary_thread: RawHandle) -> io::Result { + let previous_suspend_count = unsafe { ResumeThread(primary_thread as HANDLE) }; + if previous_suspend_count == u32::MAX { + Err(io::Error::last_os_error()) + } else { + Ok(previous_suspend_count) + } + } +} + +fn io_error_with_context(error: io::Error, context: &'static str) -> io::Error { + let kind = error.kind(); + io::Error::new(kind, anyhow::Error::new(error).context(context)) +} + +/// Shared controller for a Windows job that terminates all members when closed. +/// +/// Clones share one underlying job handle. Calling [`Self::close`] from any +/// clone closes that handle exactly once, so process wait and cancellation paths +/// can race safely without keeping the job alive through duplicated OS handles. +#[derive(Clone, Debug)] +pub struct KillOnCloseJob { + handle: Arc>>, +} + +impl KillOnCloseJob { + /// Create an unnamed, non-inheritable job configured to kill all members + /// when its sole operating-system handle is closed. + pub fn new() -> io::Result { + Self::new_with_api(&Win32JobObjectApi) + } + + fn new_with_api(api: &impl JobObjectApi) -> io::Result { + let handle = api + .create_job() + .map_err(|err| io_error_with_context(err, "failed to create job object"))?; + api.configure_kill_on_close(handle.as_raw_handle()) + .map_err(|err| { + io_error_with_context(err, "failed to configure kill-on-close job object") + })?; + Ok(Self { + handle: Arc::new(Mutex::new(Some(handle))), + }) + } + + fn assign_process_with_api( + &self, + process: RawHandle, + api: &impl JobObjectApi, + ) -> io::Result<()> { + let guard = self.handle.lock().unwrap_or_else(PoisonError::into_inner); + let Some(job) = guard.as_ref() else { + return Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "job handle is already closed", + )); + }; + api.assign_process(job.as_raw_handle(), process) + } + + /// Close the shared job handle, terminating all processes in the job. + pub fn close(&self) -> io::Result<()> { + self.take_and_close(/*exit_code*/ None) + } + + /// Terminate all job members with `exit_code`, then close the shared job + /// handle even if explicit termination reports an error. + pub fn terminate_and_close(&self, exit_code: u32) -> io::Result<()> { + self.take_and_close(Some(exit_code)) + } + + fn take_and_close(&self, exit_code: Option) -> io::Result<()> { + let mut guard = self.handle.lock().unwrap_or_else(PoisonError::into_inner); + let Some(handle) = guard.take() else { + return Ok(()); + }; + let raw_handle = handle.into_raw_handle(); + + let terminate_result = exit_code.map(|exit_code| { + let result = unsafe { TerminateJobObject(raw_handle.cast(), exit_code) }; + if result == FALSE { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } + }); + let close_result = unsafe { CloseHandle(raw_handle.cast()) }; + let close_result = if close_result == FALSE { + Err(io::Error::last_os_error()) + } else { + Ok(()) + }; + + terminate_result.unwrap_or(Ok(())).and(close_result) + } +} + +/// Owns a newly created process and its suspended primary thread until the +/// process has been assigned to a kill-on-close job and resumed. +pub struct SuspendedProcess { + process: Option, + primary_thread: Option, + process_id: u32, +} + +impl SuspendedProcess { + /// Take ownership of raw handles returned by a successful `CreateProcess*` + /// call made with `CREATE_SUSPENDED`. + /// + /// # Safety + /// + /// `process` and `primary_thread` must be valid, owned handles for the same + /// newly created process. The primary thread must still have its initial + /// suspension count, and ownership of both handles transfers to this value. + pub unsafe fn from_raw_handles( + process: RawHandle, + primary_thread: RawHandle, + process_id: u32, + ) -> Self { + Self { + process: Some(unsafe { OwnedHandle::from_raw_handle(process) }), + primary_thread: Some(unsafe { OwnedHandle::from_raw_handle(primary_thread) }), + process_id, + } + } + + /// Assign the suspended process to `job`, then resume its primary thread. + /// Any failure leaves this guard armed, so the process is terminated and + /// waited before the error is returned. + pub fn assign_and_resume(self, job: KillOnCloseJob) -> anyhow::Result { + self.assign_and_resume_with_api(job, &Win32JobObjectApi) + } + + fn assign_and_resume_with_api( + mut self, + job: KillOnCloseJob, + api: &impl JobObjectApi, + ) -> anyhow::Result { + let process = self + .process + .as_ref() + .ok_or_else(|| io::Error::other("suspended process is missing its process handle"))?; + job.assign_process_with_api(process.as_raw_handle(), api) + .context("failed to assign suspended process to job")?; + + let primary_thread = self.primary_thread.as_ref().ok_or_else(|| { + io::Error::other("suspended process is missing its primary thread handle") + })?; + let previous_suspend_count = api + .resume_thread(primary_thread.as_raw_handle()) + .context("failed to resume suspended process")?; + if previous_suspend_count != 1 { + return Err(io::Error::other(format!( + "expected suspended process thread count 1, got {previous_suspend_count}" + ))) + .context("failed to resume suspended process"); + } + + drop(self.primary_thread.take()); + let process = self + .process + .take() + .ok_or_else(|| io::Error::other("suspended process is missing its process handle"))?; + Ok(JobProcess { + process, + process_id: self.process_id, + controller: job, + }) + } +} + +impl Drop for SuspendedProcess { + fn drop(&mut self) { + let Some(process) = self.process.as_ref() else { + return; + }; + unsafe { + let _ = TerminateProcess(process.as_raw_handle().cast(), /*uExitCode*/ 1); + let _ = WaitForSingleObject(process.as_raw_handle().cast(), INFINITE); + } + } +} + +/// A running process contained by a kill-on-close Windows job. +#[derive(Debug)] +pub struct JobProcess { + process: OwnedHandle, + process_id: u32, + controller: KillOnCloseJob, +} + +impl JobProcess { + pub fn process_id(&self) -> u32 { + self.process_id + } + + pub fn as_raw_handle(&self) -> RawHandle { + self.process.as_raw_handle() + } + + pub fn controller(&self) -> KillOnCloseJob { + self.controller.clone() + } +} + +#[cfg(test)] +#[path = "job_tests.rs"] +mod tests; diff --git a/codex-rs/utils/pty/src/win/job_tests.rs b/codex-rs/utils/pty/src/win/job_tests.rs new file mode 100644 index 0000000000..74d6157a22 --- /dev/null +++ b/codex-rs/utils/pty/src/win/job_tests.rs @@ -0,0 +1,351 @@ +use std::cell::Cell; +use std::ffi::OsStr; +use std::fs; +use std::io; +use std::mem; +use std::os::windows::ffi::OsStrExt; +use std::os::windows::io::AsRawHandle; +use std::os::windows::io::OwnedHandle; +use std::os::windows::io::RawHandle; +use std::path::Path; +use std::path::PathBuf; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::thread; +use std::time::Duration; +use std::time::Instant; + +use pretty_assertions::assert_eq; +use winapi::shared::minwindef::FALSE; +use winapi::um::minwinbase::STILL_ACTIVE; +use winapi::um::processthreadsapi::CreateProcessW; +use winapi::um::processthreadsapi::GetExitCodeProcess; +use winapi::um::processthreadsapi::PROCESS_INFORMATION; +use winapi::um::processthreadsapi::STARTUPINFOW; +use winapi::um::synchapi::WaitForSingleObject; +use winapi::um::winbase::CREATE_NO_WINDOW; +use winapi::um::winbase::CREATE_SUSPENDED; +use winapi::um::winbase::WAIT_OBJECT_0; + +use super::JobObjectApi; +use super::KillOnCloseJob; +use super::SuspendedProcess; +use super::Win32JobObjectApi; + +const PROCESS_WAIT_MS: u32 = 5_000; + +struct TestDirectory { + path: PathBuf, +} + +impl TestDirectory { + fn new() -> io::Result { + static NEXT_DIRECTORY: AtomicU64 = AtomicU64::new(0); + + let path = std::env::temp_dir().join(format!( + "codex-utils-pty-job-{}-{}", + std::process::id(), + NEXT_DIRECTORY.fetch_add(1, Ordering::Relaxed) + )); + fs::create_dir(&path)?; + Ok(Self { path }) + } + + fn join(&self, name: &str) -> PathBuf { + self.path.join(name) + } +} + +impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.path); + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum FailureStage { + Create, + Configure, + Assign, + Resume, + ResumeCount(u32), +} + +struct FailingJobObjectApi { + stage: FailureStage, +} + +impl FailingJobObjectApi { + fn failure(&self, stage: FailureStage) -> io::Result<()> { + if self.stage == stage { + Err(io::Error::other(format!("forced {stage:?} failure"))) + } else { + Ok(()) + } + } +} + +impl JobObjectApi for FailingJobObjectApi { + fn create_job(&self) -> io::Result { + self.failure(FailureStage::Create)?; + Win32JobObjectApi.create_job() + } + + fn configure_kill_on_close(&self, job: RawHandle) -> io::Result<()> { + self.failure(FailureStage::Configure)?; + Win32JobObjectApi.configure_kill_on_close(job) + } + + fn assign_process(&self, job: RawHandle, process: RawHandle) -> io::Result<()> { + self.failure(FailureStage::Assign)?; + Win32JobObjectApi.assign_process(job, process) + } + + fn resume_thread(&self, primary_thread: RawHandle) -> io::Result { + if self.stage == FailureStage::Resume { + return Err(io::Error::other("forced Resume failure")); + } + if let FailureStage::ResumeCount(count) = self.stage { + return Ok(count); + } + Win32JobObjectApi.resume_thread(primary_thread) + } +} + +fn spawn_suspended(command_line: &str) -> io::Result { + let mut command_line = OsStr::new(command_line) + .encode_wide() + .chain(Some(0)) + .collect::>(); + let mut startup_info: STARTUPINFOW = unsafe { mem::zeroed() }; + startup_info.cb = mem::size_of::() as u32; + let mut process_info: PROCESS_INFORMATION = unsafe { mem::zeroed() }; + let result = unsafe { + CreateProcessW( + std::ptr::null(), + command_line.as_mut_ptr(), + std::ptr::null_mut(), + std::ptr::null_mut(), + FALSE, + CREATE_NO_WINDOW | CREATE_SUSPENDED, + std::ptr::null_mut(), + std::ptr::null(), + &mut startup_info, + &mut process_info, + ) + }; + if result == FALSE { + return Err(io::Error::last_os_error()); + } + Ok(unsafe { + SuspendedProcess::from_raw_handles( + process_info.hProcess.cast(), + process_info.hThread.cast(), + process_info.dwProcessId, + ) + }) +} + +fn create_job_then_spawn( + api: &impl JobObjectApi, + spawn: impl FnOnce(KillOnCloseJob) -> io::Result, +) -> io::Result { + KillOnCloseJob::new_with_api(api).and_then(spawn) +} + +fn batch_command(script: &Path) -> String { + format!("cmd.exe /d /c call \"{}\"", script.display()) +} + +fn spawn_marker_process(directory: &TestDirectory) -> io::Result<(SuspendedProcess, PathBuf)> { + let script = directory.join("marker.cmd"); + let marker = directory.join("marker"); + fs::write(&script, "@echo off\r\necho ran>\"%~dp0marker\"\r\n")?; + Ok((spawn_suspended(&batch_command(&script))?, marker)) +} + +fn wait_for_file(path: &Path) -> io::Result<()> { + let start = Instant::now(); + while !path.exists() { + if start.elapsed() >= Duration::from_secs(5) { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("timed out waiting for {}", path.display()), + )); + } + thread::sleep(Duration::from_millis(10)); + } + Ok(()) +} + +fn process_observer(process: &SuspendedProcess) -> io::Result { + let Some(handle) = process.process.as_ref() else { + return Err(io::Error::other("suspended process has no process handle")); + }; + handle.try_clone() +} + +fn assert_process_exited(process: &OwnedHandle) { + let result = unsafe { + WaitForSingleObject(process.as_raw_handle().cast(), /*dwMilliseconds*/ 0) + }; + assert_eq!(result, WAIT_OBJECT_0); +} + +#[test] +fn direct_job_close_terminates_root_and_grandchild() -> anyhow::Result<()> { + let directory = TestDirectory::new()?; + let root_script = directory.join("root.cmd"); + let grandchild_script = directory.join("grandchild.cmd"); + let grandchild_ready = directory.join("grandchild-ready"); + let grandchild_escaped = directory.join("grandchild-escaped"); + fs::write( + &root_script, + "@echo off\r\nstart \"\" /b cmd.exe /d /c call \"%~dp0grandchild.cmd\"\r\nping.exe -n 30 127.0.0.1 >NUL\r\n", + )?; + fs::write( + &grandchild_script, + "@echo off\r\necho ready>\"%~dp0grandchild-ready\"\r\nping.exe -n 3 127.0.0.1 >NUL\r\necho escaped>\"%~dp0grandchild-escaped\"\r\n", + )?; + + let job = KillOnCloseJob::new()?; + let process = spawn_suspended(&batch_command(&root_script))?.assign_and_resume(job)?; + + wait_for_file(&grandchild_ready)?; + + let mut exit_code = 0; + let result = unsafe { GetExitCodeProcess(process.as_raw_handle().cast(), &mut exit_code) }; + assert_ne!(result, FALSE); + assert_eq!(exit_code, STILL_ACTIVE); + + process.controller().close()?; + let result = unsafe { WaitForSingleObject(process.as_raw_handle().cast(), PROCESS_WAIT_MS) }; + assert_eq!(result, WAIT_OBJECT_0); + + thread::sleep(Duration::from_secs(4)); + assert!(grandchild_ready.exists()); + assert!(!grandchild_escaped.exists()); + + Ok(()) +} + +#[test] +fn closing_job_is_idempotent() -> anyhow::Result<()> { + let job = KillOnCloseJob::new()?; + job.close()?; + job.close()?; + Ok(()) +} + +#[test] +fn job_creation_failure_retains_stage_context() { + let spawn_reached = Cell::new(false); + let result = create_job_then_spawn( + &FailingJobObjectApi { + stage: FailureStage::Create, + }, + |_| { + spawn_reached.set(true); + Ok(()) + }, + ); + let Err(err) = result else { + panic!("forced job creation should fail"); + }; + assert!(err.to_string().contains("failed to create job object")); + assert!(!spawn_reached.get()); +} + +#[test] +fn job_configuration_failure_retains_stage_context() { + let spawn_reached = Cell::new(false); + let result = create_job_then_spawn( + &FailingJobObjectApi { + stage: FailureStage::Configure, + }, + |_| { + spawn_reached.set(true); + Ok(()) + }, + ); + let Err(err) = result else { + panic!("forced job configuration should fail"); + }; + assert!( + err.to_string() + .contains("failed to configure kill-on-close job object") + ); + assert!(!spawn_reached.get()); +} + +#[test] +fn assignment_failure_terminates_suspended_process() -> anyhow::Result<()> { + let directory = TestDirectory::new()?; + let (process, marker) = spawn_marker_process(&directory)?; + let observer = process_observer(&process)?; + let job = KillOnCloseJob::new()?; + let result = process.assign_and_resume_with_api( + job, + &FailingJobObjectApi { + stage: FailureStage::Assign, + }, + ); + let Err(err) = result else { + panic!("forced job assignment should fail"); + }; + assert!( + err.to_string() + .contains("failed to assign suspended process to job") + ); + assert_process_exited(&observer); + assert!(!marker.exists()); + Ok(()) +} + +#[test] +fn resume_failure_terminates_suspended_process() -> anyhow::Result<()> { + let directory = TestDirectory::new()?; + let (process, marker) = spawn_marker_process(&directory)?; + let observer = process_observer(&process)?; + let job = KillOnCloseJob::new()?; + let result = process.assign_and_resume_with_api( + job, + &FailingJobObjectApi { + stage: FailureStage::Resume, + }, + ); + let Err(err) = result else { + panic!("forced process resume should fail"); + }; + assert!( + err.to_string() + .contains("failed to resume suspended process") + ); + assert_process_exited(&observer); + assert!(!marker.exists()); + Ok(()) +} + +#[test] +fn unexpected_resume_count_terminates_suspended_process() -> anyhow::Result<()> { + let directory = TestDirectory::new()?; + let (process, marker) = spawn_marker_process(&directory)?; + let observer = process_observer(&process)?; + let job = KillOnCloseJob::new()?; + let result = process.assign_and_resume_with_api( + job, + &FailingJobObjectApi { + stage: FailureStage::ResumeCount(0), + }, + ); + let Err(err) = result else { + panic!("unexpected resume count should fail"); + }; + assert!( + err.to_string() + .contains("expected suspended process thread count 1") + ); + assert_process_exited(&observer); + assert!(!marker.exists()); + Ok(()) +} diff --git a/codex-rs/utils/pty/src/win/mod.rs b/codex-rs/utils/pty/src/win/mod.rs index cfc53cb51e..70c4d07b7b 100644 --- a/codex-rs/utils/pty/src/win/mod.rs +++ b/codex-rs/utils/pty/src/win/mod.rs @@ -45,10 +45,14 @@ use winapi::um::synchapi::WaitForSingleObject; use winapi::um::winbase::INFINITE; pub(crate) mod conpty; +mod job; mod procthreadattr; mod psuedocon; pub use conpty::ConPtySystem; +pub use job::JobProcess; +pub use job::KillOnCloseJob; +pub use job::SuspendedProcess; pub use psuedocon::PsuedoCon; pub use psuedocon::conpty_supported;