codex/codex-rs/windows-sandbox-service/src/package_lifecycle.rs

//! Restores owner-scoped uninstall notifications and ties file cleanup to pinned directories.
//! Registered runtime policy stays in its private module; owner tokens and pins stay here.

use std::cell::RefCell;
use std::io;
use std::os::windows::io::AsHandle;
use std::os::windows::io::AsRawHandle;
use std::os::windows::io::FromRawHandle;
use std::os::windows::io::OwnedHandle;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use anyhow::Context;
use anyhow::Result;
use anyhow::ensure;
use codex_windows_sandbox::DirectoryOpenDisposition;
use codex_windows_sandbox::PreparedWindowsSandboxCleanup;
use codex_windows_sandbox::create_directory_guard;
use codex_windows_sandbox::prepare_packaged_windows_sandbox_cleanup;
use codex_windows_sandbox::string_from_sid_bytes;
use windows::ApplicationModel::Package;
use windows::ApplicationModel::PackageCatalog;
use windows::ApplicationModel::PackageUninstallingEventArgs;
use windows::Foundation::EventRegistrationToken;
use windows::Foundation::TypedEventHandler;
use windows::Win32::System::WinRT::RO_INIT_MULTITHREADED;
use windows::Win32::System::WinRT::RoInitialize;
use windows::Win32::System::WinRT::RoUninitialize;
use windows::core::HSTRING;
use windows_sys::Win32::Foundation::HANDLE;
use windows_sys::Win32::Security as security;
use windows_sys::Win32::Storage::FileSystem as filesystem;
use windows_sys::Win32::System::RemoteDesktop::WTS_CURRENT_SERVER_HANDLE;
use windows_sys::Win32::System::RemoteDesktop::WTSEnumerateSessionsW;
use windows_sys::Win32::System::RemoteDesktop::WTSFreeMemory;
use windows_sys::Win32::System::RemoteDesktop::WTSQueryUserToken;

use crate::installation_record::InstallationRecord;

mod cleanup;
mod registered;

pub(crate) use registered::runtime_owner_removed;

struct UserInstallation {
    codex_home: Option<PathBuf>,
    // Ancestors and home remain pinned until owner-scoped cleanup finishes.
    directory_handles: Vec<OwnedHandle>,
    directory_guard: Option<OwnedHandle>,
    record: InstallationRecord,
    user_token: OwnedHandle,
    catalog: PackageCatalog,
    token: EventRegistrationToken,
}

pub(crate) struct PackageLifecycle {
    package_name: HSTRING,
    preinstalled: bool,
    uninstalling: Arc<AtomicBool>,
    installation: RefCell<Option<UserInstallation>>,
}

impl PackageLifecycle {
    pub(crate) fn new(uninstalling: Arc<AtomicBool>) -> Result<Self> {
        unsafe { RoInitialize(RO_INIT_MULTITHREADED) }
            .context("initialize the sandbox service Windows Runtime apartment")?;
        let package = Package::Current()?;
        // Windows removes OEM metadata on uninstall; capture it while the package is installed.
        let preinstalled = package
            .InstalledLocation()
            .and_then(|location| location.Path())
            .is_ok_and(|path| {
                PathBuf::from(path.to_os_string())
                    .join("microsoft.system.package.metadata")
                    .join("Custom.data")
                    .is_file()
            });
        Ok(Self {
            package_name: package.Id()?.FullName()?,
            preinstalled,
            uninstalling,
            installation: RefCell::default(),
        })
    }

    pub(crate) fn register_authenticated_user(
        &self,
        mut record: InstallationRecord,
        user_token: OwnedHandle,
    ) -> Result<InstallationRecord> {
        let mut active = self.installation.borrow_mut();
        let previous = match active.as_ref() {
            Some(installation) => Some(installation.record.clone()),
            None => crate::installation_record::load()?,
        };
        if let Some(previous) = previous {
            // A missing watcher does not retire its owner; other clients can use elevated setup.
            ensure!(
                previous.user_sid == record.user_sid && previous.codex_home == record.codex_home,
                crate::ipc::ServiceUnavailable(
                    "installation is already registered to a different owner or home"
                )
            );
            record.desktop_installation = previous
                .desktop_installation
                .or(record.desktop_installation);
        }
        record = crate::installation_record::save(record)?;
        if let Some(installation) = active.as_mut()
            && installation.codex_home.is_some()
        {
            // A restored watcher must immediately use newly registered desktop ownership.
            installation.record = record.clone();
            return Ok(record);
        }

        let saved_record = record.clone();
        with_owner_impersonation(user_token.as_raw_handle(), || {
            let mut directory_handles = Vec::new();
            let mut directory_guard = None;
            let codex_home = match crate::ipc::pin_existing_ancestors(
                &record.codex_home,
                &mut directory_handles,
            )
            .and_then(|()| {
                let home = directory_handles
                    .last()
                    .context("pin the registered home")?;
                let guard = create_directory_guard(home.as_handle())?;
                // Reject a conversion that happened before the handle-relative guard was created.
                drop(crate::ipc::pin_directory(
                    &record.codex_home,
                    filesystem::FILE_READ_ATTRIBUTES,
                    DirectoryOpenDisposition::OpenExisting,
                )?);
                directory_guard = Some(guard);
                Ok(())
            }) {
                Ok(()) => Some(record.codex_home.clone()),
                Err(error) => {
                    directory_handles.clear();
                    crate::service::log_error(
                        crate::service::EVENT_SERVICE_FAILED,
                        &format!(
                            "skipping sandbox file cleanup because the home could not be pinned: {error:#}"
                        ),
                    );
                    None
                }
            };
            if let Some(installation) = active.as_mut() {
                installation.codex_home = codex_home;
                installation.directory_handles = directory_handles;
                installation.directory_guard = directory_guard;
                installation.record = record;
                return Ok(());
            }
            let catalog = PackageCatalog::OpenForCurrentUser().context(
                crate::ipc::ServiceUnavailable("open the authenticated user's package catalog"),
            )?;
            let package_name = self.package_name.clone();
            let uninstalling = Arc::clone(&self.uninstalling);
            let owner_sid = record.user_sid.clone();
            let token = catalog
                .PackageUninstalling(&TypedEventHandler::<
                    PackageCatalog,
                    PackageUninstallingEventArgs,
                >::new(move |_, event| {
                    if let Some(event) = event
                        && event.Package()?.Id()?.FullName()? == package_name
                    {
                        registered::on_package_uninstalling(event, &owner_sid, &uninstalling)?;
                    }
                    Ok(())
                }))
                .context(crate::ipc::ServiceUnavailable(
                    "subscribe to authenticated package uninstall notifications",
                ))?;
            active.replace(UserInstallation {
                codex_home,
                directory_handles,
                directory_guard,
                record,
                user_token,
                catalog,
                token,
            });
            Ok(())
        })?;
        Ok(saved_record)
    }

    pub(crate) fn restore_logged_in_owner(&self, recorded_session_id: u32) -> Result<()> {
        let Err(error) = self.restore_authenticated_user(recorded_session_id) else {
            return Ok(());
        };

        // Session IDs can change while the service is stopped; the owner SID is durable.
        let mut sessions = std::ptr::null_mut();
        let mut count = 0;
        if unsafe {
            WTSEnumerateSessionsW(
                WTS_CURRENT_SERVER_HANDLE,
                /*reserved*/ 0,
                /*version*/ 1,
                &mut sessions,
                &mut count,
            )
        } == 0
        {
            return Err(io::Error::last_os_error()).context("list current Windows sessions");
        }
        let restored = (0..count).any(|index| {
            let session_id = unsafe { (*sessions.add(index as usize)).SessionId };
            session_id != recorded_session_id && self.restore_authenticated_user(session_id).is_ok()
        });
        unsafe { WTSFreeMemory(sessions.cast()) };
        if restored { Ok(()) } else { Err(error) }
    }

    pub(crate) fn restore_authenticated_user(&self, session_id: u32) -> Result<()> {
        if self.installation.borrow().is_some() {
            return Ok(());
        }
        let Some(record) = crate::installation_record::load()? else {
            return Ok(());
        };

        if record.runtime.is_some()
            && !crate::installation_record::is_current_package_family(&record)?
        {
            return Ok(());
        }
        let mut raw_token = std::ptr::null_mut();
        if unsafe { WTSQueryUserToken(session_id, &mut raw_token) } == 0 {
            return Err(io::Error::last_os_error()).context("open the logged-in user's token");
        }
        // SAFETY: WTSQueryUserToken transferred this handle on success.
        let token = unsafe { OwnedHandle::from_raw_handle(raw_token) };
        let user = unsafe { codex_windows_sandbox::get_user_sid_bytes(token.as_raw_handle()) }?;
        let user_sid = string_from_sid_bytes(&user).map_err(anyhow::Error::msg)?;
        ensure!(
            user_sid == record.user_sid,
            "logged-in user does not match the recorded sandbox owner"
        );

        self.register_authenticated_user(
            InstallationRecord {
                session_id,
                ..record
            },
            token,
        )
        .map(|_| ())
    }

    pub(crate) fn clean_up(&self) -> Result<()> {
        let _setup_lock =
            codex_windows_sandbox::acquire_sandbox_setup_lock(/*timeout_ms*/ 5_000)?;
        if let Some(record) = crate::installation_record::load_runtime()? {
            return registered::clean_up(self, record);
        }
        // Preserve the legacy path's ownership retirement before native preparation.
        crate::installation_record::remove()?;
        self.clean_up_resources(
            &prepare_packaged_windows_sandbox_cleanup()?,
            /*runtime*/ None,
        )
    }

    fn clean_up_resources(
        &self,
        prepared: &PreparedWindowsSandboxCleanup,
        runtime: Option<&InstallationRecord>,
    ) -> Result<()> {
        let mut installation = self.installation.borrow_mut();
        let installation = installation
            .as_mut()
            .context("the authenticated package installation was not recorded")?;
        cleanup::clean_up(
            installation,
            prepared,
            runtime,
            &self.uninstalling,
            /*preinstalled*/ self.preinstalled,
        )
    }
}

pub(crate) fn with_owner_impersonation<T>(
    user_token: HANDLE,
    operation: impl FnOnce() -> Result<T>,
) -> Result<T> {
    if unsafe { security::ImpersonateLoggedOnUser(user_token) } == 0 {
        return Err(io::Error::last_os_error()).context("impersonate the sandbox owner");
    }
    let result = operation();
    if unsafe { security::RevertToSelf() } == 0 {
        crate::service::log_error(
            crate::service::EVENT_SERVICE_FAILED,
            &format!(
                "unable to revert sandbox-owner impersonation: {}",
                io::Error::last_os_error()
            ),
        );
        // Continuing as the user would make later machine cleanup unsafe.
        std::process::abort();
    }
    result
}

impl Drop for PackageLifecycle {
    fn drop(&mut self) {
        if let Some(installation) = self.installation.get_mut().take() {
            let _ = installation
                .catalog
                .RemovePackageUninstalling(installation.token);
        }
        unsafe { RoUninitialize() };
    }
}