Fix alignment issue in windows-registry value iterator (#3052)

microsoft · May 23, 2024 · 40d35fa · 40d35fa
1 parent d06694a
commit 40d35fa
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 5 deletions.
diff --git a/crates/libs/registry/src/bindings.rs b/crates/libs/registry/src/bindings.rs
@@ -15,15 +15,21 @@ windows_targets::link!("advapi32.dll" "system" fn RegOpenKeyExW(hkey : HKEY, lps
 windows_targets::link!("advapi32.dll" "system" fn RegQueryInfoKeyW(hkey : HKEY, lpclass : PWSTR, lpcchclass : *mut u32, lpreserved : *const u32, lpcsubkeys : *mut u32, lpcbmaxsubkeylen : *mut u32, lpcbmaxclasslen : *mut u32, lpcvalues : *mut u32, lpcbmaxvaluenamelen : *mut u32, lpcbmaxvaluelen : *mut u32, lpcbsecuritydescriptor : *mut u32, lpftlastwritetime : *mut FILETIME) -> WIN32_ERROR);
 windows_targets::link!("advapi32.dll" "system" fn RegQueryValueExW(hkey : HKEY, lpvaluename : PCWSTR, lpreserved : *const u32, lptype : *mut REG_VALUE_TYPE, lpdata : *mut u8, lpcbdata : *mut u32) -> WIN32_ERROR);
 windows_targets::link!("advapi32.dll" "system" fn RegSetValueExW(hkey : HKEY, lpvaluename : PCWSTR, reserved : u32, dwtype : REG_VALUE_TYPE, lpdata : *const u8, cbdata : u32) -> WIN32_ERROR);
+windows_targets::link!("kernel32.dll" "system" fn GetProcessHeap() -> HANDLE);
+windows_targets::link!("kernel32.dll" "system" fn HeapAlloc(hheap : HANDLE, dwflags : HEAP_FLAGS, dwbytes : usize) -> *mut core::ffi::c_void);
+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_NO_MORE_ITEMS: WIN32_ERROR = 259u32;
+pub const ERROR_OUTOFMEMORY: WIN32_ERROR = 14u32;
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub struct FILETIME {
     pub dwLowDateTime: u32,
     pub dwHighDateTime: u32,
 }
+pub type HANDLE = isize;
+pub type HEAP_FLAGS = u32;
 pub type HKEY = isize;
 pub const HKEY_CLASSES_ROOT: HKEY = -2147483648i32 as _;
 pub const HKEY_CURRENT_CONFIG: HKEY = -2147483643i32 as _;

diff --git a/crates/libs/registry/src/value_iterator.rs b/crates/libs/registry/src/value_iterator.rs
@@ -6,7 +6,7 @@ pub struct ValueIterator<'a> {
     key: &'a Key,
     range: core::ops::Range<usize>,
     name: Vec<u16>,
-    value: Vec<u8>,
+    value: ValueBytes,
 }
 
 impl<'a> ValueIterator<'a> {
@@ -32,11 +32,13 @@ impl<'a> ValueIterator<'a> {
             )
         };
 
-        win32_error(result).map(|_| Self {
+        win32_error(result)?;
+
+        Ok(Self {
             key,
             range: 0..count as usize,
             name: vec![0; name_max_len as usize + 1],
-            value: vec![0; value_max_len as usize],
+            value: ValueBytes::new(value_max_len as usize)?,
         })
     }
 }
@@ -80,7 +82,6 @@ impl<'a> Iterator for ValueIterator<'a> {
                             Value::String(String::new())
                         } else {
                             let value = unsafe {
-                                // TODO: Vec<u8> does not guarantee alignment for u16
                                 core::slice::from_raw_parts(
                                     self.value.as_ptr() as *const u16,
                                     value_len as usize / 2,
@@ -95,7 +96,6 @@ impl<'a> Iterator for ValueIterator<'a> {
                             Value::MultiString(vec![])
                         } else {
                             let value = unsafe {
-                                // TODO: Vec<u8> does not guarantee alignment for u16
                                 core::slice::from_raw_parts(
                                     self.value.as_ptr() as *const u16,
                                     value_len as usize / 2,
@@ -119,3 +119,39 @@ impl<'a> Iterator for ValueIterator<'a> {
         })
     }
 }
+
+// Minimal `Vec` replacement providing `u16` alignment.
+struct ValueBytes(*mut core::ffi::c_void, usize);
+
+impl ValueBytes {
+    fn new(len: usize) -> Result<Self> {
+        // This pointer will have at least 8 byte alignment.
+        let ptr = unsafe { HeapAlloc(GetProcessHeap(), 0, len) };
+
+        if ptr.is_null() {
+            Err(Error::from_hresult(HRESULT::from_win32(ERROR_OUTOFMEMORY)))
+        } else {
+            Ok(Self(ptr, len))
+        }
+    }
+
+    fn as_mut_ptr(&mut self) -> *mut u8 {
+        self.0 as *mut u8
+    }
+}
+
+impl Drop for ValueBytes {
+    fn drop(&mut self) {
+        unsafe {
+            HeapFree(GetProcessHeap(), 0, self.0);
+        };
+    }
+}
+
+impl core::ops::Deref for ValueBytes {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        unsafe { core::slice::from_raw_parts(self.0 as *const u8, self.1) }
+    }
+}
diff --git a/crates/libs/registry/tests/bindings.txt b/crates/libs/registry/tests/bindings.txt
@@ -4,6 +4,10 @@
 --filter
     Windows.Win32.Foundation.ERROR_INVALID_DATA
     Windows.Win32.Foundation.ERROR_NO_MORE_ITEMS
+    Windows.Win32.Foundation.ERROR_OUTOFMEMORY
+    Windows.Win32.System.Memory.GetProcessHeap
+    Windows.Win32.System.Memory.HeapAlloc
+    Windows.Win32.System.Memory.HeapFree
     Windows.Win32.System.Registry.HKEY_CLASSES_ROOT
     Windows.Win32.System.Registry.HKEY_CURRENT_CONFIG
     Windows.Win32.System.Registry.HKEY_CURRENT_USER