rustpython-host_env 0.6.0

use rustpython_wtf8::Wtf8;
use std::{ffi::OsStr, io};
use widestring::{WideCStr, WideCString};
use windows_sys::{
    Win32::{
        Foundation::{
            E_POINTER, ERROR_INSUFFICIENT_BUFFER, ERROR_INVALID_FLAGS,
            ERROR_NO_UNICODE_TRANSLATION, MAX_PATH, S_OK,
        },
        Networking::WinSock::WSAStartup,
        Storage::FileSystem::{
            GetFileVersionInfoSizeW, GetFileVersionInfoW, VS_FIXEDFILEINFO, VerQueryValueW,
        },
        System::{
            Diagnostics::Debug::{
                FORMAT_MESSAGE_ALLOCATE_BUFFER, FORMAT_MESSAGE_FROM_SYSTEM,
                FORMAT_MESSAGE_IGNORE_INSERTS, FormatMessageW,
            },
            LibraryLoader::{GetModuleFileNameW, GetModuleHandleW},
            SystemInformation::{GetVersionExW, OSVERSIONINFOEXW, OSVERSIONINFOW},
            Threading::{GetCurrentThreadStackLimits, SetThreadStackGuarantee},
        },
    },
    w,
};

const NUL_ERROR: &str = "embedded null character";

/// _MAX_ENV from Windows CRT stdlib.h - maximum environment variable size
pub const _MAX_ENV: usize = 32767;
pub const HRESULT_E_POINTER: i32 = E_POINTER;
pub const HRESULT_S_OK: i32 = S_OK;
pub const CP_ACP: u32 = windows_sys::Win32::Globalization::CP_ACP;
pub const CP_OEMCP: u32 = windows_sys::Win32::Globalization::CP_OEMCP;
pub const CP_UTF7: u32 = windows_sys::Win32::Globalization::CP_UTF7;
pub const CP_UTF8: u32 = windows_sys::Win32::Globalization::CP_UTF8;
pub const MB_ERR_INVALID_CHARS: u32 = windows_sys::Win32::Globalization::MB_ERR_INVALID_CHARS;
pub const WC_ERR_INVALID_CHARS: u32 = windows_sys::Win32::Globalization::WC_ERR_INVALID_CHARS;
pub const WC_NO_BEST_FIT_CHARS: u32 = windows_sys::Win32::Globalization::WC_NO_BEST_FIT_CHARS;
pub const ERROR_INVALID_FLAGS_I32: i32 = ERROR_INVALID_FLAGS as i32;
pub const ERROR_NO_UNICODE_TRANSLATION_I32: i32 = ERROR_NO_UNICODE_TRANSLATION as i32;
pub const ERROR_INSUFFICIENT_BUFFER_I32: i32 = ERROR_INSUFFICIENT_BUFFER as i32;

pub fn init_winsock() {
    static WSA_INIT: parking_lot::Once = parking_lot::Once::new();
    WSA_INIT.call_once(|| unsafe {
        let mut wsa_data = core::mem::MaybeUninit::uninit();
        let _ = WSAStartup(0x0101, wsa_data.as_mut_ptr());
    })
}

/// Win32 BOOL convention: 0 = failure, nonzero = success.
/// Reads error via `GetLastError()` through [`io::Error::last_os_error`].
pub trait CheckWin32Bool {
    fn check_win32_bool(self) -> io::Result<()>;
}

impl CheckWin32Bool for i32 {
    #[inline]
    fn check_win32_bool(self) -> io::Result<()> {
        if self == 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }
}

/// Convenience checks for Win32 `HANDLE` return values.
pub trait CheckWin32Handle: Sized {
    /// Returns `Ok(self)` if the handle is non-NULL, otherwise the last OS error.
    /// Use for APIs whose failure sentinel is NULL (Create*Event, Create*Mutex,
    /// CreateFileMapping, GetModuleHandle, etc.).
    fn check_nonnull(self) -> io::Result<Self>;

    /// Returns `Ok(self)` unless the handle equals `INVALID_HANDLE_VALUE`.
    /// Use for APIs whose failure sentinel is `INVALID_HANDLE_VALUE`
    /// (CreateFileW, CreateNamedPipeW, etc.).
    fn check_valid(self) -> io::Result<Self>;
}

impl CheckWin32Handle for windows_sys::Win32::Foundation::HANDLE {
    #[inline]
    fn check_nonnull(self) -> io::Result<Self> {
        if self.is_null() {
            Err(io::Error::last_os_error())
        } else {
            Ok(self)
        }
    }

    #[inline]
    fn check_valid(self) -> io::Result<Self> {
        if self == windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE {
            Err(io::Error::last_os_error())
        } else {
            Ok(self)
        }
    }
}

/// Generic sentinel check for Win32 return values not covered by [`CheckWin32Bool`] or
/// [`CheckWin32Handle`]. Use for APIs that signal failure with a specific value
/// (`INVALID_FILE_ATTRIBUTES`, `WAIT_FAILED`, `GetFileVersionInfoSizeW` returning `0`, etc.).
pub trait CheckWin32Sentinel: Sized + Copy + PartialEq {
    /// Returns `Ok(self)` if `self != fail`, otherwise the last OS error.
    #[inline]
    fn check_ne(self, fail: Self) -> io::Result<Self> {
        if self == fail {
            Err(io::Error::last_os_error())
        } else {
            Ok(self)
        }
    }
}

impl CheckWin32Sentinel for u32 {}
impl CheckWin32Sentinel for i32 {}
impl CheckWin32Sentinel for u64 {}

use std::os::windows::io::FromRawHandle;
pub use std::os::windows::io::{BorrowedHandle, OwnedHandle};

/// Conversion of raw Win32 `HANDLE` into an [`OwnedHandle`] (RAII auto-close).
pub trait HandleToOwned: Sized {
    /// Wraps the handle in an [`OwnedHandle`] that calls `CloseHandle` on drop.
    /// Returns `None` if the handle is NULL or `INVALID_HANDLE_VALUE`.
    fn into_owned(self) -> Option<OwnedHandle>;
}

impl HandleToOwned for windows_sys::Win32::Foundation::HANDLE {
    #[inline]
    fn into_owned(self) -> Option<OwnedHandle> {
        if self.is_null() || self == windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE {
            None
        } else {
            // SAFETY: caller is asserting via the public `Create*`/`Open*` paths
            // that this handle is owned and unaliased.
            Some(unsafe { OwnedHandle::from_raw_handle(self.cast()) })
        }
    }
}

#[derive(Clone, Debug)]
pub struct WindowsVersionInfo {
    pub major: u32,
    pub minor: u32,
    pub build: u32,
    pub platform: u32,
    pub service_pack: String,
    pub service_pack_major: u16,
    pub service_pack_minor: u16,
    pub suite_mask: u16,
    pub product_type: u8,
}

fn get_kernel32_version() -> io::Result<(u32, u32, u32)> {
    unsafe {
        let module_name = w!("kernel32.dll");
        let h_kernel32 = GetModuleHandleW(module_name).check_nonnull()?;

        let mut kernel32_path = [0u16; MAX_PATH as usize];
        let len = GetModuleFileNameW(
            h_kernel32,
            kernel32_path.as_mut_ptr(),
            kernel32_path.len() as u32,
        );
        if len == 0 {
            return Err(io::Error::last_os_error());
        }

        let ver_block_size = GetFileVersionInfoSizeW(kernel32_path.as_ptr(), core::ptr::null_mut());
        if ver_block_size == 0 {
            return Err(io::Error::last_os_error());
        }

        let mut ver_block = vec![0u8; ver_block_size as usize];
        GetFileVersionInfoW(
            kernel32_path.as_ptr(),
            0,
            ver_block_size,
            ver_block.as_mut_ptr() as *mut _,
        )
        .check_win32_bool()?;

        let sub_block = w!("");

        let mut ffi_ptr: *mut VS_FIXEDFILEINFO = core::ptr::null_mut();
        let mut ffi_len: u32 = 0;
        VerQueryValueW(
            ver_block.as_ptr() as *const _,
            sub_block,
            &mut ffi_ptr as *mut *mut VS_FIXEDFILEINFO as *mut *mut _,
            &mut ffi_len as *mut u32,
        )
        .check_win32_bool()?;
        if ffi_ptr.is_null() {
            return Err(io::Error::last_os_error());
        }

        let ffi = *ffi_ptr;
        let real_major = (ffi.dwProductVersionMS >> 16) & 0xFFFF;
        let real_minor = ffi.dwProductVersionMS & 0xFFFF;
        let real_build = (ffi.dwProductVersionLS >> 16) & 0xFFFF;

        Ok((real_major, real_minor, real_build))
    }
}

pub fn get_windows_version() -> io::Result<WindowsVersionInfo> {
    let mut version: OSVERSIONINFOEXW = unsafe { core::mem::zeroed() };
    version.dwOSVersionInfoSize = core::mem::size_of::<OSVERSIONINFOEXW>() as u32;
    unsafe {
        let os_vi = &mut version as *mut OSVERSIONINFOEXW as *mut OSVERSIONINFOW;
        GetVersionExW(os_vi)
    }
    .check_win32_bool()?;

    // szCSDVersion is at most 128 wide chars and represented as slice so overflow is impossible.
    let service_pack = WideCStr::from_slice_truncate(&version.szCSDVersion)
        .map_err(|_| io::Error::other(NUL_ERROR))?
        .to_string()
        .map_err(|_| io::Error::other("service pack is not ASCII"))?;

    let (major, minor, build) = get_kernel32_version()?;
    Ok(WindowsVersionInfo {
        major,
        minor,
        build,
        platform: version.dwPlatformId,
        service_pack,
        service_pack_major: version.wServicePackMajor,
        service_pack_minor: version.wServicePackMinor,
        suite_mask: version.wSuiteMask,
        product_type: version.wProductType,
    })
}

pub fn current_thread_stack_bounds() -> (usize, usize) {
    let mut low: usize = 0;
    let mut high: usize = 0;
    unsafe {
        GetCurrentThreadStackLimits(&mut low as *mut usize, &mut high as *mut usize);
        let mut guarantee: u32 = 0;
        SetThreadStackGuarantee(&mut guarantee);
        low += guarantee as usize;
    }
    (low, high)
}

pub fn set_last_error(error: u32) {
    unsafe { windows_sys::Win32::Foundation::SetLastError(error) }
}

pub fn get_last_error() -> u32 {
    unsafe { windows_sys::Win32::Foundation::GetLastError() }
}

pub fn format_error_message(code: Option<u32>) -> Option<String> {
    let error_code = code.unwrap_or_else(get_last_error);
    let mut buffer: *mut u16 = core::ptr::null_mut();
    let len = unsafe {
        FormatMessageW(
            FORMAT_MESSAGE_ALLOCATE_BUFFER
                | FORMAT_MESSAGE_FROM_SYSTEM
                | FORMAT_MESSAGE_IGNORE_INSERTS,
            core::ptr::null(),
            error_code,
            0,
            &mut buffer as *mut *mut u16 as *mut u16,
            0,
            core::ptr::null(),
        )
    };

    if len == 0 || buffer.is_null() {
        return None;
    }

    let message = unsafe {
        let slice = core::slice::from_raw_parts(buffer, len as usize);
        let msg = String::from_utf16_lossy(slice).trim_end().to_string();
        windows_sys::Win32::Foundation::LocalFree(buffer as *mut _);
        msg
    };
    Some(message)
}

pub fn wide_char_to_multi_byte_len(
    code_page: u32,
    flags: u32,
    wide: &[u16],
    track_default_char: bool,
) -> io::Result<(usize, bool)> {
    let mut used_default_char = 0i32;
    let pused = if track_default_char {
        &mut used_default_char as *mut i32
    } else {
        core::ptr::null_mut()
    };
    let size = unsafe {
        windows_sys::Win32::Globalization::WideCharToMultiByte(
            code_page,
            flags,
            wide.as_ptr(),
            wide.len() as i32,
            core::ptr::null_mut(),
            0,
            core::ptr::null(),
            pused,
        )
    };
    if size <= 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok((size as usize, used_default_char != 0))
    }
}

pub fn wide_char_to_multi_byte(
    code_page: u32,
    flags: u32,
    wide: &[u16],
    out: &mut [u8],
    track_default_char: bool,
) -> io::Result<(usize, bool)> {
    let mut used_default_char = 0i32;
    let pused = if track_default_char {
        &mut used_default_char as *mut i32
    } else {
        core::ptr::null_mut()
    };
    let size = unsafe {
        windows_sys::Win32::Globalization::WideCharToMultiByte(
            code_page,
            flags,
            wide.as_ptr(),
            wide.len() as i32,
            out.as_mut_ptr().cast(),
            out.len() as i32,
            core::ptr::null(),
            pused,
        )
    };
    if size <= 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok((size as usize, used_default_char != 0))
    }
}

pub fn multi_byte_to_wide_len(code_page: u32, flags: u32, bytes: &[u8]) -> io::Result<usize> {
    let size = unsafe {
        windows_sys::Win32::Globalization::MultiByteToWideChar(
            code_page,
            flags,
            bytes.as_ptr().cast(),
            bytes.len() as i32,
            core::ptr::null_mut(),
            0,
        )
    };
    if size <= 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(size as usize)
    }
}

pub fn multi_byte_to_wide(
    code_page: u32,
    flags: u32,
    bytes: &[u8],
    out: &mut [u16],
) -> io::Result<usize> {
    let size = unsafe {
        windows_sys::Win32::Globalization::MultiByteToWideChar(
            code_page,
            flags,
            bytes.as_ptr().cast(),
            bytes.len() as i32,
            out.as_mut_ptr(),
            out.len() as i32,
        )
    };
    if size <= 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(size as usize)
    }
}

/// [`OsStr`] to [`WideCString`] for Windows FFI.
///
/// Prefer using this trait when encoding bytes to pass to Windows. Interior NULs are memory safe
/// but possibly a security hazard for FFI.
///
/// https://github.com/python/cpython/issues/111656
pub trait ToWideString {
    fn to_wide_cstring(&self) -> Result<WideCString, io::Error>;
}

impl<T> ToWideString for T
where
    T: AsRef<OsStr>,
{
    fn to_wide_cstring(&self) -> Result<WideCString, io::Error> {
        WideCString::from_os_str(self).map_err(|_| io::Error::other(NUL_ERROR))
    }
}

impl ToWideString for Wtf8 {
    fn to_wide_cstring(&self) -> Result<WideCString, io::Error> {
        // CPython's "test_invalid_cmd" test calls Popen with "pass#\0" as a command line.
        // That's technically valid since it caps the string, but CString and WideCString differ
        // in how they handle it. Rust's CString rejects any NULs whereas WideCString accepts a NUL
        // only if it appears at the end of a buffer.
        //
        // For the sake of that behavior, fail on trailing NUL.
        if self.as_bytes().last().is_some_and(|&b| b == 0) {
            return Err(io::Error::other(NUL_ERROR));
        }

        let mut buf = Vec::with_capacity(self.len() + 1);
        buf.extend(self.encode_wide());
        WideCString::from_vec(buf).map_err(|_| io::Error::other(NUL_ERROR))
    }
}