diff --git a/codex-rs/utils/pty/src/lib.rs b/codex-rs/utils/pty/src/lib.rs index f52415665..39fc9b552 100644 --- a/codex-rs/utils/pty/src/lib.rs +++ b/codex-rs/utils/pty/src/lib.rs @@ -34,4 +34,6 @@ 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::PsuedoCon; +#[cfg(windows)] pub use win::conpty::RawConPty; diff --git a/codex-rs/utils/pty/src/win/conpty.rs b/codex-rs/utils/pty/src/win/conpty.rs index ae490160a..c147b6d10 100644 --- a/codex-rs/utils/pty/src/win/conpty.rs +++ b/codex-rs/utils/pty/src/win/conpty.rs @@ -30,8 +30,8 @@ use portable_pty::PtySystem; use portable_pty::SlavePty; use portable_pty::cmdbuilder::CommandBuilder; use std::mem::ManuallyDrop; -use std::os::windows::io::AsRawHandle; use std::os::windows::io::RawHandle; +use std::ptr; use std::sync::Arc; use std::sync::Mutex; use winapi::um::wincon::COORD; @@ -82,13 +82,15 @@ impl RawConPty { self.con.raw_handle() } - pub fn into_raw_handles(self) -> (RawHandle, RawHandle, RawHandle) { + pub fn into_handles(self) -> (PsuedoCon, FileDescriptor, FileDescriptor) { let me = ManuallyDrop::new(self); - ( - me.con.raw_handle(), - me.input_write.as_raw_handle(), - me.output_read.as_raw_handle(), - ) + unsafe { + ( + ptr::read(&me.con), + ptr::read(&me.input_write), + ptr::read(&me.output_read), + ) + } } } diff --git a/codex-rs/utils/pty/src/win/mod.rs b/codex-rs/utils/pty/src/win/mod.rs index 33b6e52fb..cfc53cb51 100644 --- a/codex-rs/utils/pty/src/win/mod.rs +++ b/codex-rs/utils/pty/src/win/mod.rs @@ -49,6 +49,7 @@ mod procthreadattr; mod psuedocon; pub use conpty::ConPtySystem; +pub use psuedocon::PsuedoCon; pub use psuedocon::conpty_supported; #[derive(Debug)] diff --git a/codex-rs/windows-sandbox-rs/src/conpty/mod.rs b/codex-rs/windows-sandbox-rs/src/conpty/mod.rs index 54d1f3428..cd839952b 100644 --- a/codex-rs/windows-sandbox-rs/src/conpty/mod.rs +++ b/codex-rs/windows-sandbox-rs/src/conpty/mod.rs @@ -12,15 +12,16 @@ use crate::winutil::format_last_error; use crate::winutil::quote_windows_arg; use crate::winutil::to_wide; use anyhow::Result; +use codex_utils_pty::PsuedoCon; use codex_utils_pty::RawConPty; use std::collections::HashMap; use std::ffi::c_void; +use std::os::windows::io::IntoRawHandle; use std::path::Path; use windows_sys::Win32::Foundation::CloseHandle; use windows_sys::Win32::Foundation::GetLastError; use windows_sys::Win32::Foundation::HANDLE; use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE; -use windows_sys::Win32::System::Console::ClosePseudoConsole; use windows_sys::Win32::System::Threading::CREATE_UNICODE_ENVIRONMENT; use windows_sys::Win32::System::Threading::CreateProcessAsUserW; use windows_sys::Win32::System::Threading::EXTENDED_STARTUPINFO_PRESENT; @@ -32,10 +33,10 @@ use crate::process::make_env_block; /// Owns a ConPTY handle and its backing pipe handles. pub struct ConptyInstance { - pub hpc: HANDLE, - pub input_write: HANDLE, - pub output_read: HANDLE, - desktop: Option, + pseudoconsole: Option, + input_write: HANDLE, + output_read: HANDLE, + _desktop: Option, } impl Drop for ConptyInstance { @@ -47,19 +48,24 @@ impl Drop for ConptyInstance { if self.output_read != 0 && self.output_read != INVALID_HANDLE_VALUE { CloseHandle(self.output_read); } - if self.hpc != 0 && self.hpc != INVALID_HANDLE_VALUE { - ClosePseudoConsole(self.hpc); - } } + let _ = self.pseudoconsole.take(); } } impl ConptyInstance { - /// Consume the instance and return raw handles without closing them. - pub fn into_raw(self) -> (HANDLE, HANDLE, HANDLE, Option) { - let me = std::mem::ManuallyDrop::new(self); - let desktop = unsafe { std::ptr::read(&me.desktop) }; - (me.hpc, me.input_write, me.output_read, desktop) + pub fn raw_handle(&self) -> Option { + self.pseudoconsole + .as_ref() + .map(|pseudoconsole| pseudoconsole.raw_handle() as HANDLE) + } + + pub fn take_input_write(&mut self) -> HANDLE { + std::mem::replace(&mut self.input_write, 0) + } + + pub fn take_output_read(&mut self) -> HANDLE { + std::mem::replace(&mut self.output_read, 0) } } @@ -70,13 +76,13 @@ impl ConptyInstance { #[allow(dead_code)] pub fn create_conpty(cols: i16, rows: i16) -> Result { let raw = RawConPty::new(cols, rows)?; - let (hpc, input_write, output_read) = raw.into_raw_handles(); + let (pseudoconsole, input_write, output_read) = raw.into_handles(); Ok(ConptyInstance { - hpc: hpc as HANDLE, - input_write: input_write as HANDLE, - output_read: output_read as HANDLE, - desktop: None, + pseudoconsole: Some(pseudoconsole), + input_write: input_write.into_raw_handle() as HANDLE, + output_read: output_read.into_raw_handle() as HANDLE, + _desktop: None, }) } @@ -109,15 +115,16 @@ pub fn spawn_conpty_process_as_user( si.StartupInfo.lpDesktop = desktop.startup_info_desktop(); let raw = RawConPty::new(/*cols*/ 80, /*rows*/ 24)?; - let (hpc, input_write, output_read) = raw.into_raw_handles(); + let (pseudoconsole, input_write, output_read) = raw.into_handles(); + let hpc = pseudoconsole.raw_handle() as HANDLE; let conpty = ConptyInstance { - hpc: hpc as HANDLE, - input_write: input_write as HANDLE, - output_read: output_read as HANDLE, - desktop: Some(desktop), + pseudoconsole: Some(pseudoconsole), + input_write: input_write.into_raw_handle() as HANDLE, + output_read: output_read.into_raw_handle() as HANDLE, + _desktop: Some(desktop), }; let mut attrs = ProcThreadAttributeList::new(/*attr_count*/ 1)?; - attrs.set_pseudoconsole(conpty.hpc)?; + attrs.set_pseudoconsole(hpc)?; si.lpAttributeList = attrs.as_mut_ptr(); let mut pi: PROCESS_INFORMATION = unsafe { std::mem::zeroed() }; diff --git a/codex-rs/windows-sandbox-rs/src/elevated/command_runner_win.rs b/codex-rs/windows-sandbox-rs/src/elevated/command_runner_win.rs index b908e7a4e..80e67044e 100644 --- a/codex-rs/windows-sandbox-rs/src/elevated/command_runner_win.rs +++ b/codex-rs/windows-sandbox-rs/src/elevated/command_runner_win.rs @@ -15,7 +15,6 @@ use anyhow::Result; use codex_windows_sandbox::ErrorPayload; use codex_windows_sandbox::ExitPayload; use codex_windows_sandbox::FramedMessage; -use codex_windows_sandbox::LaunchDesktop; use codex_windows_sandbox::LocalSid; use codex_windows_sandbox::Message; use codex_windows_sandbox::OutputPayload; @@ -57,7 +56,6 @@ use windows_sys::Win32::Storage::FileSystem::FILE_GENERIC_READ; use windows_sys::Win32::Storage::FileSystem::FILE_GENERIC_WRITE; use windows_sys::Win32::Storage::FileSystem::OPEN_EXISTING; use windows_sys::Win32::System::Console::COORD; -use windows_sys::Win32::System::Console::ClosePseudoConsole; use windows_sys::Win32::System::Console::ResizePseudoConsole; use windows_sys::Win32::System::JobObjects::AssignProcessToJobObject; use windows_sys::Win32::System::JobObjects::CreateJobObjectW; @@ -87,8 +85,8 @@ struct IpcSpawnedProcess { stdout_handle: HANDLE, stderr_handle: HANDLE, stdin_handle: Option, + conpty_owner: Option, hpc_handle: Option, - _desktop_owner: Option, _pipe_handles: Option, } @@ -263,11 +261,11 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result { let effective_cwd = effective_cwd(&req.cwd, Some(log_dir.as_path())); + let mut conpty_owner = None; let mut hpc_handle: Option = None; - let mut desktop_owner = None; let mut pipe_handles = None; let (pi, stdout_handle, stderr_handle, stdin_handle) = if req.tty { - let (pi, conpty) = codex_windows_sandbox::spawn_conpty_process_as_user( + let (pi, mut conpty) = codex_windows_sandbox::spawn_conpty_process_as_user( h_token.raw(), &req.command, &effective_cwd, @@ -275,9 +273,10 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result { req.use_private_desktop, Some(log_dir.as_path()), )?; - let (hpc, input_write, output_read, desktop) = conpty.into_raw(); - hpc_handle = Some(hpc); - desktop_owner = desktop; + hpc_handle = conpty.raw_handle(); + let input_write = conpty.take_input_write(); + let output_read = conpty.take_output_read(); + conpty_owner = Some(conpty); let stdin_handle = if req.stdin_open { Some(input_write) } else { @@ -323,8 +322,8 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result { stdout_handle, stderr_handle, stdin_handle, + conpty_owner, hpc_handle, - _desktop_owner: desktop_owner, _pipe_handles: pipe_handles, }) } @@ -526,6 +525,7 @@ pub fn main() -> Result<()> { let pi = ipc_spawn.pi; let stdout_handle = ipc_spawn.stdout_handle; let stderr_handle = ipc_spawn.stderr_handle; + let mut conpty_owner = ipc_spawn.conpty_owner; let stdin_handle = ipc_spawn.stdin_handle; let hpc_handle = Arc::new(StdMutex::new(ipc_spawn.hpc_handle)); @@ -605,13 +605,10 @@ pub fn main() -> Result<()> { } } - if let Ok(mut guard) = hpc_handle.lock() - && let Some(hpc) = guard.take() - { - unsafe { - ClosePseudoConsole(hpc); - } + if let Ok(mut guard) = hpc_handle.lock() { + let _ = guard.take(); } + drop(conpty_owner.take()); let _ = out_thread.join(); if let Some(thread) = err_thread { diff --git a/codex-rs/windows-sandbox-rs/src/lib.rs b/codex-rs/windows-sandbox-rs/src/lib.rs index 16b47f293..522f8926d 100644 --- a/codex-rs/windows-sandbox-rs/src/lib.rs +++ b/codex-rs/windows-sandbox-rs/src/lib.rs @@ -97,6 +97,8 @@ pub use cap::load_or_create_cap_sids; #[cfg(target_os = "windows")] pub use cap::workspace_cap_sid_for_cwd; #[cfg(target_os = "windows")] +pub use conpty::ConptyInstance; +#[cfg(target_os = "windows")] pub use conpty::spawn_conpty_process_as_user; #[cfg(target_os = "windows")] pub use desktop::LaunchDesktop; diff --git a/codex-rs/windows-sandbox-rs/src/unified_exec/backends/legacy.rs b/codex-rs/windows-sandbox-rs/src/unified_exec/backends/legacy.rs index ba1f15a3b..8458d5c8c 100644 --- a/codex-rs/windows-sandbox-rs/src/unified_exec/backends/legacy.rs +++ b/codex-rs/windows-sandbox-rs/src/unified_exec/backends/legacy.rs @@ -1,6 +1,7 @@ use super::windows_common::finish_driver_spawn; use super::windows_common::normalize_windows_tty_input; use crate::acl::revoke_ace; +use crate::conpty::ConptyInstance; use crate::conpty::spawn_conpty_process_as_user; use crate::desktop::LaunchDesktop; use crate::logging::log_failure; @@ -33,7 +34,6 @@ use windows_sys::Win32::Foundation::HANDLE; use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE; use windows_sys::Win32::Storage::FileSystem::WriteFile; use windows_sys::Win32::System::Console::COORD; -use windows_sys::Win32::System::Console::ClosePseudoConsole; use windows_sys::Win32::System::Console::ResizePseudoConsole; use windows_sys::Win32::System::Threading::GetExitCodeProcess; use windows_sys::Win32::System::Threading::INFINITE; @@ -48,6 +48,7 @@ struct LegacyProcessHandles { output_join: std::thread::JoinHandle<()>, writer_handle: tokio::task::JoinHandle<()>, hpc: Option, + conpty_owner: Option, token_handle: HANDLE, desktop: Option, } @@ -66,8 +67,8 @@ fn spawn_legacy_process( writer_rx: mpsc::Receiver>, logs_base_dir: Option<&Path>, ) -> Result { - let (pi, output_join, writer_handle, hpc, desktop) = if tty { - let (pi, conpty) = spawn_conpty_process_as_user( + let (pi, output_join, writer_handle, hpc, conpty_owner, desktop) = if tty { + let (pi, mut conpty) = spawn_conpty_process_as_user( h_token, command, cwd, @@ -75,14 +76,14 @@ fn spawn_legacy_process( use_private_desktop, logs_base_dir, )?; - let (hpc, input_write, output_read, desktop) = conpty.into_raw(); - let output_join = spawn_output_reader(output_read, stdout_tx); + let hpc = conpty.raw_handle(); + let output_join = spawn_output_reader(conpty.take_output_read(), stdout_tx); let writer_handle = spawn_input_writer( - Some(input_write), + Some(conpty.take_input_write()), writer_rx, /*normalize_newlines*/ true, ); - (pi, output_join, writer_handle, Some(hpc), desktop) + (pi, output_join, writer_handle, hpc, Some(conpty), None) } else { let pipe_handles = spawn_process_with_pipes( h_token, @@ -120,6 +121,7 @@ fn spawn_legacy_process( output_join, writer_handle, None, + None, Some(pipe_handles.desktop), ) }; @@ -128,6 +130,7 @@ fn spawn_legacy_process( output_join, writer_handle, hpc, + conpty_owner, token_handle: h_token, desktop, }) @@ -328,6 +331,7 @@ pub(crate) async fn spawn_windows_sandbox_session_legacy( output_join, writer_handle, hpc, + mut conpty_owner, token_handle, desktop, } = match spawn_legacy_process( @@ -386,12 +390,10 @@ pub(crate) async fn spawn_windows_sandbox_session_legacy( } if let Some(hpc) = hpc_for_wait && let Ok(mut guard) = hpc.lock() - && let Some(hpc) = guard.take() { - unsafe { - ClosePseudoConsole(hpc); - } + let _ = guard.take(); } + drop(conpty_owner.take()); unsafe { if token_handle != 0 && token_handle != INVALID_HANDLE_VALUE { CloseHandle(token_handle);