window-sand-box 0.1.0

//! 私有桌面隔离
//!
//! 可选功能：在独立的 Windows 桌面中运行沙盒进程。

use anyhow::Result;
use rand::Rng;
use rand::SeedableRng;
use rand::rngs::SmallRng;
use std::ffi::c_void;
use std::path::Path;
use std::ptr;
use windows_sys::Win32::Foundation::{CloseHandle, ERROR_SUCCESS, GetLastError, HANDLE, HLOCAL, LocalFree};
use windows_sys::Win32::Security::{
    DACL_SECURITY_INFORMATION, TOKEN_DUPLICATE, TOKEN_QUERY,
};
use windows_sys::Win32::Security::Authorization::{
    EXPLICIT_ACCESS_W, GRANT_ACCESS, SE_WINDOW_OBJECT, SetEntriesInAclW, SetSecurityInfo,
    TRUSTEE_IS_SID, TRUSTEE_IS_UNKNOWN, TRUSTEE_W,
};
use windows_sys::Win32::System::StationsAndDesktops::{
    CloseDesktop, CreateDesktopW,
};
use windows_sys::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};

// 桌面访问权限（windows-sys 0.52 未导出此常量）
const DESKTOP_ALL_ACCESS: u32 = 0x0001 | 0x0002 | 0x0004 | 0x0008 | 0x0010
    | 0x0020 | 0x0040 | 0x0080 | 0x0100 | 0x0200 | 0x0400 | 0x0800 | 0x1000;

use crate::winutil::to_wide;

/// 私有桌面封装
pub struct PrivateDesktop {
    handle: isize,
    name: String,
}

impl PrivateDesktop {
    /// 创建一个新的私有桌面
    pub fn create(_log_base: Option<&Path>) -> Result<Self> {
        let mut rng = SmallRng::from_entropy();
        let name = format!("WsbxDesktop-{:x}", rng.r#gen::<u128>());
        let name_wide = to_wide(&name);

        let handle = unsafe {
            CreateDesktopW(
                name_wide.as_ptr(),
                ptr::null(),
                ptr::null_mut(),
                0,
                DESKTOP_ALL_ACCESS,
                ptr::null_mut(),
            )
        };

        if handle == 0 {
            let err = unsafe { GetLastError() };
            return Err(anyhow::anyhow!(
                "CreateDesktopW failed for {name}: {err}"
            ));
        }

        unsafe {
            if let Err(e) = grant_desktop_access(handle) {
                let _ = CloseDesktop(handle);
                return Err(e);
            }
        }

        Ok(Self { handle, name })
    }

    /// 获取桌面的完整名称（用于 STARTUPINFO）
    pub fn desktop_name(&self) -> String {
        format!("Winsta0\\{}", self.name)
    }

    pub fn handle(&self) -> isize {
        self.handle
    }
}

impl Drop for PrivateDesktop {
    fn drop(&mut self) {
        if self.handle != 0 {
            unsafe {
                let _ = CloseDesktop(self.handle);
            }
        }
    }
}

/// 授予当前用户的 Logon SID 对桌面的访问权限
unsafe fn grant_desktop_access(handle: isize) -> Result<()> {
    let mut h_token: HANDLE = 0;
    let ok = OpenProcessToken(
        GetCurrentProcess(),
        TOKEN_QUERY | TOKEN_DUPLICATE,
        &mut h_token,
    );
    if ok == 0 {
        return Err(anyhow::anyhow!("OpenProcessToken failed: {}", GetLastError()));
    }

    let logon_bytes = crate::sandbox::token::get_logon_sid_bytes(h_token);
    CloseHandle(h_token);

    // 如果没有 Logon SID（某些环境下不存在），则跳过桌面 ACL 授权
    let logon_bytes = match logon_bytes {
        Some(bytes) => bytes,
        None => return Ok(()),
    };

    let mut logon = logon_bytes;
    let psid = logon.as_mut_ptr() as *mut c_void;

    let entries = [EXPLICIT_ACCESS_W {
        grfAccessPermissions: DESKTOP_ALL_ACCESS,
        grfAccessMode: GRANT_ACCESS,
        grfInheritance: 0,
        Trustee: TRUSTEE_W {
            pMultipleTrustee: ptr::null_mut(),
            MultipleTrusteeOperation: 0,
            TrusteeForm: TRUSTEE_IS_SID,
            TrusteeType: TRUSTEE_IS_UNKNOWN,
            ptstrName: psid as *mut u16,
        },
    }];

    let mut updated_dacl: *mut windows_sys::Win32::Security::ACL = ptr::null_mut();
    let code = SetEntriesInAclW(1, entries.as_ptr(), ptr::null_mut(), &mut updated_dacl);
    if code != ERROR_SUCCESS {
        return Err(anyhow::anyhow!("SetEntriesInAclW failed: {code}"));
    }

    let code2 = SetSecurityInfo(
        handle,
        SE_WINDOW_OBJECT,
        DACL_SECURITY_INFORMATION,
        ptr::null_mut(),
        ptr::null_mut(),
        updated_dacl,
        ptr::null_mut(),
    );

    if !updated_dacl.is_null() {
        LocalFree(updated_dacl as HLOCAL);
    }

    if code2 != ERROR_SUCCESS {
        return Err(anyhow::anyhow!("SetSecurityInfo failed: {code2}"));
    }

    Ok(())
}