Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make registry key/value iterators lazy #3260

Closed
wants to merge 2 commits into from
Closed

Make registry key/value iterators lazy #3260

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d0b1106
iterators
riverar Sep 5, 2024
ba0a1d2
constant
riverar Sep 5, 2024
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
2 changes: 2 additions & 0 deletions crates/libs/registry/src/bindings.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
81 changes: 37 additions & 44 deletions crates/libs/registry/src/key_iterator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,66 +3,59 @@ use super::*;
/// An iterator of registry key names.
pub struct KeyIterator<'a> {
key: &'a Key,
range: core::ops::Range<usize>,
idx: u32,
riverar marked this conversation as resolved.
Show resolved Hide resolved
name: Vec<u16>,
}

impl<'a> KeyIterator<'a> {
pub(crate) fn new(key: &'a Key) -> Result<Self> {
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::Item> {
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();
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is recursive - not necessarily bad in this scenario but worth calling out if it's not obvious.

}

if result != ERROR_SUCCESS {
debug_assert_eq!(result, ERROR_NO_MORE_ITEMS);
return None;
}

self.idx += 1;
Some(String::from_utf16_lossy(&self.name[0..len as usize]))
}
}
40 changes: 40 additions & 0 deletions crates/libs/registry/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,3 +78,43 @@ fn from_le_bytes(ty: Type, from: &[u8]) -> Result<u64> {
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<KeyInfo> {
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,
})
}
102 changes: 47 additions & 55 deletions crates/libs/registry/src/value_iterator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,80 +3,72 @@ use super::*;
/// An iterator of registry values.
pub struct ValueIterator<'a> {
key: &'a Key,
range: core::ops::Range<usize>,
idx: u32,
name: Vec<u16>,
data: Data,
}

impl<'a> ValueIterator<'a> {
pub(crate) fn new(key: &'a Key) -> Result<Self> {
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(())
}
Comment on lines +17 to +28
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not a big deal but might be nice to default initialize and then call resize in the new method to avoid the redundant initialization here and in the other iterator.

}

impl<'a> Iterator for ValueIterator<'a> {
type Item = (String, Value);

fn next(&mut self) -> Option<Self::Item> {
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(),
},
))
}
}
11 changes: 11 additions & 0 deletions crates/tests/misc/registry/tests/keys.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,17 @@ fn keys() -> Result<()> {
let names: Vec<String> = key.keys()?.collect();
assert_eq!(names, ["one", "three", "two"]);

let mut names = Vec::<String>::new();
let iterator = key.keys()?;
for name in iterator {
if name == "one" {
key.remove_tree("three")?;
key.create("seventy")?;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe make it obvious that this requires a larger buffer...

}
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.");
Expand Down
18 changes: 18 additions & 0 deletions crates/tests/misc/registry/tests/values.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")?;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here.

}
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(())
}
2 changes: 2 additions & 0 deletions crates/tools/bindings/src/registry.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down