diff --git a/Lib/test/test_winreg.py b/Lib/test/test_winreg.py index 2c530ee75..924a96278 100644 --- a/Lib/test/test_winreg.py +++ b/Lib/test/test_winreg.py @@ -212,14 +212,10 @@ class BaseWinregTests(unittest.TestCase): class LocalWinregTests(BaseWinregTests): - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_registry_works(self): self._test_all(HKEY_CURRENT_USER) self._test_all(HKEY_CURRENT_USER, "日本-subkey") - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_registry_works_extended_functions(self): # Substitute the regular CreateKey and OpenKey calls with their # extended counterparts. @@ -232,8 +228,6 @@ class LocalWinregTests(BaseWinregTests): self._delete_test_data(HKEY_CURRENT_USER) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_named_arguments(self): self._test_named_args(HKEY_CURRENT_USER, test_key_name) # Use the regular DeleteKey to clean up @@ -251,8 +245,6 @@ class LocalWinregTests(BaseWinregTests): connect = lambda: ConnectRegistry("abcdefghijkl", HKEY_CURRENT_USER) self.assertRaises(OSError, connect) - # TODO: RUSTPYTHON - @unittest.skip("flaky") def testExpandEnvironmentStrings(self): r = ExpandEnvironmentStrings("%windir%\\test") self.assertEqual(type(r), str) @@ -299,8 +291,6 @@ class LocalWinregTests(BaseWinregTests): DeleteKey(HKEY_CURRENT_USER, test_key_name+'\\changing_value') DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_long_key(self): # Issue2810, in 2.6 and 3.1 when the key name was exactly 256 # characters, EnumKey raised "WindowsError: More data is @@ -315,8 +305,6 @@ class LocalWinregTests(BaseWinregTests): DeleteKey(HKEY_CURRENT_USER, '\\'.join((test_key_name, name))) DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_dynamic_key(self): # Issue2810, when the value is dynamically generated, these # raise "WindowsError: More data is available" in 2.6 and 3.1 @@ -351,8 +339,6 @@ class LocalWinregTests(BaseWinregTests): finally: DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_setvalueex_value_range(self): # Test for Issue #14420, accept proper ranges for SetValueEx. # Py2Reg, which gets called by SetValueEx, was using PyLong_AsLong, @@ -365,8 +351,6 @@ class LocalWinregTests(BaseWinregTests): finally: DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_setvalueex_negative_one_check(self): # Test for Issue #43984, check -1 was not set by SetValueEx. # Py2Reg, which gets called by SetValueEx, wasn't checking return @@ -384,8 +368,6 @@ class LocalWinregTests(BaseWinregTests): finally: DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_queryvalueex_return_value(self): # Test for Issue #16759, return unsigned int from QueryValueEx. # Reg2Py, which gets called by QueryValueEx, was returning a value @@ -402,8 +384,6 @@ class LocalWinregTests(BaseWinregTests): finally: DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_setvalueex_crash_with_none_arg(self): # Test for Issue #21151, segfault when None is passed to SetValueEx try: @@ -417,8 +397,6 @@ class LocalWinregTests(BaseWinregTests): finally: DeleteKey(HKEY_CURRENT_USER, test_key_name) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_read_string_containing_null(self): # Test for issue 25778: REG_SZ should not contain null characters try: @@ -443,16 +421,13 @@ class RemoteWinregTests(BaseWinregTests): @unittest.skipUnless(WIN64_MACHINE, "x64 specific registry tests") class Win64WinregTests(BaseWinregTests): - # TODO: RUSTPYTHON - @unittest.expectedFailure + def test_named_arguments(self): self._test_named_args(HKEY_CURRENT_USER, test_key_name) # Clean up and also exercise the named arguments DeleteKeyEx(key=HKEY_CURRENT_USER, sub_key=test_key_name, access=KEY_ALL_ACCESS, reserved=0) - # TODO: RUSTPYTHON - @unittest.expectedFailure @unittest.skipIf(win32_edition() in ('WindowsCoreHeadless', 'IoTEdgeOS'), "APIs not available on WindowsCoreHeadless") def test_reflection_functions(self): # Test that we can call the query, enable, and disable functions @@ -538,8 +513,6 @@ class Win64WinregTests(BaseWinregTests): DeleteKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, KEY_WOW64_32KEY, 0) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_exception_numbers(self): with self.assertRaises(FileNotFoundError) as ctx: QueryValue(HKEY_CLASSES_ROOT, 'some_value_that_does_not_exist') diff --git a/crates/vm/src/stdlib/winreg.rs b/crates/vm/src/stdlib/winreg.rs index f3744d93f..396f60a2b 100644 --- a/crates/vm/src/stdlib/winreg.rs +++ b/crates/vm/src/stdlib/winreg.rs @@ -9,25 +9,67 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef { #[pymodule] mod winreg { - use std::ffi::OsStr; - use std::os::windows::ffi::OsStrExt; - use std::ptr; - use std::sync::Arc; - - use crate::builtins::{PyInt, PyTuple}; - use crate::common::lock::PyRwLock; + use crate::builtins::{PyInt, PyTuple, PyTypeRef}; + use crate::common::hash::PyHash; + use crate::common::windows::ToWideString; + use crate::convert::TryFromObject; use crate::function::FuncArgs; + use crate::object::AsObject; use crate::protocol::PyNumberMethods; - use crate::types::AsNumber; - use crate::{PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine}; - + use crate::types::{AsNumber, Hashable}; + use crate::{Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine}; + use crossbeam_utils::atomic::AtomicCell; + use malachite_bigint::Sign; + use num_traits::ToPrimitive; + use std::ptr; use windows_sys::Win32::Foundation::{self, ERROR_MORE_DATA}; use windows_sys::Win32::System::Registry; - use num_traits::ToPrimitive; + /// Atomic HKEY handle type for lock-free thread-safe access + type AtomicHKEY = AtomicCell; - pub(crate) fn to_utf16>(s: P) -> Vec { - s.as_ref().encode_wide().chain(Some(0)).collect() + /// Convert byte slice to UTF-16 slice (zero-copy when aligned) + fn bytes_as_wide_slice(bytes: &[u8]) -> &[u16] { + // SAFETY: Windows Registry API returns properly aligned UTF-16 data. + // align_to handles any edge cases safely by returning empty prefix/suffix + // if alignment doesn't match. + let (prefix, u16_slice, suffix) = unsafe { bytes.align_to::() }; + debug_assert!( + prefix.is_empty() && suffix.is_empty(), + "Registry data should be u16-aligned" + ); + u16_slice + } + + // TODO: check if errno.rs can be reused here or not + fn os_error_from_windows_code( + vm: &VirtualMachine, + code: i32, + func_name: &str, + ) -> crate::PyRef { + use windows_sys::Win32::Foundation::{ERROR_ACCESS_DENIED, ERROR_FILE_NOT_FOUND}; + let msg = format!("[WinError {}] {}", code, func_name); + let exc_type = match code as u32 { + ERROR_FILE_NOT_FOUND => vm.ctx.exceptions.file_not_found_error.to_owned(), + ERROR_ACCESS_DENIED => vm.ctx.exceptions.permission_error.to_owned(), + _ => vm.ctx.exceptions.os_error.to_owned(), + }; + vm.new_exception_msg(exc_type, msg) + } + + /// Wrapper type for HKEY that can be created from PyHkey or int + struct HKEYArg(Registry::HKEY); + + impl TryFromObject for HKEYArg { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + // Try PyHkey first + if let Some(hkey_obj) = obj.downcast_ref::() { + return Ok(HKEYArg(hkey_obj.hkey.load())); + } + // Then try int + let handle = usize::try_from_object(vm, obj)?; + Ok(HKEYArg(handle as Registry::HKEY)) + } } // access rights @@ -49,123 +91,131 @@ mod winreg { REG_RESOURCE_REQUIREMENTS_LIST, REG_SZ, REG_WHOLE_HIVE_VOLATILE, }; - #[pyattr(once)] - fn HKEY_CLASSES_ROOT(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_CLASSES_ROOT)), - } + // Additional constants not in windows-sys + #[pyattr] + const REG_REFRESH_HIVE: u32 = 0x00000002; + #[pyattr] + const REG_NO_LAZY_FLUSH: u32 = 0x00000004; + // REG_LEGAL_OPTION is a mask of all option flags + #[pyattr] + const REG_LEGAL_OPTION: u32 = Registry::REG_OPTION_RESERVED + | Registry::REG_OPTION_NON_VOLATILE + | Registry::REG_OPTION_VOLATILE + | Registry::REG_OPTION_CREATE_LINK + | Registry::REG_OPTION_BACKUP_RESTORE + | Registry::REG_OPTION_OPEN_LINK; + // REG_LEGAL_CHANGE_FILTER is a mask of all notify flags + #[pyattr] + const REG_LEGAL_CHANGE_FILTER: u32 = Registry::REG_NOTIFY_CHANGE_NAME + | Registry::REG_NOTIFY_CHANGE_ATTRIBUTES + | Registry::REG_NOTIFY_CHANGE_LAST_SET + | Registry::REG_NOTIFY_CHANGE_SECURITY; + + // error is an alias for OSError (for backwards compatibility) + #[pyattr] + fn error(vm: &VirtualMachine) -> PyTypeRef { + vm.ctx.exceptions.os_error.to_owned() } #[pyattr(once)] - fn HKEY_CURRENT_USER(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_CURRENT_USER)), - } + fn HKEY_CLASSES_ROOT(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_CLASSES_ROOT).into_ref(&vm.ctx) } #[pyattr(once)] - fn HKEY_LOCAL_MACHINE(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_LOCAL_MACHINE)), - } + fn HKEY_CURRENT_USER(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_CURRENT_USER).into_ref(&vm.ctx) } #[pyattr(once)] - fn HKEY_USERS(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_USERS)), - } + fn HKEY_LOCAL_MACHINE(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_LOCAL_MACHINE).into_ref(&vm.ctx) } #[pyattr(once)] - fn HKEY_PERFORMANCE_DATA(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_PERFORMANCE_DATA)), - } + fn HKEY_USERS(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_USERS).into_ref(&vm.ctx) } #[pyattr(once)] - fn HKEY_CURRENT_CONFIG(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_CURRENT_CONFIG)), - } + fn HKEY_PERFORMANCE_DATA(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_PERFORMANCE_DATA).into_ref(&vm.ctx) } #[pyattr(once)] - fn HKEY_DYN_DATA(_vm: &VirtualMachine) -> PyHKEYObject { - PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(Registry::HKEY_DYN_DATA)), - } + fn HKEY_CURRENT_CONFIG(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_CURRENT_CONFIG).into_ref(&vm.ctx) + } + + #[pyattr(once)] + fn HKEY_DYN_DATA(vm: &VirtualMachine) -> PyRef { + PyHkey::new(Registry::HKEY_DYN_DATA).into_ref(&vm.ctx) } #[pyattr] - #[pyclass(name)] - #[derive(Clone, Debug, PyPayload)] - pub struct PyHKEYObject { - hkey: Arc>, + #[pyclass(name = "HKEYType")] + #[derive(Debug, PyPayload)] + struct PyHkey { + hkey: AtomicHKEY, } - // TODO: Fix - unsafe impl Send for PyHKEYObject {} - unsafe impl Sync for PyHKEYObject {} + unsafe impl Send for PyHkey {} + unsafe impl Sync for PyHkey {} - #[pyclass(with(AsNumber))] - impl PyHKEYObject { + impl PyHkey { + fn new(hkey: Registry::HKEY) -> Self { + Self { + hkey: AtomicHKEY::new(hkey), + } + } + + fn unary_fail(vm: &VirtualMachine) -> PyResult { + Err(vm.new_type_error(HKEY_ERR_MSG.to_owned())) + } + + fn binary_fail(vm: &VirtualMachine) -> PyResult { + Err(vm.new_type_error(HKEY_ERR_MSG.to_owned())) + } + + fn ternary_fail(vm: &VirtualMachine) -> PyResult { + Err(vm.new_type_error(HKEY_ERR_MSG.to_owned())) + } + } + + #[pyclass(with(AsNumber, Hashable))] + impl PyHkey { #[pygetset] fn handle(&self) -> usize { - *self.hkey.read() as usize - } - - #[pymethod] - fn __bool__(&self) -> bool { - !self.hkey.read().is_null() - } - - #[pymethod] - fn __int__(&self) -> usize { - *self.hkey.read() as usize - } - - #[pymethod] - fn __str__(&self) -> String { - format!("", *self.hkey.read() as usize) + self.hkey.load() as usize } #[pymethod] fn Close(&self, vm: &VirtualMachine) -> PyResult<()> { - let res = unsafe { Registry::RegCloseKey(*self.hkey.write()) }; - *self.hkey.write() = std::ptr::null_mut(); + // Atomically swap the handle with null and get the old value + let old_hkey = self.hkey.swap(std::ptr::null_mut()); + // Already closed - silently succeed + if old_hkey.is_null() { + return Ok(()); + } + let res = unsafe { Registry::RegCloseKey(old_hkey) }; if res == 0 { Ok(()) } else { - Err(vm.new_os_error("msg TODO".to_string())) + Err(vm.new_os_error(format!("RegCloseKey failed with error code: {res}"))) } } #[pymethod] fn Detach(&self) -> PyResult { - let hkey = *self.hkey.write(); - // std::mem::forget(self); - // TODO: Fix this - Ok(hkey as usize) + // Atomically swap the handle with null and return the old value + let old_hkey = self.hkey.swap(std::ptr::null_mut()); + Ok(old_hkey as usize) } - // fn AsHKEY(object: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - // if vm.is_none(object) { - // return Err(vm.new_type_error("cannot convert None to HKEY".to_owned())) - // } else if let Some(hkey) = object.downcast_ref::() { - // Ok(true) - // } else { - // Err(vm.new_type_error("The object is not a PyHKEY object".to_owned())) - // } - // } + #[pymethod] + fn __bool__(&self) -> bool { + !self.hkey.load().is_null() + } #[pymethod] fn __enter__(zelf: PyRef, _vm: &VirtualMachine) -> PyResult> { @@ -174,120 +224,107 @@ mod winreg { #[pymethod] fn __exit__(zelf: PyRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> { - let res = unsafe { Registry::RegCloseKey(*zelf.hkey.write()) }; - *zelf.hkey.write() = std::ptr::null_mut(); - if res == 0 { - Ok(()) - } else { - Err(vm.new_os_error("msg TODO".to_string())) + zelf.Close(vm) + } + + #[pymethod] + fn __int__(&self) -> usize { + self.hkey.load() as usize + } + + #[pymethod] + fn __str__(&self) -> String { + format!("", self.hkey.load()) + } + } + + impl Drop for PyHkey { + fn drop(&mut self) { + let hkey = self.hkey.swap(std::ptr::null_mut()); + if !hkey.is_null() { + unsafe { Registry::RegCloseKey(hkey) }; } } } - impl Drop for PyHKEYObject { - fn drop(&mut self) { - unsafe { - let hkey = *self.hkey.write(); - if !hkey.is_null() { - Registry::RegCloseKey(hkey); - } - } + impl Hashable for PyHkey { + // CPython uses PyObject_GenericHash which hashes the object's address + fn hash(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + Ok(zelf.get_id() as PyHash) } } pub const HKEY_ERR_MSG: &str = "bad operand type"; - impl PyHKEYObject { - pub fn new(hkey: *mut std::ffi::c_void) -> Self { - Self { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(hkey)), - } - } - - pub fn unary_fail(vm: &VirtualMachine) -> PyResult { - Err(vm.new_type_error(HKEY_ERR_MSG.to_owned())) - } - - pub fn binary_fail(vm: &VirtualMachine) -> PyResult { - Err(vm.new_type_error(HKEY_ERR_MSG.to_owned())) - } - - pub fn ternary_fail(vm: &VirtualMachine) -> PyResult { - Err(vm.new_type_error(HKEY_ERR_MSG.to_owned())) - } - } - - impl AsNumber for PyHKEYObject { + impl AsNumber for PyHkey { fn as_number() -> &'static PyNumberMethods { static AS_NUMBER: PyNumberMethods = PyNumberMethods { - add: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - subtract: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - multiply: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - remainder: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - divmod: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - power: Some(|_a, _b, _c, vm| PyHKEYObject::ternary_fail(vm)), - negative: Some(|_a, vm| PyHKEYObject::unary_fail(vm)), - positive: Some(|_a, vm| PyHKEYObject::unary_fail(vm)), - absolute: Some(|_a, vm| PyHKEYObject::unary_fail(vm)), + add: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + subtract: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + multiply: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + remainder: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + divmod: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + power: Some(|_a, _b, _c, vm| PyHkey::ternary_fail(vm)), + negative: Some(|_a, vm| PyHkey::unary_fail(vm)), + positive: Some(|_a, vm| PyHkey::unary_fail(vm)), + absolute: Some(|_a, vm| PyHkey::unary_fail(vm)), boolean: Some(|a, vm| { - if let Some(a) = a.downcast_ref::() { + if let Some(a) = a.downcast_ref::() { Ok(a.__bool__()) } else { - PyHKEYObject::unary_fail(vm)?; + PyHkey::unary_fail(vm)?; unreachable!() } }), - invert: Some(|_a, vm| PyHKEYObject::unary_fail(vm)), - lshift: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - rshift: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - and: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - xor: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), - or: Some(|_a, _b, vm| PyHKEYObject::binary_fail(vm)), + invert: Some(|_a, vm| PyHkey::unary_fail(vm)), + lshift: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + rshift: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + and: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + xor: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), + or: Some(|_a, _b, vm| PyHkey::binary_fail(vm)), int: Some(|a, vm| { - if let Some(a) = a.downcast_ref::() { + if let Some(a) = a.downcast_ref::() { Ok(vm.new_pyobj(a.__int__())) } else { - PyHKEYObject::unary_fail(vm)?; + PyHkey::unary_fail(vm)?; unreachable!() } }), - float: Some(|_a, vm| PyHKEYObject::unary_fail(vm)), + float: Some(|_a, vm| PyHkey::unary_fail(vm)), ..PyNumberMethods::NOT_IMPLEMENTED }; &AS_NUMBER } } - // TODO: Computer name can be `None` #[pyfunction] fn ConnectRegistry( computer_name: Option, - key: PyRef, + key: PyRef, vm: &VirtualMachine, - ) -> PyResult { + ) -> PyResult { if let Some(computer_name) = computer_name { let mut ret_key = std::ptr::null_mut(); - let wide_computer_name = to_utf16(computer_name); + let wide_computer_name = computer_name.to_wide_with_nul(); let res = unsafe { Registry::RegConnectRegistryW( wide_computer_name.as_ptr(), - *key.hkey.read(), + key.hkey.load(), &mut ret_key, ) }; if res == 0 { - Ok(PyHKEYObject::new(ret_key)) + Ok(PyHkey::new(ret_key)) } else { Err(vm.new_os_error(format!("error code: {res}"))) } } else { let mut ret_key = std::ptr::null_mut(); let res = unsafe { - Registry::RegConnectRegistryW(std::ptr::null_mut(), *key.hkey.read(), &mut ret_key) + Registry::RegConnectRegistryW(std::ptr::null_mut(), key.hkey.load(), &mut ret_key) }; if res == 0 { - Ok(PyHKEYObject::new(ret_key)) + Ok(PyHkey::new(ret_key)) } else { Err(vm.new_os_error(format!("error code: {res}"))) } @@ -295,18 +332,14 @@ mod winreg { } #[pyfunction] - fn CreateKey( - key: PyRef, - sub_key: String, - vm: &VirtualMachine, - ) -> PyResult { - let wide_sub_key = to_utf16(sub_key); + fn CreateKey(key: PyRef, sub_key: String, vm: &VirtualMachine) -> PyResult { + let wide_sub_key = sub_key.to_wide_with_nul(); let mut out_key = std::ptr::null_mut(); let res = unsafe { - Registry::RegCreateKeyW(*key.hkey.read(), wide_sub_key.as_ptr(), &mut out_key) + Registry::RegCreateKeyW(key.hkey.load(), wide_sub_key.as_ptr(), &mut out_key) }; if res == 0 { - Ok(PyHKEYObject::new(out_key)) + Ok(PyHkey::new(out_key)) } else { Err(vm.new_os_error(format!("error code: {res}"))) } @@ -315,7 +348,7 @@ mod winreg { #[derive(FromArgs, Debug)] struct CreateKeyExArgs { #[pyarg(any)] - key: PyRef, + key: PyRef, #[pyarg(any)] sub_key: String, #[pyarg(any, default = 0)] @@ -325,11 +358,11 @@ mod winreg { } #[pyfunction] - fn CreateKeyEx(args: CreateKeyExArgs, vm: &VirtualMachine) -> PyResult { - let wide_sub_key = to_utf16(args.sub_key); - let mut res: *mut std::ffi::c_void = core::ptr::null_mut(); + fn CreateKeyEx(args: CreateKeyExArgs, vm: &VirtualMachine) -> PyResult { + let wide_sub_key = args.sub_key.to_wide_with_nul(); + let mut res: Registry::HKEY = core::ptr::null_mut(); let err = unsafe { - let key = *args.key.hkey.read(); + let key = args.key.hkey.load(); Registry::RegCreateKeyExW( key, wide_sub_key.as_ptr(), @@ -343,9 +376,9 @@ mod winreg { ) }; if err == 0 { - Ok(PyHKEYObject { + Ok(PyHkey { #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(res)), + hkey: AtomicHKEY::new(res), }) } else { Err(vm.new_os_error(format!("error code: {err}"))) @@ -353,9 +386,26 @@ mod winreg { } #[pyfunction] - fn DeleteKey(key: PyRef, sub_key: String, vm: &VirtualMachine) -> PyResult<()> { - let wide_sub_key = to_utf16(sub_key); - let res = unsafe { Registry::RegDeleteKeyW(*key.hkey.read(), wide_sub_key.as_ptr()) }; + fn CloseKey(hkey: PyRef, vm: &VirtualMachine) -> PyResult<()> { + hkey.Close(vm) + } + + #[pyfunction] + fn DeleteKey(key: PyRef, sub_key: String, vm: &VirtualMachine) -> PyResult<()> { + let wide_sub_key = sub_key.to_wide_with_nul(); + let res = unsafe { Registry::RegDeleteKeyW(key.hkey.load(), wide_sub_key.as_ptr()) }; + if res == 0 { + Ok(()) + } else { + Err(vm.new_os_error(format!("error code: {res}"))) + } + } + + #[pyfunction] + fn DeleteValue(key: PyRef, value: Option, vm: &VirtualMachine) -> PyResult<()> { + let wide_value = value.map(|v| v.to_wide_with_nul()); + let value_ptr = wide_value.as_ref().map_or(std::ptr::null(), |v| v.as_ptr()); + let res = unsafe { Registry::RegDeleteValueW(key.hkey.load(), value_ptr) }; if res == 0 { Ok(()) } else { @@ -366,21 +416,21 @@ mod winreg { #[derive(FromArgs, Debug)] struct DeleteKeyExArgs { #[pyarg(any)] - key: PyRef, + key: PyRef, #[pyarg(any)] sub_key: String, + #[pyarg(any, default = windows_sys::Win32::System::Registry::KEY_WOW64_64KEY)] + access: u32, #[pyarg(any, default = 0)] reserved: u32, - #[pyarg(any, default = windows_sys::Win32::System::Registry::KEY_WOW64_32KEY)] - access: u32, } #[pyfunction] fn DeleteKeyEx(args: DeleteKeyExArgs, vm: &VirtualMachine) -> PyResult<()> { - let wide_sub_key = to_utf16(args.sub_key); + let wide_sub_key = args.sub_key.to_wide_with_nul(); let res = unsafe { Registry::RegDeleteKeyExW( - *args.key.hkey.read(), + args.key.hkey.load(), wide_sub_key.as_ptr(), args.access, args.reserved, @@ -393,36 +443,41 @@ mod winreg { } } - // #[pyfunction] - // fn EnumKey(key: PyRef, index: i32, vm: &VirtualMachine) -> PyResult { - // let mut tmpbuf = [0u16; 257]; - // let mut len = std::mem::sizeof(tmpbuf.len())/std::mem::sizeof(tmpbuf[0]); - // let res = unsafe { - // Registry::RegEnumKeyExW( - // *key.hkey.read(), - // index as u32, - // tmpbuf.as_mut_ptr(), - // &mut len, - // std::ptr::null_mut(), - // std::ptr::null_mut(), - // std::ptr::null_mut(), - // std::ptr::null_mut(), - // ) - // }; - // if res != 0 { - // return Err(vm.new_os_error(format!("error code: {res}"))); - // } - // let s = String::from_utf16(&tmpbuf[..len as usize]) - // .map_err(|e| vm.new_value_error(format!("UTF16 error: {e}")))?; - // Ok(s) - // } + #[pyfunction] + fn EnumKey(key: PyRef, index: i32, vm: &VirtualMachine) -> PyResult { + // The Windows docs claim that the max key name length is 255 + // characters, plus a terminating nul character. However, + // empirical testing demonstrates that it is possible to + // create a 256 character key that is missing the terminating + // nul. RegEnumKeyEx requires a 257 character buffer to + // retrieve such a key name. + let mut tmpbuf = [0u16; 257]; + let mut len = tmpbuf.len() as u32; + let res = unsafe { + Registry::RegEnumKeyExW( + key.hkey.load(), + index as u32, + tmpbuf.as_mut_ptr(), + &mut len, + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + if res != 0 { + return Err(vm.new_os_error(format!("error code: {res}"))); + } + String::from_utf16(&tmpbuf[..len as usize]) + .map_err(|e| vm.new_value_error(format!("UTF16 error: {e}"))) + } #[pyfunction] - fn EnumValue(hkey: PyRef, index: u32, vm: &VirtualMachine) -> PyResult { + fn EnumValue(hkey: PyRef, index: u32, vm: &VirtualMachine) -> PyResult { // Query registry for the required buffer sizes. let mut ret_value_size: u32 = 0; let mut ret_data_size: u32 = 0; - let hkey: *mut std::ffi::c_void = *hkey.hkey.read(); + let hkey: Registry::HKEY = hkey.hkey.load(); let rc = unsafe { Registry::RegQueryInfoKeyW( hkey, @@ -457,6 +512,7 @@ mod winreg { loop { let mut current_value_size = ret_value_size; let mut current_data_size = ret_data_size; + let mut reg_type: u32 = 0; let rc = unsafe { Registry::RegEnumValueW( hkey, @@ -464,11 +520,7 @@ mod winreg { ret_value_buf.as_mut_ptr(), &mut current_value_size as *mut u32, ptr::null_mut(), - { - // typ will hold the registry data type. - let mut t = 0u32; - &mut t - }, + &mut reg_type as *mut u32, ret_data_buf.as_mut_ptr(), &mut current_data_size as *mut u32, ) @@ -488,22 +540,6 @@ mod winreg { return Err(vm.new_os_error(format!("RegEnumValueW failed with error code {rc}"))); } - // At this point, current_value_size and current_data_size have been updated. - // Retrieve the registry type. - let mut reg_type: u32 = 0; - unsafe { - Registry::RegEnumValueW( - hkey, - index, - ret_value_buf.as_mut_ptr(), - &mut current_value_size as *mut u32, - ptr::null_mut(), - &mut reg_type as *mut u32, - ret_data_buf.as_mut_ptr(), - &mut current_data_size as *mut u32, - ) - }; - // Convert the registry value name from UTF‑16. let name_len = ret_value_buf .iter() @@ -529,8 +565,8 @@ mod winreg { } #[pyfunction] - fn FlushKey(key: PyRef, vm: &VirtualMachine) -> PyResult<()> { - let res = unsafe { Registry::RegFlushKey(*key.hkey.read()) }; + fn FlushKey(key: PyRef, vm: &VirtualMachine) -> PyResult<()> { + let res = unsafe { Registry::RegFlushKey(key.hkey.load()) }; if res == 0 { Ok(()) } else { @@ -540,16 +576,15 @@ mod winreg { #[pyfunction] fn LoadKey( - key: PyRef, + key: PyRef, sub_key: String, file_name: String, vm: &VirtualMachine, ) -> PyResult<()> { - let sub_key = to_utf16(sub_key); - let file_name = to_utf16(file_name); - let res = unsafe { - Registry::RegLoadKeyW(*key.hkey.read(), sub_key.as_ptr(), file_name.as_ptr()) - }; + let sub_key = sub_key.to_wide_with_nul(); + let file_name = file_name.to_wide_with_nul(); + let res = + unsafe { Registry::RegLoadKeyW(key.hkey.load(), sub_key.as_ptr(), file_name.as_ptr()) }; if res == 0 { Ok(()) } else { @@ -560,7 +595,7 @@ mod winreg { #[derive(Debug, FromArgs)] struct OpenKeyArgs { #[pyarg(any)] - key: PyRef, + key: PyRef, #[pyarg(any)] sub_key: String, #[pyarg(any, default = 0)] @@ -571,28 +606,32 @@ mod winreg { #[pyfunction] #[pyfunction(name = "OpenKeyEx")] - fn OpenKey(args: OpenKeyArgs, vm: &VirtualMachine) -> PyResult { - let wide_sub_key = to_utf16(args.sub_key); - let res: *mut *mut std::ffi::c_void = core::ptr::null_mut(); + fn OpenKey(args: OpenKeyArgs, vm: &VirtualMachine) -> PyResult { + let wide_sub_key = args.sub_key.to_wide_with_nul(); + let mut res: Registry::HKEY = std::ptr::null_mut(); let err = unsafe { - let key = *args.key.hkey.read(); - Registry::RegOpenKeyExW(key, wide_sub_key.as_ptr(), args.reserved, args.access, res) + let key = args.key.hkey.load(); + Registry::RegOpenKeyExW( + key, + wide_sub_key.as_ptr(), + args.reserved, + args.access, + &mut res, + ) }; if err == 0 { - unsafe { - Ok(PyHKEYObject { - #[allow(clippy::arc_with_non_send_sync)] - hkey: Arc::new(PyRwLock::new(*res)), - }) - } + Ok(PyHkey { + #[allow(clippy::arc_with_non_send_sync)] + hkey: AtomicHKEY::new(res), + }) } else { - Err(vm.new_os_error(format!("error code: {err}"))) + Err(os_error_from_windows_code(vm, err as i32, "RegOpenKeyEx")) } } #[pyfunction] - fn QueryInfoKey(key: PyRef, vm: &VirtualMachine) -> PyResult> { - let key = *key.hkey.read(); + fn QueryInfoKey(key: HKEYArg, vm: &VirtualMachine) -> PyResult> { + let key = key.0; let mut lpcsubkeys: u32 = 0; let mut lpcvalues: u32 = 0; let mut lpftlastwritetime: Foundation::FILETIME = unsafe { std::mem::zeroed() }; @@ -627,38 +666,110 @@ mod winreg { } #[pyfunction] - fn QueryValue(key: PyRef, sub_key: String, vm: &VirtualMachine) -> PyResult<()> { - let key = *key.hkey.read(); - let mut lpcbdata: i32 = 0; - // let mut lpdata = 0; - let wide_sub_key = to_utf16(sub_key); - let err = unsafe { - Registry::RegQueryValueW( - key, - wide_sub_key.as_ptr(), - std::ptr::null_mut(), - &mut lpcbdata, - ) - }; + fn QueryValue(key: HKEYArg, sub_key: Option, vm: &VirtualMachine) -> PyResult { + let hkey = key.0; - if err != 0 { - return Err(vm.new_os_error(format!("error code: {err}"))); + if hkey == Registry::HKEY_PERFORMANCE_DATA { + return Err(os_error_from_windows_code( + vm, + Foundation::ERROR_INVALID_HANDLE as i32, + "RegQueryValue", + )); } - Ok(()) + // Open subkey if provided and non-empty + let child_key = if let Some(ref sk) = sub_key { + if !sk.is_empty() { + let wide_sub_key = sk.to_wide_with_nul(); + let mut out_key = std::ptr::null_mut(); + let res = unsafe { + Registry::RegOpenKeyExW( + hkey, + wide_sub_key.as_ptr(), + 0, + Registry::KEY_QUERY_VALUE, + &mut out_key, + ) + }; + if res != 0 { + return Err(os_error_from_windows_code(vm, res as i32, "RegOpenKeyEx")); + } + Some(out_key) + } else { + None + } + } else { + None + }; + + let target_key = child_key.unwrap_or(hkey); + let mut buf_size: u32 = 256; + let mut buffer: Vec = vec![0; buf_size as usize]; + let mut reg_type: u32 = 0; + + // Loop to handle ERROR_MORE_DATA + let result = loop { + let mut size = buf_size; + let res = unsafe { + Registry::RegQueryValueExW( + target_key, + std::ptr::null(), // NULL value name for default value + std::ptr::null_mut(), + &mut reg_type, + buffer.as_mut_ptr(), + &mut size, + ) + }; + if res == ERROR_MORE_DATA { + buf_size *= 2; + buffer.resize(buf_size as usize, 0); + continue; + } + if res == Foundation::ERROR_FILE_NOT_FOUND { + // Return empty string if there's no default value + break Ok(String::new()); + } + if res != 0 { + break Err(os_error_from_windows_code( + vm, + res as i32, + "RegQueryValueEx", + )); + } + if reg_type != Registry::REG_SZ { + break Err(os_error_from_windows_code( + vm, + Foundation::ERROR_INVALID_DATA as i32, + "RegQueryValue", + )); + } + + // Convert UTF-16 to String + let u16_slice = bytes_as_wide_slice(&buffer[..size as usize]); + let len = u16_slice + .iter() + .position(|&c| c == 0) + .unwrap_or(u16_slice.len()); + break String::from_utf16(&u16_slice[..len]) + .map_err(|e| vm.new_value_error(format!("UTF16 error: {e}"))); + }; + + // Close child key if we opened one + if let Some(ck) = child_key { + unsafe { Registry::RegCloseKey(ck) }; + } + + result } #[pyfunction] - fn QueryValueEx( - key: PyRef, - name: String, - vm: &VirtualMachine, - ) -> PyResult { - let wide_name = to_utf16(name); - let mut buf_size = 0; + fn QueryValueEx(key: HKEYArg, name: String, vm: &VirtualMachine) -> PyResult> { + let hkey = key.0; + let wide_name = name.to_wide_with_nul(); + let mut buf_size: u32 = 0; let res = unsafe { Registry::RegQueryValueExW( - *key.hkey.read(), + hkey, wide_name.as_ptr(), std::ptr::null_mut(), std::ptr::null_mut(), @@ -666,35 +777,62 @@ mod winreg { &mut buf_size, ) }; - // TODO: res == ERROR_MORE_DATA - if res != 0 { - return Err(vm.new_os_error(format!("error code: {res}"))); + // Handle ERROR_MORE_DATA by using a default buffer size + if res == ERROR_MORE_DATA || buf_size == 0 { + buf_size = 256; + } else if res != 0 { + return Err(os_error_from_windows_code( + vm, + res as i32, + "RegQueryValueEx", + )); } - let mut retBuf = Vec::with_capacity(buf_size as usize); + + let mut ret_buf = vec![0u8; buf_size as usize]; let mut typ = 0; - let res = unsafe { - Registry::RegQueryValueExW( - *key.hkey.read(), - wide_name.as_ptr(), - std::ptr::null_mut(), - &mut typ, - retBuf.as_mut_ptr(), - &mut buf_size, - ) - }; - // TODO: res == ERROR_MORE_DATA - if res != 0 { - return Err(vm.new_os_error(format!("error code: {res}"))); + let mut ret_size: u32; + + // Loop to handle ERROR_MORE_DATA + loop { + ret_size = buf_size; + let res = unsafe { + Registry::RegQueryValueExW( + hkey, + wide_name.as_ptr(), + std::ptr::null_mut(), + &mut typ, + ret_buf.as_mut_ptr(), + &mut ret_size, + ) + }; + + if res != ERROR_MORE_DATA { + if res != 0 { + return Err(os_error_from_windows_code( + vm, + res as i32, + "RegQueryValueEx", + )); + } + break; + } + + // Double buffer size and retry + buf_size *= 2; + ret_buf.resize(buf_size as usize, 0); } - let obj = reg_to_py(vm, retBuf.as_slice(), typ)?; - Ok(obj) + + // Only pass the bytes actually returned by the API + let obj = reg_to_py(vm, &ret_buf[..ret_size as usize], typ)?; + // Return tuple (value, type) + Ok(vm.ctx.new_tuple(vec![obj, vm.ctx.new_int(typ).into()])) } #[pyfunction] - fn SaveKey(key: PyRef, file_name: String, vm: &VirtualMachine) -> PyResult<()> { - let file_name = to_utf16(file_name); + fn SaveKey(key: PyRef, file_name: String, vm: &VirtualMachine) -> PyResult<()> { + let file_name = file_name.to_wide_with_nul(); let res = unsafe { - Registry::RegSaveKeyW(*key.hkey.read(), file_name.as_ptr(), std::ptr::null_mut()) + Registry::RegSaveKeyW(key.hkey.load(), file_name.as_ptr(), std::ptr::null_mut()) }; if res == 0 { Ok(()) @@ -705,7 +843,7 @@ mod winreg { #[pyfunction] fn SetValue( - key: PyRef, + key: PyRef, sub_key: String, typ: u32, value: String, @@ -715,32 +853,63 @@ mod winreg { return Err(vm.new_type_error("type must be winreg.REG_SZ")); } - let wide_sub_key = to_utf16(sub_key); - - // TODO: Value check - if *key.hkey.read() == Registry::HKEY_PERFORMANCE_DATA { - return Err(vm.new_os_error("Cannot set value on HKEY_PERFORMANCE_DATA")); + let hkey = key.hkey.load(); + if hkey == Registry::HKEY_PERFORMANCE_DATA { + return Err(os_error_from_windows_code( + vm, + Foundation::ERROR_INVALID_HANDLE as i32, + "RegSetValue", + )); } - // if (sub_key && sub_key[0]) { - // // TODO: create key - // } + // Create subkey if sub_key is non-empty + let child_key = if !sub_key.is_empty() { + let wide_sub_key = sub_key.to_wide_with_nul(); + let mut out_key = std::ptr::null_mut(); + let res = unsafe { + Registry::RegCreateKeyExW( + hkey, + wide_sub_key.as_ptr(), + 0, + std::ptr::null(), + 0, + Registry::KEY_SET_VALUE, + std::ptr::null(), + &mut out_key, + std::ptr::null_mut(), + ) + }; + if res != 0 { + return Err(os_error_from_windows_code(vm, res as i32, "RegCreateKeyEx")); + } + Some(out_key) + } else { + None + }; + let target_key = child_key.unwrap_or(hkey); + // Convert value to UTF-16 for Wide API + let wide_value = value.to_wide_with_nul(); let res = unsafe { Registry::RegSetValueExW( - *key.hkey.read(), - wide_sub_key.as_ptr(), + target_key, + std::ptr::null(), // value name is NULL 0, typ, - value.as_ptr(), - value.len() as u32, + wide_value.as_ptr() as *const u8, + (wide_value.len() * 2) as u32, // byte count ) }; + // Close child key if we created one + if let Some(ck) = child_key { + unsafe { Registry::RegCloseKey(ck) }; + } + if res == 0 { Ok(()) } else { - Err(vm.new_os_error(format!("error code: {res}"))) + Err(os_error_from_windows_code(vm, res as i32, "RegSetValueEx")) } } @@ -764,11 +933,7 @@ mod winreg { } } REG_SZ | REG_EXPAND_SZ => { - // Treat the data as a UTF-16 string. - let u16_count = ret_data.len() / 2; - let u16_slice = unsafe { - std::slice::from_raw_parts(ret_data.as_ptr() as *const u16, u16_count) - }; + let u16_slice = bytes_as_wide_slice(ret_data); // Only use characters up to the first NUL. let len = u16_slice .iter() @@ -782,24 +947,32 @@ mod winreg { if ret_data.is_empty() { Ok(vm.ctx.new_list(vec![]).into()) } else { - let u16_count = ret_data.len() / 2; - let u16_slice = unsafe { - std::slice::from_raw_parts(ret_data.as_ptr() as *const u16, u16_count) + let u16_slice = bytes_as_wide_slice(ret_data); + let u16_count = u16_slice.len(); + + // Remove trailing null if present (like countStrings) + let len = if u16_count > 0 && u16_slice[u16_count - 1] == 0 { + u16_count - 1 + } else { + u16_count }; + let mut strings: Vec = Vec::new(); let mut start = 0; - for (i, &c) in u16_slice.iter().enumerate() { - if c == 0 { - // An empty string signals the end. - if start == i { - break; - } + for i in 0..len { + if u16_slice[i] == 0 { let s = String::from_utf16(&u16_slice[start..i]) .map_err(|e| vm.new_value_error(format!("UTF16 error: {e}")))?; strings.push(vm.ctx.new_str(s).into()); start = i + 1; } } + // Handle last string if not null-terminated + if start < len { + let s = String::from_utf16(&u16_slice[start..len]) + .map_err(|e| vm.new_value_error(format!("UTF16 error: {e}")))?; + strings.push(vm.ctx.new_str(s).into()); + } Ok(vm.ctx.new_list(strings).into()) } } @@ -817,44 +990,97 @@ mod winreg { fn py2reg(value: PyObjectRef, typ: u32, vm: &VirtualMachine) -> PyResult>> { match typ { REG_DWORD => { - let val = value.downcast_ref::(); - if val.is_none() { - return Err(vm.new_type_error("value must be an integer".to_string())); + if vm.is_none(&value) { + return Ok(Some(0u32.to_le_bytes().to_vec())); } - let val = val.unwrap().as_bigint().to_u32().unwrap(); + let val = value + .downcast_ref::() + .ok_or_else(|| vm.new_type_error("value must be an integer".to_string()))?; + let bigint = val.as_bigint(); + // Check for negative value - raise OverflowError + if bigint.sign() == Sign::Minus { + return Err(vm.new_overflow_error("int too big to convert".to_string())); + } + let val = bigint + .to_u32() + .ok_or_else(|| vm.new_overflow_error("int too big to convert".to_string()))?; Ok(Some(val.to_le_bytes().to_vec())) } REG_QWORD => { - let val = value.downcast_ref::(); - if val.is_none() { - return Err(vm.new_type_error("value must be an integer".to_string())); + if vm.is_none(&value) { + return Ok(Some(0u64.to_le_bytes().to_vec())); } - let val = val.unwrap().as_bigint().to_u64().unwrap(); + let val = value + .downcast_ref::() + .ok_or_else(|| vm.new_type_error("value must be an integer".to_string()))?; + let bigint = val.as_bigint(); + // Check for negative value - raise OverflowError + if bigint.sign() == Sign::Minus { + return Err(vm.new_overflow_error("int too big to convert".to_string())); + } + let val = bigint + .to_u64() + .ok_or_else(|| vm.new_overflow_error("int too big to convert".to_string()))?; Ok(Some(val.to_le_bytes().to_vec())) } - // REG_SZ is fallthrough - REG_EXPAND_SZ => { - Err(vm - .new_type_error("TODO: RUSTPYTHON REG_EXPAND_SZ is not supported".to_string())) + REG_SZ | REG_EXPAND_SZ => { + if vm.is_none(&value) { + // Return empty string as UTF-16 null terminator + return Ok(Some(vec![0u8, 0u8])); + } + let s = value + .downcast::() + .map_err(|_| vm.new_type_error("value must be a string".to_string()))?; + let wide = s.as_str().to_wide_with_nul(); + // Convert Vec to Vec + let bytes: Vec = wide.iter().flat_map(|&c| c.to_le_bytes()).collect(); + Ok(Some(bytes)) } REG_MULTI_SZ => { - Err(vm.new_type_error("TODO: RUSTPYTHON REG_MULTI_SZ is not supported".to_string())) + if vm.is_none(&value) { + // Empty list = double null terminator + return Ok(Some(vec![0u8, 0u8, 0u8, 0u8])); + } + let list = value.downcast::().map_err(|_| { + vm.new_type_error("value must be a list of strings".to_string()) + })?; + + let mut bytes: Vec = Vec::new(); + for item in list.borrow_vec().iter() { + let s = item + .downcast_ref::() + .ok_or_else(|| { + vm.new_type_error("list items must be strings".to_string()) + })?; + let wide = s.as_str().to_wide_with_nul(); + bytes.extend(wide.iter().flat_map(|&c| c.to_le_bytes())); + } + // Add final null terminator (double null at end) + bytes.extend([0u8, 0u8]); + Ok(Some(bytes)) } - // REG_BINARY is fallthrough + // REG_BINARY and other types _ => { if vm.is_none(&value) { return Ok(None); } - Err(vm.new_type_error("TODO: RUSTPYTHON Not supported".to_string())) + // Try to get bytes + if let Some(bytes) = value.downcast_ref::() { + return Ok(Some(bytes.as_bytes().to_vec())); + } + Err(vm.new_type_error(format!( + "Objects of type '{}' can not be used as binary registry values", + value.class().name() + ))) } } } #[pyfunction] fn SetValueEx( - key: PyRef, + key: PyRef, value_name: String, - _reserved: u32, + _reserved: PyObjectRef, typ: u32, value: PyObjectRef, vm: &VirtualMachine, @@ -863,10 +1089,10 @@ mod winreg { Ok(Some(v)) => { let len = v.len() as u32; let ptr = v.as_ptr(); - let wide_value_name = to_utf16(value_name); + let wide_value_name = value_name.to_wide_with_nul(); let res = unsafe { Registry::RegSetValueExW( - *key.hkey.read(), + key.hkey.load(), wide_value_name.as_ptr(), 0, typ, @@ -881,10 +1107,10 @@ mod winreg { Ok(None) => { let len = 0; let ptr = std::ptr::null(); - let wide_value_name = to_utf16(value_name); + let wide_value_name = value_name.to_wide_with_nul(); let res = unsafe { Registry::RegSetValueExW( - *key.hkey.read(), + key.hkey.load(), wide_value_name.as_ptr(), 0, typ, @@ -896,14 +1122,14 @@ mod winreg { return Err(vm.new_os_error(format!("error code: {res}"))); } } - Err(_) => return Err(vm.new_type_error("value must be an integer".to_string())), + Err(e) => return Err(e), } Ok(()) } #[pyfunction] - fn EnableReflectionKey(key: PyRef, vm: &VirtualMachine) -> PyResult<()> { - let res = unsafe { Registry::RegEnableReflectionKey(*key.hkey.read()) }; + fn DisableReflectionKey(key: PyRef, vm: &VirtualMachine) -> PyResult<()> { + let res = unsafe { Registry::RegDisableReflectionKey(key.hkey.load()) }; if res == 0 { Ok(()) } else { @@ -912,21 +1138,56 @@ mod winreg { } #[pyfunction] - fn ExpandEnvironmentStrings(i: String) -> PyResult { - let mut out = vec![0; 1024]; - let r = unsafe { - windows_sys::Win32::System::Environment::ExpandEnvironmentStringsA( - i.as_ptr(), - out.as_mut_ptr(), - out.len() as u32, + fn EnableReflectionKey(key: PyRef, vm: &VirtualMachine) -> PyResult<()> { + let res = unsafe { Registry::RegEnableReflectionKey(key.hkey.load()) }; + if res == 0 { + Ok(()) + } else { + Err(vm.new_os_error(format!("error code: {res}"))) + } + } + + #[pyfunction] + fn QueryReflectionKey(key: PyRef, vm: &VirtualMachine) -> PyResult { + let mut result: windows_sys::Win32::Foundation::BOOL = 0; + let res = unsafe { Registry::RegQueryReflectionKey(key.hkey.load(), &mut result) }; + if res == 0 { + Ok(result != 0) + } else { + Err(vm.new_os_error(format!("error code: {res}"))) + } + } + + #[pyfunction] + fn ExpandEnvironmentStrings(i: String, vm: &VirtualMachine) -> PyResult { + let wide_input = i.to_wide_with_nul(); + + // First call with size=0 to get required buffer size + let required_size = unsafe { + windows_sys::Win32::System::Environment::ExpandEnvironmentStringsW( + wide_input.as_ptr(), + std::ptr::null_mut(), + 0, ) }; - let s = String::from_utf8(out[..r as usize].to_vec()) - .unwrap() - .replace("\0", "") - .replace("\x02", "") - .to_string(); + if required_size == 0 { + return Err(vm.new_os_error("ExpandEnvironmentStringsW failed".to_string())); + } - Ok(s) + // Allocate buffer with exact size and expand + let mut out = vec![0u16; required_size as usize]; + let r = unsafe { + windows_sys::Win32::System::Environment::ExpandEnvironmentStringsW( + wide_input.as_ptr(), + out.as_mut_ptr(), + required_size, + ) + }; + if r == 0 { + return Err(vm.new_os_error("ExpandEnvironmentStringsW failed".to_string())); + } + + let len = out.iter().position(|&c| c == 0).unwrap_or(out.len()); + String::from_utf16(&out[..len]).map_err(|e| vm.new_value_error(format!("UTF16 error: {e}"))) } }