//! Windows command runner used by the **elevated** sandbox path.
//!
//! The CLI launches this binary under the sandbox user when Windows sandbox level is
//! Elevated. It connects to the IPC pipes, reads the framed `SpawnRequest`, derives a
//! restricted token from the sandbox user, and spawns the child process via ConPTY
//! (`tty=true`) or pipes (`tty=false`). It then streams output frames back to the parent,
//! accepts stdin/terminate frames, and emits a final exit frame. The legacy restricted‑token
//! path spawns the child directly and does not use this runner.

#![allow(unsafe_op_in_unsafe_fn)]

mod cwd_junction;

use anyhow::Context;
use anyhow::Result;
use codex_windows_sandbox::ErrorPayload;
use codex_windows_sandbox::ExitPayload;
use codex_windows_sandbox::FramedMessage;
use codex_windows_sandbox::LocalSid;
use codex_windows_sandbox::Message;
use codex_windows_sandbox::OutputPayload;
use codex_windows_sandbox::OutputStream;
use codex_windows_sandbox::PipeSpawnHandles;
use codex_windows_sandbox::ResizePayload;
use codex_windows_sandbox::SandboxPolicy;
use codex_windows_sandbox::SpawnReady;
use codex_windows_sandbox::SpawnRequest;
use codex_windows_sandbox::StderrMode;
use codex_windows_sandbox::StdinMode;
use codex_windows_sandbox::allow_null_device;
use codex_windows_sandbox::create_readonly_token_with_caps_and_user_from;
use codex_windows_sandbox::create_workspace_write_token_with_caps_and_user_from;
use codex_windows_sandbox::decode_bytes;
use codex_windows_sandbox::encode_bytes;
use codex_windows_sandbox::get_current_token_for_restriction;
use codex_windows_sandbox::hide_current_user_profile_dir;
use codex_windows_sandbox::log_note;
use codex_windows_sandbox::parse_policy;
use codex_windows_sandbox::read_frame;
use codex_windows_sandbox::read_handle_loop;
use codex_windows_sandbox::spawn_process_with_pipes;
use codex_windows_sandbox::to_wide;
use codex_windows_sandbox::write_frame;
use std::ffi::OsStr;
use std::fs::File;
use std::os::windows::io::FromRawHandle;
use std::path::Path;
use std::path::PathBuf;
use std::ptr;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use windows_sys::Win32::Foundation::CloseHandle;
use windows_sys::Win32::Foundation::ERROR_FILE_NOT_FOUND;
use windows_sys::Win32::Foundation::GetLastError;
use windows_sys::Win32::Foundation::HANDLE;
use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
use windows_sys::Win32::Storage::FileSystem::CreateFileW;
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::ResizePseudoConsole;
use windows_sys::Win32::System::JobObjects::AssignProcessToJobObject;
use windows_sys::Win32::System::JobObjects::CreateJobObjectW;
use windows_sys::Win32::System::JobObjects::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
use windows_sys::Win32::System::JobObjects::JOBOBJECT_EXTENDED_LIMIT_INFORMATION;
use windows_sys::Win32::System::JobObjects::JobObjectExtendedLimitInformation;
use windows_sys::Win32::System::JobObjects::SetInformationJobObject;
use windows_sys::Win32::System::Threading::GetExitCodeProcess;
use windows_sys::Win32::System::Threading::GetProcessId;
use windows_sys::Win32::System::Threading::INFINITE;
use windows_sys::Win32::System::Threading::MUTEX_ALL_ACCESS;
use windows_sys::Win32::System::Threading::OpenMutexW;
use windows_sys::Win32::System::Threading::PROCESS_INFORMATION;
use windows_sys::Win32::System::Threading::TerminateProcess;
use windows_sys::Win32::System::Threading::WaitForSingleObject;

const READ_ACL_MUTEX_NAME: &str = "Local\\CodexSandboxReadAcl";
const WAIT_TIMEOUT: u32 = 0x0000_0102;

struct IpcSpawnedProcess {
    log_dir: PathBuf,
    pi: PROCESS_INFORMATION,
    stdout_handle: HANDLE,
    stderr_handle: HANDLE,
    stdin_handle: Option<HANDLE>,
    conpty_owner: Option<codex_windows_sandbox::ConptyInstance>,
    hpc_handle: Option<HANDLE>,
    _pipe_handles: Option<PipeSpawnHandles>,
}

/// Small RAII wrapper for raw Win32 handles.
///
/// The elevated runner has a few early-return paths where we acquire a token, job, or pipe
/// handle and then may fail while preparing the child. Keeping those handles in a guard makes
/// the error paths read more directly and closes the gaps that were previously leaking them.
struct OwnedWinHandle(HANDLE);

impl OwnedWinHandle {
    fn new(handle: HANDLE) -> Self {
        Self(handle)
    }

    fn raw(&self) -> HANDLE {
        self.0
    }

    fn into_raw(mut self) -> HANDLE {
        // Transfer ownership to the caller. After this point the caller is responsible for
        // eventually closing the returned HANDLE.
        let handle = self.0;
        self.0 = 0;
        handle
    }
}

impl Drop for OwnedWinHandle {
    fn drop(&mut self) {
        if self.0 != 0 && self.0 != INVALID_HANDLE_VALUE {
            unsafe {
                CloseHandle(self.0);
            }
        }
    }
}

unsafe fn create_job_kill_on_close() -> Result<HANDLE> {
    let h_job = OwnedWinHandle::new(CreateJobObjectW(std::ptr::null_mut(), std::ptr::null()));
    if h_job.raw() == 0 {
        return Err(anyhow::anyhow!("CreateJobObjectW failed"));
    }
    let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
    limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
    let ok = SetInformationJobObject(
        h_job.raw(),
        JobObjectExtendedLimitInformation,
        &mut limits as *mut _ as *mut _,
        std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
    );
    if ok == 0 {
        return Err(anyhow::anyhow!("SetInformationJobObject failed"));
    }
    Ok(h_job.into_raw())
}

/// Open a named pipe created by the parent process.
fn open_pipe(name: &str, access: u32) -> Result<HANDLE> {
    let path = to_wide(name);
    let handle = unsafe {
        CreateFileW(
            path.as_ptr(),
            access,
            0,
            std::ptr::null_mut(),
            OPEN_EXISTING,
            0,
            0,
        )
    };
    if handle == windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE {
        let err = unsafe { GetLastError() };
        return Err(anyhow::anyhow!("CreateFileW failed for pipe {name}: {err}"));
    }
    Ok(handle)
}

/// Send an error frame back to the parent process.
fn send_error(writer: &Arc<StdMutex<File>>, code: &str, message: String) -> Result<()> {
    let msg = FramedMessage {
        version: 1,
        message: Message::Error {
            payload: ErrorPayload {
                message,
                code: code.to_string(),
            },
        },
    };
    if let Ok(mut guard) = writer.lock() {
        write_frame(&mut *guard, &msg)?;
    }
    Ok(())
}

/// Read and validate the initial spawn request frame.
fn read_spawn_request(reader: &mut File) -> Result<SpawnRequest> {
    let Some(msg) = read_frame(reader)? else {
        anyhow::bail!("runner: pipe closed before spawn_request");
    };
    if msg.version != 1 {
        anyhow::bail!("runner: unsupported protocol version {}", msg.version);
    }
    match msg.message {
        Message::SpawnRequest { payload } => Ok(*payload),
        other => anyhow::bail!("runner: expected spawn_request, got {other:?}"),
    }
}

fn read_acl_mutex_exists() -> Result<bool> {
    let name = to_wide(OsStr::new(READ_ACL_MUTEX_NAME));
    let handle = unsafe { OpenMutexW(MUTEX_ALL_ACCESS, 0, name.as_ptr()) };
    if handle == 0 {
        let err = unsafe { GetLastError() };
        if err == ERROR_FILE_NOT_FOUND {
            return Ok(false);
        }
        return Err(anyhow::anyhow!("OpenMutexW failed: {err}"));
    }
    unsafe {
        CloseHandle(handle);
    }
    Ok(true)
}

/// Pick an effective CWD, using a junction if the ACL helper is active.
fn effective_cwd(req_cwd: &Path, log_dir: Option<&Path>) -> PathBuf {
    let use_junction = match read_acl_mutex_exists() {
        Ok(exists) => exists,
        Err(err) => {
            log_note(
                &format!(
                    "junction: failed to probe ACL mutex state: {err}; defaulting to junction cwd"
                ),
                log_dir,
            );
            true
        }
    };
    if use_junction {
        cwd_junction::create_cwd_junction(req_cwd, log_dir).unwrap_or_else(|| req_cwd.to_path_buf())
    } else {
        req_cwd.to_path_buf()
    }
}

fn spawn_ipc_process(req: &SpawnRequest) -> Result<IpcSpawnedProcess> {
    let log_dir = req.codex_home.clone();
    hide_current_user_profile_dir(req.codex_home.as_path());
    let policy = parse_policy(&req.policy_json_or_preset).context("parse policy_json_or_preset")?;
    let mut cap_psids: Vec<LocalSid> = Vec::new();
    for sid in &req.cap_sids {
        cap_psids.push(
            LocalSid::from_string(sid)
                .context("ConvertStringSidToSidW failed for capability SID")?,
        );
    }
    if cap_psids.is_empty() {
        anyhow::bail!("runner: empty capability SID list");
    }

    // The token helpers still take raw SID pointers, but we keep ownership in `LocalSid`
    // wrappers for as long as possible. That way any failure after SID parsing but before the
    // child is fully spawned still releases the backing LocalAlloc memory automatically.
    let cap_psid_ptrs: Vec<*mut _> = cap_psids.iter().map(LocalSid::as_ptr).collect();
    let base = OwnedWinHandle::new(unsafe { get_current_token_for_restriction()? });
    let h_token = OwnedWinHandle::new(unsafe {
        match &policy {
            SandboxPolicy::ReadOnly { .. } => {
                create_readonly_token_with_caps_and_user_from(base.raw(), &cap_psid_ptrs)
            }
            SandboxPolicy::WorkspaceWrite { .. } => {
                create_workspace_write_token_with_caps_and_user_from(base.raw(), &cap_psid_ptrs)
            }
            SandboxPolicy::DangerFullAccess | SandboxPolicy::ExternalSandbox { .. } => {
                unreachable!()
            }
        }
    }?);
    unsafe {
        // These ACL adjustments need the raw SID values, but ownership stays with `cap_psids`.
        // We do not manually `LocalFree` anything here; the wrappers handle every return path.
        allow_null_device(cap_psid_ptrs[0]);
        for psid in &cap_psid_ptrs {
            allow_null_device(*psid);
        }
    }

    let effective_cwd = effective_cwd(&req.cwd, Some(log_dir.as_path()));

    let mut conpty_owner = None;
    let mut hpc_handle: Option<HANDLE> = None;
    let mut pipe_handles = None;
    let (pi, stdout_handle, stderr_handle, stdin_handle) = if req.tty {
        let (pi, mut conpty) = codex_windows_sandbox::spawn_conpty_process_as_user(
            h_token.raw(),
            &req.command,
            &effective_cwd,
            &req.env,
            req.use_private_desktop,
            Some(log_dir.as_path()),
        )?;
        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 {
            unsafe {
                CloseHandle(input_write);
            }
            None
        };
        (
            pi,
            output_read,
            windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE,
            stdin_handle,
        )
    } else {
        let stdin_mode = if req.stdin_open {
            StdinMode::Open
        } else {
            StdinMode::Closed
        };
        let spawned_pipes: PipeSpawnHandles = spawn_process_with_pipes(
            h_token.raw(),
            &req.command,
            &effective_cwd,
            &req.env,
            stdin_mode,
            StderrMode::Separate,
            req.use_private_desktop,
            Some(log_dir.as_path()),
        )?;
        let pi = spawned_pipes.process;
        let stdout_handle = spawned_pipes.stdout_read;
        let stderr_handle = spawned_pipes
            .stderr_read
            .unwrap_or(windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE);
        let stdin_handle = spawned_pipes.stdin_write;
        pipe_handles = Some(spawned_pipes);
        (pi, stdout_handle, stderr_handle, stdin_handle)
    };
    Ok(IpcSpawnedProcess {
        log_dir,
        pi,
        stdout_handle,
        stderr_handle,
        stdin_handle,
        conpty_owner,
        hpc_handle,
        _pipe_handles: pipe_handles,
    })
}

/// Stream stdout/stderr from the child into Output frames.
fn spawn_output_reader(
    writer: Arc<StdMutex<File>>,
    handle: HANDLE,
    stream: OutputStream,
    log_dir: Option<PathBuf>,
) -> std::thread::JoinHandle<()> {
    read_handle_loop(handle, move |chunk| {
        let msg = FramedMessage {
            version: 1,
            message: Message::Output {
                payload: OutputPayload {
                    data_b64: encode_bytes(chunk),
                    stream,
                },
            },
        };
        if let Ok(mut guard) = writer.lock()
            && let Err(err) = write_frame(&mut *guard, &msg)
        {
            log_note(
                &format!("runner output write failed: {err}"),
                log_dir.as_deref(),
            );
        }
    })
}

/// Read stdin/terminate frames and forward to the child process.
fn spawn_input_loop(
    mut reader: File,
    stdin_handle: Option<HANDLE>,
    hpc_handle: Arc<StdMutex<Option<HANDLE>>>,
    process_handle: Arc<StdMutex<Option<HANDLE>>>,
    log_dir: Option<PathBuf>,
) -> std::thread::JoinHandle<()> {
    std::thread::spawn(move || {
        let mut stdin_handle = stdin_handle;
        loop {
            let msg = match read_frame(&mut reader) {
                Ok(Some(v)) => v,
                Ok(None) => break,
                Err(_) => break,
            };
            match msg.message {
                Message::Stdin { payload } => {
                    let Ok(bytes) = decode_bytes(&payload.data_b64) else {
                        continue;
                    };
                    if let Some(handle) = stdin_handle {
                        let mut offset = 0usize;
                        // `WriteFile` can report success after consuming only part of the buffer
                        // when the target is a pipe. Treat this like a normal partial write and
                        // keep advancing until every decoded stdin byte has been forwarded.
                        //
                        // If the child closes stdin or the pipe enters an error state, we log
                        // that fact, close our local HANDLE, and stop trying to forward later
                        // `Stdin` frames. That prevents silent truncation while also avoiding an
                        // endless stream of failing writes after the child is already gone.
                        while offset < bytes.len() {
                            let chunk = &bytes[offset..];
                            let chunk_len = chunk.len().min(u32::MAX as usize);
                            let mut written = 0u32;
                            let ok = unsafe {
                                windows_sys::Win32::Storage::FileSystem::WriteFile(
                                    handle,
                                    chunk.as_ptr(),
                                    chunk_len as u32,
                                    &mut written,
                                    ptr::null_mut(),
                                )
                            };
                            if ok == 0 {
                                log_note(
                                    &format!(
                                        "runner stdin write failed after {offset} bytes: {}",
                                        unsafe { GetLastError() }
                                    ),
                                    log_dir.as_deref(),
                                );
                                unsafe {
                                    CloseHandle(handle);
                                }
                                stdin_handle = None;
                                break;
                            }
                            if written == 0 {
                                log_note(
                                    "runner stdin write made no progress; closing child stdin",
                                    log_dir.as_deref(),
                                );
                                unsafe {
                                    CloseHandle(handle);
                                }
                                stdin_handle = None;
                                break;
                            }
                            offset += written as usize;
                        }
                    }
                }
                Message::CloseStdin { .. } => {
                    if let Some(handle) = stdin_handle.take() {
                        unsafe {
                            CloseHandle(handle);
                        }
                    }
                }
                Message::Resize {
                    payload: ResizePayload { rows, cols },
                } => {
                    if let Ok(guard) = hpc_handle.lock()
                        && let Some(hpc) = guard.as_ref()
                    {
                        unsafe {
                            let _ = ResizePseudoConsole(
                                *hpc,
                                COORD {
                                    X: cols as i16,
                                    Y: rows as i16,
                                },
                            );
                        }
                    }
                }
                Message::Terminate { .. } => {
                    if let Ok(guard) = process_handle.lock()
                        && let Some(handle) = guard.as_ref()
                    {
                        unsafe {
                            let _ = TerminateProcess(*handle, 1);
                        }
                    }
                }
                Message::SpawnRequest { .. } => {}
                Message::SpawnReady { .. } => {}
                Message::Output { .. } => {}
                Message::Exit { .. } => {}
                Message::Error { .. } => {}
            }
        }
        if let Some(handle) = stdin_handle {
            unsafe {
                CloseHandle(handle);
            }
        }
    })
}

/// Entry point for the Windows command runner process.
pub fn main() -> Result<()> {
    let mut pipe_in = None;
    let mut pipe_out = None;
    for arg in std::env::args().skip(1) {
        if let Some(rest) = arg.strip_prefix("--pipe-in=") {
            pipe_in = Some(rest.to_string());
        } else if let Some(rest) = arg.strip_prefix("--pipe-out=") {
            pipe_out = Some(rest.to_string());
        }
    }

    let Some(pipe_in) = pipe_in else {
        anyhow::bail!("runner: no pipe-in provided");
    };
    let Some(pipe_out) = pipe_out else {
        anyhow::bail!("runner: no pipe-out provided");
    };

    // Open both pipe ends under guards first so a failure on the second open cannot leak the
    // first HANDLE. Only after both opens succeed do we transfer ownership into `File`, which
    // then becomes responsible for closing them.
    let h_pipe_in = OwnedWinHandle::new(open_pipe(&pipe_in, FILE_GENERIC_READ)?);
    let h_pipe_out = OwnedWinHandle::new(open_pipe(&pipe_out, FILE_GENERIC_WRITE)?);
    let mut pipe_read = unsafe { File::from_raw_handle(h_pipe_in.into_raw() as _) };
    let pipe_write = Arc::new(StdMutex::new(unsafe {
        File::from_raw_handle(h_pipe_out.into_raw() as _)
    }));

    let req = match read_spawn_request(&mut pipe_read) {
        Ok(v) => v,
        Err(err) => {
            let _ = send_error(&pipe_write, "spawn_failed", err.to_string());
            return Err(err);
        }
    };

    let ipc_spawn = match spawn_ipc_process(&req) {
        Ok(value) => value,
        Err(err) => {
            let _ = send_error(&pipe_write, "spawn_failed", err.to_string());
            return Err(err);
        }
    };
    let log_dir = Some(ipc_spawn.log_dir.as_path());
    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));

    let h_job = unsafe { create_job_kill_on_close().ok() };
    if let Some(job) = h_job {
        unsafe {
            let _ = AssignProcessToJobObject(job, pi.hProcess);
        }
    }

    let process_handle = Arc::new(StdMutex::new(Some(pi.hProcess)));

    let msg = FramedMessage {
        version: 1,
        message: Message::SpawnReady {
            payload: SpawnReady {
                process_id: unsafe { GetProcessId(pi.hProcess) },
            },
        },
    };
    if let Err(err) = if let Ok(mut guard) = pipe_write.lock() {
        write_frame(&mut *guard, &msg)
    } else {
        anyhow::bail!("runner spawn_ready write failed: pipe_write lock poisoned");
    } {
        let _ = send_error(&pipe_write, "spawn_failed", err.to_string());
        return Err(err);
    }
    let log_dir_owned = log_dir.map(Path::to_path_buf);
    let out_thread = spawn_output_reader(
        Arc::clone(&pipe_write),
        stdout_handle,
        OutputStream::Stdout,
        log_dir_owned.clone(),
    );
    let err_thread = if stderr_handle != windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE {
        Some(spawn_output_reader(
            Arc::clone(&pipe_write),
            stderr_handle,
            OutputStream::Stderr,
            log_dir_owned.clone(),
        ))
    } else {
        None
    };

    let _input_thread = spawn_input_loop(
        pipe_read,
        stdin_handle,
        Arc::clone(&hpc_handle),
        Arc::clone(&process_handle),
        log_dir_owned,
    );

    let timeout = req.timeout_ms.map(|ms| ms as u32).unwrap_or(INFINITE);
    let wait_res = unsafe { WaitForSingleObject(pi.hProcess, timeout) };
    let timed_out = wait_res == WAIT_TIMEOUT;

    let exit_code: i32;
    unsafe {
        if timed_out {
            let _ = TerminateProcess(pi.hProcess, 1);
            exit_code = 128 + 64;
        } else {
            let mut raw_exit: u32 = 1;
            GetExitCodeProcess(pi.hProcess, &mut raw_exit);
            exit_code = raw_exit as i32;
        }
        if pi.hThread != 0 {
            CloseHandle(pi.hThread);
        }
        if pi.hProcess != 0 {
            CloseHandle(pi.hProcess);
        }
        if let Some(job) = h_job {
            CloseHandle(job);
        }
    }

    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 {
        let _ = thread.join();
    }

    let exit_msg = FramedMessage {
        version: 1,
        message: Message::Exit {
            payload: ExitPayload {
                exit_code,
                timed_out,
            },
        },
    };
    if let Ok(mut guard) = pipe_write.lock()
        && let Err(err) = write_frame(&mut *guard, &exit_msg)
    {
        log_note(&format!("runner exit write failed: {err}"), log_dir);
    }

    std::process::exit(exit_code);
}
