From f6c49f4f5eadd7129c9fe220110e99542f823f74 Mon Sep 17 00:00:00 2001
From: Kenny Kerr <kenny@kennykerr.ca>
Date: Wed, 3 Jul 2024 10:46:10 -0500
Subject: [PATCH] Ensure that `HSTRING` builder provides initialized memory
 (#3141)

---
 crates/libs/strings/src/hstring_builder.rs |  8 ++++++-
 crates/libs/strings/src/hstring_header.rs  | 28 ++++++++++------------
 crates/tests/strings/tests/hstring.rs      |  4 ++++
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/crates/libs/strings/src/hstring_builder.rs b/crates/libs/strings/src/hstring_builder.rs
index 446451f0df..23679bb36a 100644
--- a/crates/libs/strings/src/hstring_builder.rs
+++ b/crates/libs/strings/src/hstring_builder.rs
@@ -9,7 +9,13 @@ pub struct HStringBuilder(*mut HStringHeader);
 impl HStringBuilder {
     /// Creates a preallocated `HSTRING` value.
     pub fn new(len: usize) -> Result<Self> {
-        Ok(Self(HStringHeader::alloc(len.try_into()?)?))
+        let header = HStringHeader::alloc(len.try_into()?)?;
+
+        if len > 0 {
+            unsafe { core::ptr::write_bytes((*header).data, 0, len) };
+        }
+
+        Ok(Self(header))
     }
 
     /// Shortens the string by removing any trailing 0 characters.
diff --git a/crates/libs/strings/src/hstring_header.rs b/crates/libs/strings/src/hstring_header.rs
index 84c3851519..dac8760a12 100644
--- a/crates/libs/strings/src/hstring_header.rs
+++ b/crates/libs/strings/src/hstring_header.rs
@@ -14,18 +14,18 @@ pub struct HStringHeader {
 }
 
 impl HStringHeader {
-    pub fn alloc(len: u32) -> Result<*mut HStringHeader> {
+    pub fn alloc(len: u32) -> Result<*mut Self> {
         if len == 0 {
             return Ok(core::ptr::null_mut());
         }
 
         // Allocate enough space for header and two bytes per character.
         // The space for the terminating null character is already accounted for inside of `HStringHeader`.
-        let bytes = core::mem::size_of::<HStringHeader>() + 2 * len as usize;
+        let bytes = core::mem::size_of::<Self>() + 2 * len as usize;
 
         #[cfg(windows)]
-        let header = unsafe { bindings::HeapAlloc(bindings::GetProcessHeap(), 0, bytes) }
-            as *mut HStringHeader;
+        let header =
+            unsafe { bindings::HeapAlloc(bindings::GetProcessHeap(), 0, bytes) } as *mut Self;
 
         #[cfg(not(windows))]
         let header = unsafe {
@@ -33,7 +33,7 @@ impl HStringHeader {
                 fn malloc(bytes: usize) -> *mut core::ffi::c_void;
             }
 
-            malloc(bytes) as *mut HStringHeader
+            malloc(bytes) as *mut Self
         };
 
         if header.is_null() {
@@ -42,7 +42,7 @@ impl HStringHeader {
 
         unsafe {
             // Use `ptr::write` (since `header` is unintialized). `HStringHeader` is safe to be all zeros.
-            header.write(core::mem::MaybeUninit::<HStringHeader>::zeroed().assume_init());
+            header.write(core::mem::MaybeUninit::<Self>::zeroed().assume_init());
             (*header).len = len;
             (*header).count = RefCount::new(1);
             (*header).data = &mut (*header).buffer_start;
@@ -51,17 +51,13 @@ impl HStringHeader {
         Ok(header)
     }
 
-    pub unsafe fn free(header: *mut HStringHeader) {
+    pub unsafe fn free(header: *mut Self) {
         if header.is_null() {
             return;
         }
 
-        let header = header as *mut _;
-
         #[cfg(windows)]
-        {
-            bindings::HeapFree(bindings::GetProcessHeap(), 0, header);
-        }
+        bindings::HeapFree(bindings::GetProcessHeap(), 0, header as *mut _);
 
         #[cfg(not(windows))]
         {
@@ -69,18 +65,18 @@ impl HStringHeader {
                 fn free(ptr: *mut core::ffi::c_void);
             }
 
-            free(header);
+            free(header as *mut _);
         }
     }
 
-    pub fn duplicate(&self) -> Result<*mut HStringHeader> {
+    pub fn duplicate(&self) -> Result<*mut Self> {
         if self.flags & HSTRING_REFERENCE_FLAG == 0 {
             // If this is not a "fast pass" string then simply increment the reference count.
             self.count.add_ref();
-            Ok(self as *const HStringHeader as *mut HStringHeader)
+            Ok(self as *const Self as *mut Self)
         } else {
             // Otherwise, allocate a new string and copy the value into the new string.
-            let copy = HStringHeader::alloc(self.len)?;
+            let copy = Self::alloc(self.len)?;
             // SAFETY: since we are duplicating the string it is safe to copy all data from self to the initialized `copy`.
             // We copy `len + 1` characters since `len` does not account for the terminating null character.
             unsafe {
diff --git a/crates/tests/strings/tests/hstring.rs b/crates/tests/strings/tests/hstring.rs
index 13597f4b90..bb7a35a586 100644
--- a/crates/tests/strings/tests/hstring.rs
+++ b/crates/tests/strings/tests/hstring.rs
@@ -47,5 +47,9 @@ fn hstring_builder() -> Result<()> {
     assert_eq!(h.len(), 5);
     assert_eq!(h.as_wide(), HELLO);
 
+    // HStringBuilder will initialize memory to zero.
+    let b = HStringBuilder::new(5)?;
+    assert_eq!(*b, [0, 0, 0, 0, 0]);
+
     Ok(())
 }