Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 32 additions & 21 deletions crates/vm/src/stdlib/_winapi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ mod _winapi {
Py, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine,
builtins::PyStrRef,
common::lock::{PyMutex, PyMutexGuard},
convert::ToPyException,
convert::{ToPyException, ToPyObject},
exceptions::nul_char_error,
function::{ArgMapping, ArgSequence, OptionalArg},
types::Constructor,
Expand Down Expand Up @@ -72,8 +72,10 @@ mod _winapi {
const INFINITE: u32 = host_winapi::INFINITE_TIMEOUT;

#[pyfunction]
fn CloseHandle(handle: WinHandle) -> WindowsSysResult<i32> {
WindowsSysResult(host_winapi::close_handle(handle.0))
fn CloseHandle(handle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::from_bitfield(&handle, vm, "CloseHandle", "argument")?;
WindowsSysResult(host_winapi::close_handle(value.0)).into_pyresult(vm)?;
Ok(vm.ctx.none())
}

#[pyfunction]
Expand All @@ -96,9 +98,11 @@ mod _winapi {
}

#[pyfunction]
fn DeregisterEventSource(handle: WinHandle, vm: &VirtualMachine) -> PyResult<()> {
vm.allow_threads(|| host_winapi::deregister_event_source(handle.0))
.map_err(|error| error.to_pyexception(vm))
fn DeregisterEventSource(handle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::from_bitfield(&handle, vm, "DeregisterEventSource", "argument")?;
vm.allow_threads(|| host_winapi::deregister_event_source(value.0))
.map_err(|error| error.to_pyexception(vm))?;
Ok(vm.ctx.none())
}

#[pyfunction]
Expand Down Expand Up @@ -152,13 +156,12 @@ mod _winapi {
}

#[pyfunction]
fn GetStdHandle(
std_handle: host_winapi::StdHandle,
vm: &VirtualMachine,
) -> PyResult<Option<WinHandle>> {
host_winapi::get_std_handle(std_handle)
fn GetStdHandle(std_handle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::unsigned_long_bitfield(&std_handle, vm, "GetStdHandle", "argument")?;
let handle = host_winapi::get_std_handle(value)
.map(|handle| handle.map(WinHandle))
.map_err(|e| e.to_pyexception(vm))
.map_err(|e| e.to_pyexception(vm))?;
Ok(handle.to_pyobject(vm))
}

#[pyfunction]
Expand Down Expand Up @@ -444,8 +447,10 @@ mod _winapi {
}

#[pyfunction]
fn GetExitCodeProcess(h: WinHandle, vm: &VirtualMachine) -> PyResult<u32> {
host_winapi::get_exit_code_process(h.0).map_err(|e| e.to_pyexception(vm))
fn GetExitCodeProcess(h: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::from_bitfield(&h, vm, "GetExitCodeProcess", "argument")?;
let code = host_winapi::get_exit_code_process(value.0).map_err(|e| e.to_pyexception(vm))?;
Ok(vm.ctx.new_int(code).into())
}

#[pyfunction]
Expand Down Expand Up @@ -496,16 +501,17 @@ mod _winapi {
}

#[pyfunction]
fn GetModuleFileName(handle: isize, vm: &VirtualMachine) -> PyResult<String> {
fn GetModuleFileName(handle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::from_bitfield(&handle, vm, "GetModuleFileName", "argument")?;
let mut path: Vec<u16> = vec![0; host_winapi::MAX_PATH_USIZE];

let length = host_winapi::get_module_file_name(handle as _, &mut path);
let length = host_winapi::get_module_file_name(value.0, &mut path);
if length == 0 {
return Err(vm.new_runtime_error("GetModuleFileName failed"));
}

let (path, _) = path.split_at(length as usize);
Ok(String::from_utf16(path).unwrap())
Ok(vm.ctx.new_str(String::from_utf16(path).unwrap()).into())
}

#[pyfunction]
Expand Down Expand Up @@ -1118,14 +1124,19 @@ mod _winapi {

// UnmapViewOfFile - Unmap a mapped view of a file.
#[pyfunction]
fn UnmapViewOfFile(address: isize, vm: &VirtualMachine) -> PyResult<()> {
host_winapi::unmap_view_of_file(address).map_err(|e| e.to_pyexception(vm))
fn UnmapViewOfFile(address: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::from_bitfield(&address, vm, "UnmapViewOfFile", "argument")?;
host_winapi::unmap_view_of_file(value.0 as isize).map_err(|e| e.to_pyexception(vm))?;
Ok(vm.ctx.none())
}

// VirtualQuerySize - Return the size of a memory region.
#[pyfunction]
fn VirtualQuerySize(address: isize, vm: &VirtualMachine) -> PyResult<usize> {
host_winapi::virtual_query_size(address).map_err(|e| e.to_pyexception(vm))
fn VirtualQuerySize(address: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let value = WinHandle::from_bitfield(&address, vm, "VirtualQuerySize", "argument")?;
let size =
host_winapi::virtual_query_size(value.0 as isize).map_err(|e| e.to_pyexception(vm))?;
Ok(vm.ctx.new_int(size).into())
}

#[pyfunction]
Expand Down
93 changes: 92 additions & 1 deletion crates/vm/src/windows.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,104 @@
use crate::{
PyObjectRef, PyResult, TryFromObject, VirtualMachine,
PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
builtins::PyInt,
convert::{ToPyObject, ToPyResult},
};
use num_traits::ToPrimitive;
use rustpython_host_env::nt as host_nt;

/// Windows HANDLE wrapper for Python interop
#[derive(Clone, Copy)]
pub struct WinHandle(pub host_nt::Handle);

impl WinHandle {
/// Convert an F_HANDLE bitfield after the caller has finished binding.
/// The caller must keep the original argument alive through its operation.
pub(crate) fn from_bitfield(
obj: &PyObject,
vm: &VirtualMachine,
callee: &str,
argument: &str,
) -> PyResult<Self> {
let value = Self::parse_bitfield::<{ usize::BITS }>(obj, vm, callee, argument)?;
Ok(Self(value as host_nt::Handle))
}

/// Convert a Windows `unsigned long` (`k`) after argument binding.
pub(crate) fn unsigned_long_bitfield(
obj: &PyObject,
vm: &VirtualMachine,
callee: &str,
argument: &str,
) -> PyResult<u32> {
Ok(Self::parse_bitfield::<32>(obj, vm, callee, argument)? as u32)
}

fn parse_bitfield<const WIDTH: u32>(
obj: &PyObject,
vm: &VirtualMachine,
callee: &str,
argument: &str,
) -> PyResult<usize> {
let (handle, fits) = {
let owned;
let int = if let Some(int) = obj.downcast_ref::<PyInt>() {
int
} else if let Some(index) = obj.try_index_opt(vm) {
owned = index?;
&owned
} else {
let actual = if vm.is_none(obj) {
b"None".to_vec()
} else {
let name = obj.class().slot_name();
name.as_bytes()[..name.len().min(50)].to_vec()
};
let mut message = format!("{callee}() {argument} must be int, not ").into_bytes();
message.extend_from_slice(&actual);
return Err(match String::from_utf8(message) {
Ok(message) => vm.new_type_error(message),
Err(error) => {
let start = error.utf8_error().valid_up_to();
let bytes = error.into_bytes();
let end = bytes.len();
vm.new_unicode_decode_error(
vm.ctx.new_str("utf-8"),
vm.ctx.new_bytes(bytes),
start,
end,
vm.ctx.new_str("unexpected end of data"),
)
}
});
};
if WIDTH == 32 {
if let Some(value) = int.as_bigint().to_i32() {
(value as u32 as usize, true)
} else if let Some(value) = int.as_bigint().to_u32() {
(value as usize, true)
} else {
(int.as_u32_mask() as usize, false)
}
} else if let Some(value) = int.as_bigint().to_isize() {
(value as usize, true)
} else if let Some(value) = int.as_bigint().to_usize() {
(value, true)
} else {
(int.as_u64_mask() as usize, false)
}
};
if !fits {
crate::stdlib::_warnings::warn(
vm.ctx.exceptions.deprecation_warning,
"integer value out of range".to_owned(),
1,
vm,
)?;
}
Ok(handle)
}
}

pub(crate) trait WindowsSysResultValue {
type Ok: ToPyObject;

Expand Down
Loading