checkpoint

codex-rs/windows-sandbox-rs/src/process.rs

@@ -21,11 +21,12 @@ use windows_sys::Win32::System::JobObjects::JobObjectExtendedLimitInformation;
 use windows_sys::Win32::System::JobObjects::SetInformationJobObject;
 use windows_sys::Win32::System::JobObjects::JOBOBJECT_EXTENDED_LIMIT_INFORMATION;
 use windows_sys::Win32::System::JobObjects::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
-use windows_sys::Win32::System::Threading::CreateProcessAsUserW;
 use windows_sys::Win32::System::Threading::GetExitCodeProcess;
+use windows_sys::Win32::System::Threading::CreateProcessWithTokenW;
 use windows_sys::Win32::System::Threading::WaitForSingleObject;
 use windows_sys::Win32::System::Threading::CREATE_UNICODE_ENVIRONMENT;
 use windows_sys::Win32::System::Threading::INFINITE;
+use windows_sys::Win32::System::Threading::LOGON_WITH_PROFILE;
 use windows_sys::Win32::System::Threading::PROCESS_INFORMATION;
 use windows_sys::Win32::System::Threading::STARTF_USESTDHANDLES;
 use windows_sys::Win32::System::Threading::STARTUPINFOW;
@@ -79,6 +80,7 @@ fn quote_arg(a: &str) -> String {
     out.push('"');
     out
 }
+#[allow(dead_code)]
 unsafe fn ensure_inheritable_stdio(si: &mut STARTUPINFOW) -> Result<()> {
     for kind in [STD_INPUT_HANDLE, STD_OUTPUT_HANDLE, STD_ERROR_HANDLE] {
         let h = GetStdHandle(kind);
@@ -96,12 +98,16 @@ unsafe fn ensure_inheritable_stdio(si: &mut STARTUPINFOW) -> Result<()> {
     Ok(())
 }
 
+/// # Safety
+/// Caller must provide a valid primary token handle (`h_token`) with appropriate access,
+/// and the `argv`, `cwd`, and `env_map` must remain valid for the duration of the call.
 pub unsafe fn create_process_as_user(
     h_token: HANDLE,
     argv: &[String],
     cwd: &Path,
     env_map: &HashMap<String, String>,
     logs_base_dir: Option<&Path>,
+    stdio: Option<(HANDLE, HANDLE, HANDLE)>,
 ) -> Result<(PROCESS_INFORMATION, STARTUPINFOW)> {
     let cmdline_str = argv
         .iter()
@@ -117,17 +123,22 @@ pub unsafe fn create_process_as_user(
     // Point explicitly at the interactive desktop.
     let desktop = to_wide("Winsta0\\Default");
     si.lpDesktop = desktop.as_ptr() as *mut u16;
-    ensure_inheritable_stdio(&mut si)?;
+    if let Some((stdin_h, stdout_h, stderr_h)) = stdio {
+        si.dwFlags |= STARTF_USESTDHANDLES;
+        si.hStdInput = stdin_h;
+        si.hStdOutput = stdout_h;
+        si.hStdError = stderr_h;
+    } else {
+        ensure_inheritable_stdio(&mut si)?;
+    }
     let mut pi: PROCESS_INFORMATION = std::mem::zeroed();
-    let ok = CreateProcessAsUserW(
+    let ok = CreateProcessWithTokenW(
         h_token,
+        LOGON_WITH_PROFILE,
         std::ptr::null(),
         cmdline.as_mut_ptr(),
-        std::ptr::null_mut(),
-        std::ptr::null_mut(),
-        1,
         CREATE_UNICODE_ENVIRONMENT,
-        env_block.as_ptr() as *mut c_void,
+        env_block.as_ptr() as *const c_void,
         to_wide(cwd).as_ptr(),
         &si,
         &mut pi,
@@ -135,7 +146,7 @@ pub unsafe fn create_process_as_user(
     if ok == 0 {
         let err = GetLastError() as i32;
         let msg = format!(
-            "CreateProcessAsUserW failed: {} ({}) | cwd={} | cmd={} | env_u16_len={} | si_flags={}",
+            "CreateProcessWithTokenW failed: {} ({}) | cwd={} | cmd={} | env_u16_len={} | si_flags={}",
             err,
             format_last_error(err),
             cwd.display(),
@@ -144,11 +155,14 @@ pub unsafe fn create_process_as_user(
             si.dwFlags,
         );
         logging::debug_log(&msg, logs_base_dir);
-        return Err(anyhow!("CreateProcessAsUserW failed: {}", err));
+        return Err(anyhow!("CreateProcessWithTokenW failed: {}", err));
     }
     Ok((pi, si))
 }
 
+/// # Safety
+/// Caller must provide valid process information handles.
+#[allow(dead_code)]
 pub unsafe fn wait_process_and_exitcode(pi: &PROCESS_INFORMATION) -> Result<i32> {
     let res = WaitForSingleObject(pi.hProcess, INFINITE);
     if res != 0 {
@@ -161,6 +175,9 @@ pub unsafe fn wait_process_and_exitcode(pi: &PROCESS_INFORMATION) -> Result<i32>
     Ok(code as i32)
 }
 
+/// # Safety
+/// Caller must close the returned job handle.
+#[allow(dead_code)]
 pub unsafe fn create_job_kill_on_close() -> Result<HANDLE> {
     let h = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null());
     if h == 0 {
@@ -183,6 +200,9 @@ pub unsafe fn create_job_kill_on_close() -> Result<HANDLE> {
     Ok(h)
 }
 
+/// # Safety
+/// Caller must pass valid handles for a job object and a process.
+#[allow(dead_code)]
 pub unsafe fn assign_to_job(h_job: HANDLE, h_process: HANDLE) -> Result<()> {
     if AssignProcessToJobObject(h_job, h_process) == 0 {
         return Err(anyhow!(