From d0b110632d26e4934452c69c47422909de2fe6fd Mon Sep 17 00:00:00 2001 From: Rafael Rivera Date: Thu, 5 Sep 2024 14:35:17 -0700 Subject: [PATCH 1/2] iterators --- crates/libs/registry/src/bindings.rs | 2 + crates/libs/registry/src/key_iterator.rs | 81 ++++++++-------- crates/libs/registry/src/lib.rs | 40 ++++++++ crates/libs/registry/src/value_iterator.rs | 102 ++++++++++----------- crates/tests/misc/registry/tests/keys.rs | 11 +++ crates/tests/misc/registry/tests/values.rs | 18 ++++ crates/tools/bindings/src/registry.txt | 2 + 7 files changed, 157 insertions(+), 99 deletions(-) diff --git a/crates/libs/registry/src/bindings.rs b/crates/libs/registry/src/bindings.rs index ea270ec24f..f20db0161a 100644 --- a/crates/libs/registry/src/bindings.rs +++ b/crates/libs/registry/src/bindings.rs @@ -20,7 +20,9 @@ windows_targets::link!("kernel32.dll" "system" fn HeapAlloc(hheap : HANDLE, dwfl windows_targets::link!("kernel32.dll" "system" fn HeapFree(hheap : HANDLE, dwflags : HEAP_FLAGS, lpmem : *const core::ffi::c_void) -> BOOL); pub type BOOL = i32; pub const ERROR_INVALID_DATA: WIN32_ERROR = 13u32; +pub const ERROR_MORE_DATA: WIN32_ERROR = 234u32; pub const ERROR_NO_MORE_ITEMS: WIN32_ERROR = 259u32; +pub const ERROR_SUCCESS: WIN32_ERROR = 0u32; #[repr(C)] #[derive(Clone, Copy)] pub struct FILETIME { diff --git a/crates/libs/registry/src/key_iterator.rs b/crates/libs/registry/src/key_iterator.rs index aca76c5ff1..638cba7279 100644 --- a/crates/libs/registry/src/key_iterator.rs +++ b/crates/libs/registry/src/key_iterator.rs @@ -3,66 +3,59 @@ use super::*; /// An iterator of registry key names. pub struct KeyIterator<'a> { key: &'a Key, - range: core::ops::Range, + idx: u32, name: Vec, } impl<'a> KeyIterator<'a> { pub(crate) fn new(key: &'a Key) -> Result { - let mut count = 0; - let mut max_len = 0; + let info = get_key_info(key)?; - let result = unsafe { - RegQueryInfoKeyW( - key.0, - null_mut(), - null_mut(), - null_mut(), - &mut count, - &mut max_len, - null_mut(), - null_mut(), - null_mut(), - null_mut(), - null_mut(), - null_mut(), - ) - }; - - win32_error(result).map(|_| Self { + Ok(Self { key, - range: 0..count as usize, - name: vec![0; max_len as usize + 1], + idx: 0, + name: vec![0; (info.subkey_name_max + 1) as usize], }) } + + fn resize(&mut self) -> Result<()> { + let info = get_key_info(self.key)?; + self.name.resize((info.subkey_name_max + 1) as usize, 0); + Ok(()) + } } impl<'a> Iterator for KeyIterator<'a> { type Item = String; fn next(&mut self) -> Option { - self.range.next().and_then(|index| { - let mut len = self.name.len() as u32; - - let result = unsafe { - RegEnumKeyExW( - self.key.0, - index as u32, - self.name.as_mut_ptr(), - &mut len, - null(), - null_mut(), - null_mut(), - null_mut(), - ) - }; + let mut len = self.name.len() as u32; + let result = unsafe { + RegEnumKeyExW( + self.key.0, + self.idx, + self.name.as_mut_ptr(), + &mut len, + null_mut(), + null_mut(), + null_mut(), + null_mut(), + ) + }; - if result != 0 { - debug_assert_eq!(result, ERROR_NO_MORE_ITEMS); - None - } else { - Some(String::from_utf16_lossy(&self.name[0..len as usize])) + if result == ERROR_MORE_DATA { + if self.resize().is_err() { + return None; } - }) + return self.next(); + } + + if result != 0 { + debug_assert_eq!(result, ERROR_NO_MORE_ITEMS); + return None; + } + + self.idx += 1; + Some(String::from_utf16_lossy(&self.name[0..len as usize])) } } diff --git a/crates/libs/registry/src/lib.rs b/crates/libs/registry/src/lib.rs index cf35f9d628..109244d817 100644 --- a/crates/libs/registry/src/lib.rs +++ b/crates/libs/registry/src/lib.rs @@ -78,3 +78,43 @@ fn from_le_bytes(ty: Type, from: &[u8]) -> Result { fn as_bytes(value: &HSTRING) -> &[u8] { unsafe { core::slice::from_raw_parts(value.as_ptr() as *const _, (value.len() + 1) * 2) } } + +struct KeyInfo { + /// Size of the key's subkey with the longest name, in Unicode characters, not including the terminating null character. + subkey_name_max: u32, + + /// Size of the key's longest value name, in Unicode characters. The size does not include the terminating null character. + value_name_max: u32, + + /// Size of the longest data component among the key's values, in bytes. + value_data_max: u32, +} + +fn get_key_info(key: &Key) -> Result { + let mut subkey_name_max = 0; + let mut value_name_max = 0; + let mut value_data_max = 0; + + let result = unsafe { + RegQueryInfoKeyW( + key.0, + null_mut(), + null_mut(), + null_mut(), + null_mut(), + &mut subkey_name_max, + null_mut(), + null_mut(), + &mut value_name_max, + &mut value_data_max, + null_mut(), + null_mut(), + ) + }; + + win32_error(result).map(|_| KeyInfo { + subkey_name_max, + value_name_max, + value_data_max, + }) +} diff --git a/crates/libs/registry/src/value_iterator.rs b/crates/libs/registry/src/value_iterator.rs index e11d319b5a..12ec679ede 100644 --- a/crates/libs/registry/src/value_iterator.rs +++ b/crates/libs/registry/src/value_iterator.rs @@ -3,80 +3,72 @@ use super::*; /// An iterator of registry values. pub struct ValueIterator<'a> { key: &'a Key, - range: core::ops::Range, + idx: u32, name: Vec, data: Data, } impl<'a> ValueIterator<'a> { pub(crate) fn new(key: &'a Key) -> Result { - let mut count = 0; - let mut name_max_len = 0; - let mut value_max_len = 0; - - let result = unsafe { - RegQueryInfoKeyW( - key.0, - null_mut(), - null_mut(), - null_mut(), - null_mut(), - null_mut(), - null_mut(), - &mut count, - &mut name_max_len, - &mut value_max_len, - null_mut(), - null_mut(), - ) - }; - - win32_error(result)?; + let info = get_key_info(key)?; Ok(Self { key, - range: 0..count as usize, - name: vec![0; name_max_len as usize + 1], - data: Data::new(value_max_len as usize), + idx: 0, + name: vec![0; (info.value_name_max + 1) as usize], + data: Data::new(info.value_data_max as usize), }) } + + fn resize(&mut self) -> Result<()> { + let info = get_key_info(self.key)?; + self.name.resize((info.value_name_max + 1) as usize, 0); + self.data = Data::new(info.value_data_max as usize); + Ok(()) + } } impl<'a> Iterator for ValueIterator<'a> { type Item = (String, Value); fn next(&mut self) -> Option { - self.range.next().and_then(|index| { - let mut ty = 0; - let mut name_len = self.name.len() as u32; - let mut data_len = self.data.len() as u32; + let mut ty = 0; + let mut name_len = self.name.len() as u32; + let mut data_len = self.data.len() as u32; - let result = unsafe { - RegEnumValueW( - self.key.0, - index as u32, - self.name.as_mut_ptr(), - &mut name_len, - core::ptr::null(), - &mut ty, - self.data.as_mut_ptr(), - &mut data_len, - ) - }; + let result = unsafe { + RegEnumValueW( + self.key.0, + self.idx, + self.name.as_mut_ptr(), + &mut name_len, + core::ptr::null(), + &mut ty, + self.data.as_mut_ptr(), + &mut data_len, + ) + }; - if result != 0 { - debug_assert_eq!(result, ERROR_NO_MORE_ITEMS); - None - } else { - let name = String::from_utf16_lossy(&self.name[0..name_len as usize]); - Some(( - name, - Value { - data: Data::from_slice(&self.data[0..data_len as usize]), - ty: ty.into(), - }, - )) + if result == ERROR_MORE_DATA { + if self.resize().is_err() { + return None; } - }) + return self.next(); + } + + if result != ERROR_SUCCESS { + debug_assert_eq!(result, ERROR_NO_MORE_ITEMS); + return None; + } + + self.idx += 1; + let name = String::from_utf16_lossy(&self.name[0..name_len as usize]); + Some(( + name, + Value { + data: Data::from_slice(&self.data[0..data_len as usize]), + ty: ty.into(), + }, + )) } } diff --git a/crates/tests/misc/registry/tests/keys.rs b/crates/tests/misc/registry/tests/keys.rs index 4a8114bf33..9fbac8a1c8 100644 --- a/crates/tests/misc/registry/tests/keys.rs +++ b/crates/tests/misc/registry/tests/keys.rs @@ -14,6 +14,17 @@ fn keys() -> Result<()> { let names: Vec = key.keys()?.collect(); assert_eq!(names, ["one", "three", "two"]); + let mut names = Vec::::new(); + let iterator = key.keys()?; + for name in iterator { + if name == "one" { + key.remove_tree("three")?; + key.create("seventy")?; + } + names.push(name.clone()); + } + assert_eq!(names, ["one", "seventy", "two"]); + let err = key.open("missing").unwrap_err(); assert_eq!(err.code(), HRESULT(0x80070002u32 as i32)); // HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND) assert_eq!(err.message(), "The system cannot find the file specified."); diff --git a/crates/tests/misc/registry/tests/values.rs b/crates/tests/misc/registry/tests/values.rs index 9ad13514c7..eca900a732 100644 --- a/crates/tests/misc/registry/tests/values.rs +++ b/crates/tests/misc/registry/tests/values.rs @@ -22,5 +22,23 @@ fn values() -> Result<()> { ] ); + let mut names = Vec::<(String, Value)>::new(); + let iterator = key.values()?; + for (name, value) in iterator { + if name == "string" { + key.set_string("string-two", "hello world two")?; + } + names.push((name, value)); + } + assert_eq!( + names, + [ + ("u32".to_string(), Value::from(123u32)), + ("u64".to_string(), Value::from(456u64)), + ("string".to_string(), Value::from("hello world")), + ("string-two".to_string(), Value::from("hello world two")), + ] + ); + Ok(()) } diff --git a/crates/tools/bindings/src/registry.txt b/crates/tools/bindings/src/registry.txt index 917ea742a7..3cc71397c3 100644 --- a/crates/tools/bindings/src/registry.txt +++ b/crates/tools/bindings/src/registry.txt @@ -3,7 +3,9 @@ --filter Windows.Win32.Foundation.ERROR_INVALID_DATA + Windows.Win32.Foundation.ERROR_MORE_DATA Windows.Win32.Foundation.ERROR_NO_MORE_ITEMS + Windows.Win32.Foundation.ERROR_SUCCESS Windows.Win32.System.Memory.GetProcessHeap Windows.Win32.System.Memory.HeapAlloc Windows.Win32.System.Memory.HeapFree From ba0a1d27f495bedb21431a77a282a5e4135e7662 Mon Sep 17 00:00:00 2001 From: Rafael Rivera Date: Thu, 5 Sep 2024 16:22:58 -0700 Subject: [PATCH 2/2] constant --- crates/libs/registry/src/key_iterator.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/libs/registry/src/key_iterator.rs b/crates/libs/registry/src/key_iterator.rs index 638cba7279..a9eaecce86 100644 --- a/crates/libs/registry/src/key_iterator.rs +++ b/crates/libs/registry/src/key_iterator.rs @@ -50,7 +50,7 @@ impl<'a> Iterator for KeyIterator<'a> { return self.next(); } - if result != 0 { + if result != ERROR_SUCCESS { debug_assert_eq!(result, ERROR_NO_MORE_ITEMS); return None; }