From 98c45cfa4d0921e88451e187581e5dfe633a4b73 Mon Sep 17 00:00:00 2001
From: Kenny Kerr <kenny@kennykerr.ca>
Date: Fri, 21 Jun 2024 07:56:19 -0500
Subject: [PATCH] registry type and bytes

---
 crates/libs/registry/src/key.rs           | 62 +++++++++++++++++++++--
 crates/libs/registry/src/lib.rs           |  3 ++
 crates/libs/registry/src/type.rs          | 21 ++++++++
 crates/tests/registry/tests/bad_string.rs | 36 +++++++++++++
 4 files changed, 117 insertions(+), 5 deletions(-)
 create mode 100644 crates/libs/registry/src/type.rs
 create mode 100644 crates/tests/registry/tests/bad_string.rs

diff --git a/crates/libs/registry/src/key.rs b/crates/libs/registry/src/key.rs
index 4b0076fca8..978890f239 100644
--- a/crates/libs/registry/src/key.rs
+++ b/crates/libs/registry/src/key.rs
@@ -115,6 +115,34 @@ impl Key {
         unsafe { self.set_value(name, REG_BINARY, value.as_ptr() as _, value.len()) }
     }
 
+    /// Gets the type for the name in the registry key.
+    pub fn get_type<T: AsRef<str>>(&self, name: T) -> Result<Type> {
+        let name = pcwstr(name);
+        let mut ty = 0;
+
+        let result = unsafe {
+            RegQueryValueExW(
+                self.0,
+                name.as_ptr(),
+                null(),
+                &mut ty,
+                null_mut(),
+                null_mut(),
+            )
+        };
+
+        win32_error(result)?;
+
+        Ok(match ty {
+            REG_DWORD => Type::U32,
+            REG_QWORD => Type::U64,
+            REG_BINARY => Type::Bytes,
+            REG_SZ | REG_EXPAND_SZ => Type::String,
+            REG_MULTI_SZ => Type::MultiString,
+            rest => Type::Unknown(rest),
+        })
+    }
+
     /// Gets the value for the name in the registry key.
     pub fn get_value<T: AsRef<str>>(&self, name: T) -> Result<Value> {
         let name = pcwstr(name);
@@ -252,11 +280,35 @@ impl Key {
 
     /// Gets the value for the name in the registry key.
     pub fn get_bytes<T: AsRef<str>>(&self, name: T) -> Result<Vec<u8>> {
-        if let Value::Bytes(value) = self.get_value(name)? {
-            Ok(value)
-        } else {
-            Err(invalid_data())
-        }
+        let name = pcwstr(name);
+        let mut len = 0;
+
+        let result = unsafe {
+            RegQueryValueExW(
+                self.0,
+                name.as_ptr(),
+                null(),
+                null_mut(),
+                null_mut(),
+                &mut len,
+            )
+        };
+
+        win32_error(result)?;
+        let mut value = vec![0u8; len as usize];
+
+        let result = unsafe {
+            RegQueryValueExW(
+                self.0,
+                name.as_ptr(),
+                null(),
+                null_mut(),
+                value.as_mut_ptr() as _,
+                &mut len,
+            )
+        };
+
+        win32_error(result).map(|_| value)
     }
 
     /// Gets the value for the name in the registry key.
diff --git a/crates/libs/registry/src/lib.rs b/crates/libs/registry/src/lib.rs
index 759397acbf..1d11207d5d 100644
--- a/crates/libs/registry/src/lib.rs
+++ b/crates/libs/registry/src/lib.rs
@@ -24,6 +24,9 @@ pub use key_iterator::KeyIterator;
 mod value_iterator;
 pub use value_iterator::ValueIterator;
 
+mod r#type;
+pub use r#type::Type;
+
 pub use windows_result::Result;
 use windows_result::*;
 
diff --git a/crates/libs/registry/src/type.rs b/crates/libs/registry/src/type.rs
new file mode 100644
index 0000000000..050e59b11b
--- /dev/null
+++ b/crates/libs/registry/src/type.rs
@@ -0,0 +1,21 @@
+/// The possible types that a registry value could have.
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub enum Type {
+    /// A 32-bit unsigned integer value.
+    U32,
+
+    /// A 64-bit unsigned integer value.
+    U64,
+
+    /// A string value.
+    String,
+
+    /// An array u8 bytes.
+    Bytes,
+
+    /// An array of string values.
+    MultiString,
+
+    /// An unknown or unsupported type.
+    Unknown(u32),
+}
diff --git a/crates/tests/registry/tests/bad_string.rs b/crates/tests/registry/tests/bad_string.rs
new file mode 100644
index 0000000000..412d8c33a4
--- /dev/null
+++ b/crates/tests/registry/tests/bad_string.rs
@@ -0,0 +1,36 @@
+use windows::{core::w, Win32::System::Registry::*};
+use windows_registry::*;
+
+#[test]
+fn bad_string() -> Result<()> {
+    let bad_string_bytes = vec![
+        0x00, 0xD8, // leading surrogate
+        0x01, 0x01, // bogus trailing surrogate
+        0x00, 0x00, // null
+    ];
+
+    let test_key = "software\\windows-rs\\tests\\bad_string";
+    _ = CURRENT_USER.remove_tree(test_key);
+    let key = CURRENT_USER.create(test_key)?;
+
+    unsafe {
+        RegSetValueExW(
+            HKEY(key.as_raw()),
+            w!("name"),
+            0,
+            REG_SZ,
+            Some(&bad_string_bytes),
+        )
+        .ok()?
+    };
+
+    let ty = key.get_type("name")?;
+    assert_eq!(ty, Type::String);
+
+    let value_as_string = key.get_string("name")?;
+    assert_eq!(value_as_string, "�ā");
+
+    let value_as_bytes = key.get_bytes("name")?;
+    assert_eq!(value_as_bytes, bad_string_bytes);
+    Ok(())
+}