window-sand-box 0.1.1

//! Windows 受限令牌创建
//!
//! 核心安全机制：使用 `CreateRestrictedToken` 创建受限令牌，
//! 
//! # Safety
//!
//! 本模块大量使用 Windows API，所有 `pub unsafe fn` 的调用者必须遵守各自的 Safety 约定。
//! 模块内 `unsafe fn` 内部直接调用其他 unsafe 操作是设计使然（Windows API 封装层），
//! 无需在每个调用点额外包裹 `unsafe {{}}`。

#![allow(unsafe_op_in_unsafe_fn)]
//! 通过 restricting SID 列表控制进程的写入权限。

use anyhow::{Result, anyhow};
use std::ffi::c_void;
use std::ptr;
use windows_sys::Win32::Foundation::{CloseHandle, GetLastError, HANDLE, HLOCAL, LUID, LocalFree};
use windows_sys::Win32::Security::{
    AdjustTokenPrivileges, CopySid,
    CreateRestrictedToken, CreateWellKnownSid,
    GetLengthSid, GetTokenInformation, LookupPrivilegeValueW,
    SID_AND_ATTRIBUTES, SetTokenInformation,
    TOKEN_ADJUST_DEFAULT, TOKEN_ADJUST_PRIVILEGES,
    TOKEN_ASSIGN_PRIMARY, TOKEN_DUPLICATE, TOKEN_PRIVILEGES, TOKEN_QUERY,
    TokenDefaultDacl, TokenGroups,
};
use windows_sys::Win32::Security::Authorization::{
    ConvertStringSidToSidW, EXPLICIT_ACCESS_W, GRANT_ACCESS, SetEntriesInAclW,
    TRUSTEE_IS_SID, TRUSTEE_IS_UNKNOWN, TRUSTEE_W,
};
use windows_sys::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};

// 令牌标志
const DISABLE_MAX_PRIVILEGE: u32 = 0x01;
const LUA_TOKEN: u32 = 0x04;
const WRITE_RESTRICTED: u32 = 0x08;
const SE_GROUP_LOGON_ID: u32 = 0xC0000000;
const SE_PRIVILEGE_ENABLED: u32 = 0x0000_0002;
const GENERIC_ALL: u32 = 0x1000_0000;
const WIN_WORLD_SID: i32 = 1;

/// 管理通过 `ConvertStringSidToSidW` 分配的 SID，自动释放
pub struct LocalSid {
    psid: *mut c_void,
}

impl LocalSid {
    /// 从 SID 字符串创建
    pub fn from_string(sid: &str) -> Result<Self> {
        let wide: Vec<u16> = crate::winutil::to_wide(sid);
        let mut psid: *mut c_void = ptr::null_mut();
        let ok = unsafe { ConvertStringSidToSidW(wide.as_ptr(), &mut psid) };
        if ok == 0 || psid.is_null() {
            return Err(anyhow!("ConvertStringSidToSidW failed for: {sid}"));
        }
        Ok(Self { psid })
    }

    pub fn as_ptr(&self) -> *mut c_void {
        self.psid
    }
}

impl Drop for LocalSid {
    fn drop(&mut self) {
        if !self.psid.is_null() {
            unsafe { LocalFree(self.psid as HLOCAL); }
        }
    }
}

// # Safety: `LocalSid` 包含裸指针，Send 是安全的因为所有权在线程间转移时
// 只有一个线程持有。但 NOT Sync：多个线程通过共享引用访问时可能同时读取
// 指针指向的 SID 数据，而 as_ptr() 返回 *mut c_void 允许未同步的可变访问，
// 且 drop 中的 LocalFree 与并发读取构成数据竞争。
// 如需跨线程共享，使用 Arc<Mutex<LocalSid>> 等同步原语包装。
unsafe impl Send for LocalSid {}

/// 获取 Everyone SID
unsafe fn world_sid() -> Result<Vec<u8>> {
    let mut size: u32 = 0;
    CreateWellKnownSid(WIN_WORLD_SID, ptr::null_mut(), ptr::null_mut(), &mut size);
    let mut buf: Vec<u8> = vec![0u8; size as usize];
    let ok = CreateWellKnownSid(
        WIN_WORLD_SID,
        ptr::null_mut(),
        buf.as_mut_ptr() as *mut c_void,
        &mut size,
    );
    if ok == 0 {
        return Err(anyhow!("CreateWellKnownSid failed: {}", GetLastError()));
    }
    Ok(buf)
}

/// 获取当前进程令牌
unsafe fn get_current_token() -> Result<HANDLE> {
    let desired = TOKEN_DUPLICATE | TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY
        | TOKEN_ADJUST_DEFAULT | TOKEN_ADJUST_PRIVILEGES;
    let mut h: HANDLE = 0;
    let ok = OpenProcessToken(GetCurrentProcess(), desired, &mut h);
    if ok == 0 {
        return Err(anyhow!("OpenProcessToken failed: {}", GetLastError()));
    }
    Ok(h)
}

/// 获取令牌的 Logon SID 字节（可能不存在）
///
/// 某些环境（如通过 cargo run 启动、某些服务账户）的令牌中没有 Logon SID，
/// 此时返回 None，调用者应优雅降级。
pub(crate) unsafe fn get_logon_sid_bytes(h_token: HANDLE) -> Option<Vec<u8>> {
    let mut needed: u32 = 0;
    GetTokenInformation(h_token, TokenGroups, ptr::null_mut(), 0, &mut needed);
    if needed == 0 {
        return None;
    }
    let mut buf: Vec<u8> = vec![0u8; needed as usize];
    let ok = GetTokenInformation(
        h_token,
        TokenGroups,
        buf.as_mut_ptr() as *mut c_void,
        needed,
        &mut needed,
    );
    if ok == 0 {
        return None;
    }

    let group_count = ptr::read_unaligned(buf.as_ptr() as *const u32) as usize;

    // TOKEN_GROUPS 的内存布局：GroupCount (u32) 后接 SID_AND_ATTRIBUTES[]
    // Windows 此处无对齐填充，使用 read_unaligned 安全读取每个条目
    let groups_start = buf.as_ptr().add(std::mem::size_of::<u32>());
    for i in 0..group_count {
        let offset = i * std::mem::size_of::<SID_AND_ATTRIBUTES>();
        let entry_ptr = groups_start.add(offset) as *const SID_AND_ATTRIBUTES;
        let entry: SID_AND_ATTRIBUTES = ptr::read_unaligned(entry_ptr);
        if (entry.Attributes & SE_GROUP_LOGON_ID) == SE_GROUP_LOGON_ID {
            let sid_len = GetLengthSid(entry.Sid);
            if sid_len == 0 {
                continue;
            }
            let mut out = vec![0u8; sid_len as usize];
            if CopySid(sid_len, out.as_mut_ptr() as *mut c_void, entry.Sid) == 0 {
                continue;
            }
            return Some(out);
        }
    }
    None // Logon SID not present on this token
}

/// 设置令牌的默认 DACL
unsafe fn set_default_dacl(h_token: HANDLE, sids: &[*mut c_void]) -> Result<()> {
    if sids.is_empty() {
        return Ok(());
    }
    let entries: Vec<EXPLICIT_ACCESS_W> = sids
        .iter()
        .map(|sid| EXPLICIT_ACCESS_W {
            grfAccessPermissions: GENERIC_ALL,
            grfAccessMode: GRANT_ACCESS,
            grfInheritance: 0,
            Trustee: TRUSTEE_W {
                pMultipleTrustee: ptr::null_mut(),
                MultipleTrusteeOperation: 0,
                TrusteeForm: TRUSTEE_IS_SID,
                TrusteeType: TRUSTEE_IS_UNKNOWN,
                ptstrName: *sid as *mut u16,
            },
        })
        .collect();

    let mut p_new_dacl: *mut windows_sys::Win32::Security::ACL = ptr::null_mut();
    let res = SetEntriesInAclW(
        entries.len() as u32,
        entries.as_ptr(),
        ptr::null_mut(),
        &mut p_new_dacl,
    );
    if res != 0 {
        return Err(anyhow!("SetEntriesInAclW failed: {res}"));
    }

    #[repr(C)]
    struct TokenDefaultDaclInfo {
        default_dacl: *mut windows_sys::Win32::Security::ACL,
    }

    let mut info = TokenDefaultDaclInfo {
        default_dacl: p_new_dacl,
    };
    let ok = SetTokenInformation(
        h_token,
        TokenDefaultDacl,
        &mut info as *mut _ as *mut c_void,
        std::mem::size_of::<TokenDefaultDaclInfo>() as u32,
    );
    if ok == 0 {
        let err = GetLastError();
        if !p_new_dacl.is_null() {
            LocalFree(p_new_dacl as HLOCAL);
        }
        return Err(anyhow!("SetTokenInformation(TokenDefaultDacl) failed: {err}"));
    }
    if !p_new_dacl.is_null() {
        LocalFree(p_new_dacl as HLOCAL);
    }
    Ok(())
}

/// 在令牌上启用一个特权
unsafe fn enable_privilege(h_token: HANDLE, name: &str) -> Result<()> {
    let mut luid = LUID { LowPart: 0, HighPart: 0 };
    let wide_name: Vec<u16> = crate::winutil::to_wide(name);
    let ok = LookupPrivilegeValueW(ptr::null(), wide_name.as_ptr(), &mut luid);
    if ok == 0 {
        return Err(anyhow!("LookupPrivilegeValueW failed: {}", GetLastError()));
    }
    let mut tp: TOKEN_PRIVILEGES = std::mem::zeroed();
    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    let ok2 = AdjustTokenPrivileges(h_token, 0, &tp, 0, ptr::null_mut(), ptr::null_mut());
    if ok2 == 0 {
        return Err(anyhow!("AdjustTokenPrivileges failed: {}", GetLastError()));
    }
    Ok(())
}

/// 创建受限令牌的底层实现
///
/// 抽取 `create_restricted_token` 和 `create_readonly_restricted_token` 的公共逻辑。
///
/// # Safety
///
/// 调用者必须确保以下约定：
/// - `cap_sids` 中的每个指针必须指向有效的、已初始化的 SID 结构体，
///   且指针在函数执行期间（直到 `CreateRestrictedToken` 返回前）必须保持有效。
/// - `blacklist_sid` 如果为 `Some(ptr)`，`ptr` 必须指向有效的 SID 结构体，
///   且指针在函数执行期间必须保持有效。
/// - 所有 SID 指针必须是通过 `ConvertStringSidToSidW`、`AllocateAndInitializeSid`
///   或类似 Windows API 合法分配的，并且其内存布局符合 Windows SID 规范。
/// - 如果指针在函数返回后不再使用（不泄漏），调用者无需额外清理——本函数仅
///   在 `CreateRestrictedToken` 中引用这些 SID，不获取所有权。
///
/// # 参数
/// - `cap_sids`: Capability SID 指针列表，空切片表示只读令牌
/// - `blacklist_sid`: 可选的全局黑名单 SID
///
/// # 返回
/// 受限令牌句柄（调用者需要 CloseHandle）
unsafe fn create_token_inner(
    cap_sids: &[*mut c_void],
    blacklist_sid: Option<*mut c_void>,
) -> Result<HANDLE> {
    let base = get_current_token()?;

    let mut logon_sid_bytes = get_logon_sid_bytes(base);
    let has_logon = logon_sid_bytes.is_some();
    let psid_logon = logon_sid_bytes.as_mut().map_or(std::ptr::null_mut(), |b| b.as_mut_ptr() as *mut c_void);

    let mut everyone = world_sid()?;
    let psid_everyone = everyone.as_mut_ptr() as *mut c_void;

    // restricting SID 列表 = cap_sids + [logon（可选）, everyone] + [blacklist]（可选）
    let logon_extra = if has_logon { 1 } else { 0 };
    let bl_extra = if blacklist_sid.is_some() { 1 } else { 0 };
    let total_entries = cap_sids.len() + 1 + logon_extra + bl_extra;
    let mut entries: Vec<SID_AND_ATTRIBUTES> = vec![std::mem::zeroed(); total_entries];

    let mut idx = 0;
    for &psid in cap_sids {
        entries[idx].Sid = psid;
        idx += 1;
    }
    if has_logon {
        entries[idx].Sid = psid_logon;
        idx += 1;
    }
    entries[idx].Sid = psid_everyone;
    idx += 1;
    if let Some(bl_sid) = blacklist_sid {
        entries[idx].Sid = bl_sid;
    }

    let mut new_token: HANDLE = 0;
    let flags = DISABLE_MAX_PRIVILEGE | LUA_TOKEN | WRITE_RESTRICTED;
    let ok = CreateRestrictedToken(
        base,
        flags,
        0, ptr::null(),
        0, ptr::null(),
        entries.len() as u32,
        entries.as_ptr(),
        &mut new_token,
    );
    if ok == 0 {
        let err = GetLastError();
        CloseHandle(base);
        return Err(anyhow!("CreateRestrictedToken failed: {err}"));
    }

    // 默认 DACL：Everyone + Logon SID（若有）+ capability SIDs
    let mut dacl_sids: Vec<*mut c_void> = Vec::with_capacity(cap_sids.len() + 1 + logon_extra);
    if has_logon {
        dacl_sids.push(psid_logon);
    }
    dacl_sids.push(psid_everyone);
    dacl_sids.extend_from_slice(cap_sids);

    if let Err(e) = set_default_dacl(new_token, &dacl_sids) {
        CloseHandle(new_token);
        CloseHandle(base);
        return Err(e);
    }

    if let Err(e) = enable_privilege(new_token, "SeChangeNotifyPrivilege") {
        CloseHandle(new_token);
        CloseHandle(base);
        return Err(e);
    }

    CloseHandle(base);
    Ok(new_token)
}

/// 创建受限令牌（含写入 capability）
///
/// # Safety
///
/// 调用者必须确保：
/// - `cap_sids` 中的每个指针必须指向有效的 SID 结构体，
///   且指针在 `create_token_inner` 执行期间保持有效。
/// - `blacklist_sid` 如果为 `Some(ptr)`，`ptr` 必须指向有效的 SID 结构体。
/// - 调用者负责在不再需要令牌时通过 `CloseHandle` 关闭返回的令牌句柄。
///
/// # 参数
/// - `cap_sids`: Capability SID 指针列表（用于 Allow ACE）
/// - `blacklist_sid`: 可选的全局黑名单 SID（用于 Deny ALL ACE，不加入默认 DACL）
///
/// # 返回
/// 受限令牌句柄（调用者需要 CloseHandle）
pub unsafe fn create_restricted_token(
    cap_sids: &[*mut c_void],
    blacklist_sid: Option<*mut c_void>,
) -> Result<HANDLE> {
    create_token_inner(cap_sids, blacklist_sid)
}

/// 创建只读受限令牌（没有写入 capability）
///
/// # Safety
///
/// 本函数会创建新的受限令牌，调用者负责在不再需要时通过 `CloseHandle` 关闭返回的令牌句柄。
pub unsafe fn create_readonly_restricted_token() -> Result<HANDLE> {
    create_token_inner(&[], None)
}